feat(ingest/dremio): Retrieve default_schema for SQL views (#11832)

datahub-project · Nov 29, 2024 · 74a8488 · 74a8488
1 parent c42f779
commit 74a8488
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py
@@ -774,3 +774,14 @@ def process_source_and_containers(source):
                 containers.extend(future.result())
 
         return containers
+
+    def get_context_for_vds(self, resource_id: str) -> str:
+        context_array = self.get(
+            url=f"/catalog/{resource_id}",
+        ).get("sqlContext")
+        if context_array:
+            return ".".join(
+                f'"{part}"' if "." in part else f"{part}" for part in context_array
+            )
+        else:
+            return ""
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_entities.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_entities.py
@@ -200,6 +200,7 @@ class DremioDataset:
     columns: List[DremioDatasetColumn]
     sql_definition: Optional[str]
     dataset_type: DremioDatasetType
+    default_schema: Optional[str]
     owner: Optional[str]
     owner_type: Optional[str]
     created: str
@@ -235,6 +236,9 @@ def __init__(
 
         if self.sql_definition:
             self.dataset_type = DremioDatasetType.VIEW
+            self.default_schema = api_operations.get_context_for_vds(
+                resource_id=self.resource_id
+            )
         else:
             self.dataset_type = DremioDatasetType.TABLE
 

diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_source.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_source.py
@@ -417,6 +417,7 @@ def process_dataset(
                     view_urn=dataset_urn,
                     view_definition=dataset_info.sql_definition,
                     default_db=self.default_db,
+                    default_schema=dataset_info.default_schema,
                 )
 
         elif dataset_info.dataset_type == DremioDatasetType.TABLE: