Skip to content

Commit

Permalink
feat(ingest): add external url for snowflake objects (#6580)
Browse files Browse the repository at this point in the history
  • Loading branch information
@mayurinehate
mayurinehate authored Dec 2, 2022
1 parent 43775ec commit 1689212
Show file tree
Hide file tree
Showing 6 changed files with 1,000 additions and 13 deletions.
5 changes: 5 additions & 0 deletions metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,11 @@ class SnowflakeV2Config(SnowflakeConfig, SnowflakeUsageConfig):
description="For details, refer [Classification](../../../../metadata-ingestion/docs/dev_guides/classification.md).",
)

include_external_url: bool = Field(
default=True,
description="Whether to populate Snowsight url for Snowflake Objects",
)

@root_validator(pre=False)
def validate_unsupported_configs(cls, values: Dict) -> Dict:

Expand Down
59 changes: 59 additions & 0 deletions metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_utils.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import logging
from enum import Enum
from functools import lru_cache
from typing import Any, Optional

from snowflake.connector import SnowflakeConnection
Expand All @@ -12,6 +14,18 @@
from datahub.metadata.com.linkedin.pegasus2avro.events.metadata import ChangeType
from datahub.metadata.schema_classes import _Aspect

logger: logging.Logger = logging.getLogger(__name__)


class SnowflakeCloudProvider(str, Enum):
AWS = "aws"
GCP = "gcp"
AZURE = "azure"


SNOWFLAKE_DEFAULT_CLOUD_REGION_ID = "us-west-2"
SNOWFLAKE_DEFAULT_CLOUD = SnowflakeCloudProvider.AWS


# Required only for mypy, since we are using mixin classes, and not inheritance.
# Reference - https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes
Expand Down Expand Up @@ -59,6 +73,51 @@ class SnowflakeCommonMixin:

platform = "snowflake"

@staticmethod
@lru_cache(maxsize=128)
def create_snowsight_base_url(account_id: str) -> Optional[str]:
cloud: Optional[str] = None
account_locator: Optional[str] = None
cloud_region_id: Optional[str] = None
privatelink: bool = False

if "." not in account_id: # e.g. xy12345
account_locator = account_id.lower()
cloud_region_id = SNOWFLAKE_DEFAULT_CLOUD_REGION_ID
else:
parts = account_id.split(".")
if len(parts) == 2: # e.g. xy12345.us-east-1
account_locator = parts[0].lower()
cloud_region_id = parts[1].lower()
elif len(parts) == 3 and parts[2] in (
SnowflakeCloudProvider.AWS,
SnowflakeCloudProvider.GCP,
SnowflakeCloudProvider.AZURE,
):
# e.g. xy12345.ap-south-1.aws or xy12345.us-central1.gcp or xy12345.west-us-2.azure
# NOT xy12345.us-west-2.privatelink or xy12345.eu-central-1.privatelink
account_locator = parts[0].lower()
cloud_region_id = parts[1].lower()
cloud = parts[2].lower()
elif len(parts) == 3 and parts[2] == "privatelink":
account_locator = parts[0].lower()
cloud_region_id = parts[1].lower()
privatelink = True
else:
logger.warning(
f"Could not create Snowsight base url for account {account_id}."
)
return None

if not privatelink and (cloud is None or cloud == SNOWFLAKE_DEFAULT_CLOUD):
return f"https://app.snowflake.com/{cloud_region_id}/{account_locator}/"
elif privatelink:
return f"https://app.{account_locator}.{cloud_region_id}.privatelink.snowflakecomputing.com/"
return f"https://app.snowflake.com/{cloud_region_id}.{cloud}/{account_locator}/"

def get_snowsight_base_url(self: SnowflakeCommonProtocol) -> Optional[str]:
return SnowflakeCommonMixin.create_snowsight_base_url(self.config.get_account())

def _is_dataset_pattern_allowed(
self: SnowflakeCommonProtocol,
dataset_name: Optional[str],
Expand Down
37 changes: 37 additions & 0 deletions metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -646,6 +646,14 @@ def gen_dataset_workunits(
description=table.comment,
qualifiedName=dataset_name,
customProperties={**upstream_column_props},
externalUrl=self.get_external_url_for_table(
table.name,
schema_name,
db_name,
"table" if isinstance(table, SnowflakeTable) else "view",
)
if self.config.include_external_url
else None,
)
yield self.wrap_aspect_as_workunit(
"dataset", dataset_urn, "datasetProperties", dataset_properties
Expand Down Expand Up @@ -889,6 +897,9 @@ def gen_database_containers(
description=database.comment,
sub_types=[SqlContainerSubTypes.DATABASE],
domain_urn=domain_urn,
external_url=self.get_external_url_for_database(database.name)
if self.config.include_external_url
else None,
)

self.stale_entity_removal_handler.add_entity_to_state(
Expand Down Expand Up @@ -922,6 +933,9 @@ def gen_schema_containers(
description=schema.comment,
sub_types=[SqlContainerSubTypes.SCHEMA],
parent_container_key=database_container_key,
external_url=self.get_external_url_for_schema(schema.name, db_name)
if self.config.include_external_url
else None,
)

for wu in container_workunits:
Expand Down Expand Up @@ -1077,3 +1091,26 @@ def get_sample_values_for_table(self, conn, table_name, schema_name, db_name):
df = pd.DataFrame(dat, columns=[col.name for col in cur.description])

return df

# domain is either "view" or "table"
def get_external_url_for_table(
self, table_name: str, schema_name: str, db_name: str, domain: str
) -> Optional[str]:
base_url = self.get_snowsight_base_url()
if base_url is not None:
return f"{base_url}#/data/databases/{db_name}/schemas/{schema_name}/{domain}/{table_name}/"
return None

def get_external_url_for_schema(
self, schema_name: str, db_name: str
) -> Optional[str]:
base_url = self.get_snowsight_base_url()
if base_url is not None:
return f"{base_url}#/data/databases/{db_name}/schemas/{schema_name}/"
return None

def get_external_url_for_database(self, db_name: str) -> Optional[str]:
base_url = self.get_snowsight_base_url()
if base_url is not None:
return f"{base_url}#/data/databases/{db_name}/"
return None
24 changes: 12 additions & 12 deletions metadata-ingestion/tests/integration/snowflake-beta/snowflake_beta_golden.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
"changeType": "UPSERT",
"aspectName": "containerProperties",
"aspect": {
"value": "{\"customProperties\": {\"platform\": \"snowflake\", \"instance\": \"PROD\", \"database\": \"test_db\"}, \"name\": \"TEST_DB\", \"description\": \"Comment for TEST_DB\"}",
"value": "{\"customProperties\": {\"platform\": \"snowflake\", \"instance\": \"PROD\", \"database\": \"test_db\"}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/\", \"name\": \"TEST_DB\", \"description\": \"Comment for TEST_DB\"}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -61,7 +61,7 @@
"changeType": "UPSERT",
"aspectName": "containerProperties",
"aspect": {
"value": "{\"customProperties\": {\"platform\": \"snowflake\", \"instance\": \"PROD\", \"database\": \"test_db\", \"schema\": \"test_schema\"}, \"name\": \"TEST_SCHEMA\", \"description\": \"comment for TEST_DB.TEST_SCHEMA\"}",
"value": "{\"customProperties\": {\"platform\": \"snowflake\", \"instance\": \"PROD\", \"database\": \"test_db\", \"schema\": \"test_schema\"}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/\", \"name\": \"TEST_SCHEMA\", \"description\": \"comment for TEST_DB.TEST_SCHEMA\"}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -159,7 +159,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_1\", \"qualifiedName\": \"test_db.test_schema.table_1\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_1/\", \"name\": \"TABLE_1\", \"qualifiedName\": \"test_db.test_schema.table_1\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -243,7 +243,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_2\", \"qualifiedName\": \"test_db.test_schema.table_2\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_2/\", \"name\": \"TABLE_2\", \"qualifiedName\": \"test_db.test_schema.table_2\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -327,7 +327,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_3\", \"qualifiedName\": \"test_db.test_schema.table_3\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_3/\", \"name\": \"TABLE_3\", \"qualifiedName\": \"test_db.test_schema.table_3\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -411,7 +411,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_4\", \"qualifiedName\": \"test_db.test_schema.table_4\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_4/\", \"name\": \"TABLE_4\", \"qualifiedName\": \"test_db.test_schema.table_4\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -495,7 +495,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_5\", \"qualifiedName\": \"test_db.test_schema.table_5\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_5/\", \"name\": \"TABLE_5\", \"qualifiedName\": \"test_db.test_schema.table_5\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -579,7 +579,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_6\", \"qualifiedName\": \"test_db.test_schema.table_6\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_6/\", \"name\": \"TABLE_6\", \"qualifiedName\": \"test_db.test_schema.table_6\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -663,7 +663,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_7\", \"qualifiedName\": \"test_db.test_schema.table_7\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_7/\", \"name\": \"TABLE_7\", \"qualifiedName\": \"test_db.test_schema.table_7\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -747,7 +747,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_8\", \"qualifiedName\": \"test_db.test_schema.table_8\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_8/\", \"name\": \"TABLE_8\", \"qualifiedName\": \"test_db.test_schema.table_8\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -831,7 +831,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_9\", \"qualifiedName\": \"test_db.test_schema.table_9\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_9/\", \"name\": \"TABLE_9\", \"qualifiedName\": \"test_db.test_schema.table_9\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down Expand Up @@ -915,7 +915,7 @@
"changeType": "UPSERT",
"aspectName": "datasetProperties",
"aspect": {
"value": "{\"customProperties\": {}, \"name\": \"TABLE_10\", \"qualifiedName\": \"test_db.test_schema.table_10\", \"description\": \"Comment for Table\", \"tags\": []}",
"value": "{\"customProperties\": {}, \"externalUrl\": \"https://app.snowflake.com/ap-south-1/abc12345/#/data/databases/TEST_DB/schemas/TEST_SCHEMA/table/TABLE_10/\", \"name\": \"TABLE_10\", \"qualifiedName\": \"test_db.test_schema.table_10\", \"description\": \"Comment for Table\", \"tags\": []}",
"contentType": "application/json"
},
"systemMetadata": {
Expand Down
Loading

0 comments on commit 1689212

Please sign in to comment.