revert catalog for non-athena platform

datahub-project · Jan 15, 2025 · 0626fda · 0626fda
1 parent e2b1a9a
commit 0626fda
Show file tree

Hide file tree

Showing 10 changed files with 126 additions and 95 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py b/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py
@@ -157,7 +157,7 @@ class GlueSourceConfig(
         default=None,
         description="The aws account id where the target glue catalog lives. If None, datahub will ingest glue in aws caller's account.",
     )
-    catalog_name: Optional[str] = Field(
+    athena_catalog_name: Optional[str] = Field(
         default=None, description="The aws athena catalog name"
     )
     ignore_resource_links: Optional[bool] = Field(
@@ -229,25 +229,41 @@ def platform_validator(cls, v: str) -> str:
             )
 
     def __init__(self, **data: Any):
+        """Post init configuration operations."""
         super().__init__(**data)
-        if self.catalog_id:
-            current_account_id = self.sts_client.get_caller_identity().get("Account")
-            if self.catalog_id == current_account_id:
-                self.catalog_name = DEFAULT_CATALOG_NAME
+        self._set_athena_catalog_name()
+
+    def _set_athena_catalog_name(self) -> None:
+        """Set the correct athena catalog name or raise an exception in case of misconfiguration."""
+        if self.platform == "athena":
+            if self.catalog_id:
+                current_account_id = self.sts_client.get_caller_identity().get(
+                    "Account"
+                )
+                if self.catalog_id == current_account_id:
+                    self.athena_catalog_name = DEFAULT_CATALOG_NAME
+                else:
+                    self._validate_athena_catalog_name()
             else:
-                self._validate_catalog_name()
+                self.athena_catalog_name = DEFAULT_CATALOG_NAME
         else:
-            self.catalog_name = DEFAULT_CATALOG_NAME
+            self.athena_catalog_name = None
+
+    def _validate_athena_catalog_name(self) -> None:
+        """Validate if athena catalog name is set correctly.
 
-    def _validate_catalog_name(self) -> None:
+        This method helps to avoid issue when the `athena_catalog_name` does not exist in a specified AWS account.
+        """
         effective_catalog_id = (
-            self.athena_client.get_data_catalog(Name=self.catalog_name)["DataCatalog"]
+            self.athena_client.get_data_catalog(Name=self.athena_catalog_name)[
+                "DataCatalog"
+            ]
             .get("Parameters", {})
             .get("catalog-id", "")
         )
         if effective_catalog_id != self.catalog_id:
             raise ValueError(
-                f"Catalog configuration mismatch for catalog name {self.catalog_name}."
+                f"Catalog configuration mismatch for catalog name {self.athena_catalog_name}."
                 f"Effective catalog_id: {effective_catalog_id}, configured catalog_id: {self.catalog_id}."
             )
 
@@ -495,8 +511,16 @@ def process_dataflow_node(
 
             # if data object is Glue table
             if "database" in node_args and "table_name" in node_args:
-                full_table_name = f"{self.source_config.catalog_name}.{node_args['database']}.{node_args['table_name']}"
-
+                full_table_name = ".".join(
+                    filter(
+                        None,
+                        [
+                            self.source_config.athena_catalog_name,
+                            node_args["database"],
+                            node_args["table_name"],
+                        ],
+                    )
+                )
                 # we know that the table will already be covered when ingesting Glue tables
                 node_urn = make_dataset_urn_with_platform_instance(
                     platform=self.platform,
@@ -1120,8 +1144,15 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
     def _gen_table_wu(self, table: TablePaginatorTypeDef) -> Iterable[MetadataWorkUnit]:
         database_name = table["DatabaseName"]
         table_name = table["Name"]
-        full_table_name = (
-            f"{self.source_config.catalog_name}.{database_name}.{table_name}"
+        full_table_name = ".".join(
+            filter(
+                None,
+                [
+                    self.source_config.athena_catalog_name,
+                    database_name,
+                    table_name,
+                ],
+            )
         )
         self.report.report_table_scanned()
         if not self.source_config.database_pattern.allowed(

diff --git a/metadata-ingestion/tests/unit/glue/glue_deleted_actor_mces_golden.json b/metadata-ingestion/tests/unit/glue/glue_deleted_actor_mces_golden.json
@@ -96,7 +96,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -133,7 +133,7 @@
                 },
                 {
                     "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
-                        "schemaName": "awsdatacatalog.test-database.test_jsons_markers",
+                        "schemaName": "test-database.test_jsons_markers",
                         "platform": "urn:li:dataPlatform:glue",
                         "version": 0,
                         "created": {
@@ -250,7 +250,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -269,7 +269,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {
@@ -286,7 +286,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
     "changeType": "UPSERT",
     "aspectName": "browsePathsV2",
     "aspect": {
@@ -309,7 +309,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -346,7 +346,7 @@
                 },
                 {
                     "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
-                        "schemaName": "awsdatacatalog.test-database.test_parquet",
+                        "schemaName": "test-database.test_parquet",
                         "platform": "urn:li:dataPlatform:glue",
                         "version": 0,
                         "created": {
@@ -464,7 +464,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -483,7 +483,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {
@@ -500,7 +500,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
     "changeType": "UPSERT",
     "aspectName": "browsePathsV2",
     "aspect": {
@@ -522,7 +522,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.flights-database.avro,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,flights-database.avro,PROD)",
     "changeType": "UPSERT",
     "aspectName": "status",
     "aspect": {

diff --git a/metadata-ingestion/tests/unit/glue/glue_delta_mces_golden.json b/metadata-ingestion/tests/unit/glue/glue_delta_mces_golden.json
@@ -58,7 +58,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.awsdatacatalog.delta-database.delta_table_1,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.delta-database.delta_table_1,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -82,7 +82,7 @@
                 },
                 {
                     "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
-                        "schemaName": "awsdatacatalog.delta-database.delta_table_1",
+                        "schemaName": "delta-database.delta_table_1",
                         "platform": "urn:li:dataPlatform:glue",
                         "version": 0,
                         "created": {
@@ -1525,7 +1525,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.awsdatacatalog.delta-database.delta_table_1,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.delta-database.delta_table_1,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -1538,7 +1538,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.awsdatacatalog.delta-database.delta_table_1,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.delta-database.delta_table_1,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {

diff --git a/metadata-ingestion/tests/unit/glue/glue_malformed_delta_mces_golden.json b/metadata-ingestion/tests/unit/glue/glue_malformed_delta_mces_golden.json
@@ -58,7 +58,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.awsdatacatalog.delta-database.delta_table_1,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.delta-database.delta_table_1,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -105,7 +105,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.awsdatacatalog.delta-database.delta_table_1,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.delta-database.delta_table_1,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -118,7 +118,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.awsdatacatalog.delta-database.delta_table_1,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,delta_platform_instance.delta-database.delta_table_1,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {

diff --git a/metadata-ingestion/tests/unit/glue/glue_mces_golden.json b/metadata-ingestion/tests/unit/glue/glue_mces_golden.json
@@ -167,7 +167,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.flights-database.avro,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,flights-database.avro,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -205,7 +205,7 @@
                 },
                 {
                     "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
-                        "schemaName": "awsdatacatalog.flights-database.avro",
+                        "schemaName": "flights-database.avro",
                         "platform": "urn:li:dataPlatform:glue",
                         "version": 0,
                         "created": {
@@ -370,7 +370,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.flights-database.avro,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,flights-database.avro,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -383,7 +383,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.flights-database.avro,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,flights-database.avro,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {
@@ -395,7 +395,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -432,7 +432,7 @@
                 },
                 {
                     "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
-                        "schemaName": "awsdatacatalog.test-database.test_jsons_markers",
+                        "schemaName": "test-database.test_jsons_markers",
                         "platform": "urn:li:dataPlatform:glue",
                         "version": 0,
                         "created": {
@@ -555,7 +555,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -568,7 +568,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_jsons_markers,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_jsons_markers,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {
@@ -580,7 +580,7 @@
 {
     "proposedSnapshot": {
         "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
-            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
             "aspects": [
                 {
                     "com.linkedin.pegasus2avro.common.Status": {
@@ -617,7 +617,7 @@
                 },
                 {
                     "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
-                        "schemaName": "awsdatacatalog.test-database.test_parquet",
+                        "schemaName": "test-database.test_parquet",
                         "platform": "urn:li:dataPlatform:glue",
                         "version": 0,
                         "created": {
@@ -741,7 +741,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
     "changeType": "UPSERT",
     "aspectName": "subTypes",
     "aspect": {
@@ -754,7 +754,7 @@
 },
 {
     "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)",
     "changeType": "UPSERT",
     "aspectName": "container",
     "aspect": {
@@ -896,7 +896,7 @@
                 {
                     "com.linkedin.pegasus2avro.datajob.DataJobInputOutput": {
                         "inputDatasets": [
-                            "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.flights-database.avro,PROD)"
+                            "urn:li:dataset:(urn:li:dataPlatform:glue,flights-database.avro,PROD)"
                         ],
                         "outputDatasets": [],
                         "inputDatajobs": []
@@ -1085,7 +1085,7 @@
                 {
                     "com.linkedin.pegasus2avro.datajob.DataJobInputOutput": {
                         "inputDatasets": [
-                            "urn:li:dataset:(urn:li:dataPlatform:glue,awsdatacatalog.test-database.test_parquet,PROD)"
+                            "urn:li:dataset:(urn:li:dataPlatform:glue,test-database.test_parquet,PROD)"
                         ],
                         "outputDatasets": [],
                         "inputDatajobs": []