feat(ingest/dbt-cloud): update metadata_endpoint inference (#11041)

metadata-ingestion/src/datahub/cli/get_cli.py

@@ -56,6 +56,7 @@ def urn(ctx: Any, urn: Optional[str], aspect: List[str], details: bool) -> None:
                 entity_urn=urn,
                 aspects=aspect,
                 typed=False,
+                details=details,
             ),
             sort_keys=True,
             indent=2,

metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py

@@ -40,8 +40,7 @@ class DBTCloudConfig(DBTCommonConfig):
 
     metadata_endpoint: str = Field(
         default="https://metadata.cloud.getdbt.com/graphql",
-        description="The dbt Cloud metadata API endpoint. This is deprecated, and will be removed in a future release. Please use access_url instead.",
-        deprecated=True,
+        description="The dbt Cloud metadata API endpoint. If not provided, we will try to infer it from the access_url.",
     )
 
     token: str = Field(
@@ -66,13 +65,39 @@ class DBTCloudConfig(DBTCommonConfig):
     @root_validator(pre=True)
     def set_metadata_endpoint(cls, values: dict) -> dict:
         if values.get("access_url") and not values.get("metadata_endpoint"):
-            parsed_uri = urlparse(values["access_url"])
-            values[
-                "metadata_endpoint"
-            ] = f"{parsed_uri.scheme}://metadata.{parsed_uri.netloc}/graphql"
+            metadata_endpoint = infer_metadata_endpoint(values["access_url"])
+            if metadata_endpoint is None:
+                raise ValueError(
+                    "Unable to infer the metadata endpoint from the access URL. Please provide a metadata endpoint."
+                )
+            values["metadata_endpoint"] = metadata_endpoint
         return values
 
 
+def infer_metadata_endpoint(access_url: str) -> Optional[str]:
+    # See https://docs.getdbt.com/docs/cloud/about-cloud/access-regions-ip-addresses#api-access-urls
+    # and https://docs.getdbt.com/docs/dbt-cloud-apis/discovery-querying#discovery-api-endpoints
+
+    try:
+        parsed_uri = urlparse(access_url)
+        assert parsed_uri.scheme is not None
+        assert parsed_uri.hostname is not None
+    except Exception as e:
+        logger.debug(f"Unable to parse access URL {access_url}: {e}", exc_info=e)
+        return None
+
+    if parsed_uri.hostname.endswith(".dbt.com"):
+        # For cell-based deployments.
+        # prefix.region.dbt.com -> prefix.metadata.region.dbt.com
+        hostname_parts = parsed_uri.hostname.split(".", maxsplit=1)
+        return f"{parsed_uri.scheme}://{hostname_parts[0]}.metadata.{hostname_parts[1]}/graphql"
+    elif parsed_uri.hostname.endswith(".getdbt.com"):
+        return f"{parsed_uri.scheme}://metadata.{parsed_uri.netloc}/graphql"
+    else:
+        # The self-hosted variants also have the metadata. prefix.
+        return f"{parsed_uri.scheme}://metadata.{parsed_uri.netloc}/graphql"
+
+
 _DBT_GRAPHQL_COMMON_FIELDS = """
   runId
   accountId

metadata-ingestion/tests/unit/test_dbt_source.py

@@ -7,7 +7,10 @@
 
 from datahub.emitter import mce_builder
 from datahub.ingestion.api.common import PipelineContext
-from datahub.ingestion.source.dbt.dbt_cloud import DBTCloudConfig
+from datahub.ingestion.source.dbt.dbt_cloud import (
+    DBTCloudConfig,
+    infer_metadata_endpoint,
+)
 from datahub.ingestion.source.dbt.dbt_core import (
     DBTCoreConfig,
     DBTCoreSource,
@@ -366,7 +369,7 @@ def test_dbt_entity_emission_configuration_helpers():
 
 def test_dbt_cloud_config_access_url():
     config_dict = {
-        "access_url": "https://my-dbt-cloud.dbt.com",
+        "access_url": "https://emea.getdbt.com",
         "token": "dummy_token",
         "account_id": "123456",
         "project_id": "1234567",
@@ -375,8 +378,8 @@ def test_dbt_cloud_config_access_url():
         "target_platform": "dummy_platform",
     }
     config = DBTCloudConfig.parse_obj(config_dict)
-    assert config.access_url == "https://my-dbt-cloud.dbt.com"
-    assert config.metadata_endpoint == "https://metadata.my-dbt-cloud.dbt.com/graphql"
+    assert config.access_url == "https://emea.getdbt.com"
+    assert config.metadata_endpoint == "https://metadata.emea.getdbt.com/graphql"
 
 
 def test_dbt_cloud_config_with_defined_metadata_endpoint():
@@ -398,6 +401,20 @@ def test_dbt_cloud_config_with_defined_metadata_endpoint():
     )
 
 
+def test_infer_metadata_endpoint() -> None:
+    assert (
+        infer_metadata_endpoint("https://cloud.getdbt.com")
+        == "https://metadata.cloud.getdbt.com/graphql"
+    )
+    assert (
+        infer_metadata_endpoint("https://prefix.us1.dbt.com")
+        == "https://prefix.metadata.us1.dbt.com/graphql"
+    )
+    assert (
+        infer_metadata_endpoint("http://dbt.corp.internal")
+    ) == "http://metadata.dbt.corp.internal/graphql"
+
+
 def test_dbt_time_parsing() -> None:
     time_formats = [
         "2024-03-28T05:56:15.236210Z",