feat(ingest): capture table descriptions

datahub-project · Mar 2, 2021 · cd569c0 · cd569c0
1 parent 83cd89a
commit cd569c0
Show file tree

Hide file tree

Showing 6 changed files with 647 additions and 303 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql_common.py
@@ -2,7 +2,7 @@
 import time
 from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, List, Optional
+from typing import Any, Iterable, List, Optional
 
 from sqlalchemy import create_engine
 from sqlalchemy.engine import reflection
@@ -26,6 +26,7 @@
     SchemaMetadata,
     StringTypeClass,
 )
+from datahub.metadata.schema_classes import DatasetPropertiesClass
 
 logger = logging.getLogger(__name__)
 
@@ -144,7 +145,7 @@ def __init__(self, config, ctx, platform: str):
         self.platform = platform
         self.report = SQLSourceReport()
 
-    def get_workunits(self):
+    def get_workunits(self) -> Iterable[SqlWorkUnit]:
         env: str = "PROD"
         sql_config = self.config
         platform = self.platform
@@ -163,10 +164,26 @@ def get_workunits(self):
 
                 if sql_config.table_pattern.allowed(dataset_name):
                     columns = inspector.get_columns(table, schema)
+                    try:
+                        description: Optional[str] = inspector.get_table_comment(
+                            table, schema
+                        )["text"]
+                    except NotImplementedError:
+                        description = None
+
+                    # TODO: capture inspector.get_pk_constraint
+                    # TODO: capture inspector.get_sorted_table_and_fkc_names
+
                     mce = MetadataChangeEvent()
 
                     dataset_snapshot = DatasetSnapshot()
                     dataset_snapshot.urn = f"urn:li:dataset:(urn:li:dataPlatform:{platform},{dataset_name},{env})"
+                    if description is not None:
+                        dataset_properties = DatasetPropertiesClass(
+                            description=description,
+                            # uri=dataset_name,
+                        )
+                        dataset_snapshot.aspects.append(dataset_properties)
                     schema_metadata = get_schema_metadata(
                         self.report, dataset_name, platform, columns
                     )