From a8149d020231bd0a80660082d1d9a20314f90d72 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Tue, 3 May 2022 14:23:01 +0530 Subject: [PATCH] fix(snowflake): get external tables when there is default namespace --- .../datahub/ingestion/source/sql/snowflake.py | 37 ++++++++++++++----- .../ingestion/source_report/sql/snowflake.py | 5 ++- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py index 7c549408380cb9..d991373f11c86d 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py @@ -93,11 +93,32 @@ def get_metadata_engine( **self.config.options, ) - def inspect_version(self) -> Any: + def inspect_session_metadata(self) -> Any: db_engine = self.get_metadata_engine() - logger.info("Checking current version") - for db_row in db_engine.execute("select CURRENT_VERSION()"): - self.report.saas_version = db_row[0] + try: + logger.info("Checking current version") + for db_row in db_engine.execute("select CURRENT_VERSION()"): + self.report.saas_version = db_row[0] + except Exception as e: + self.report.report_failure("version", f"Error: {e}") + try: + logger.info("Checking current warehouse") + for db_row in db_engine.execute("select current_warehouse()"): + self.report.default_warehouse = db_row[0] + except Exception as e: + self.report.report_failure("current_warehouse", f"Error: {e}") + try: + logger.info("Checking current database") + for db_row in db_engine.execute("select current_database()"): + self.report.default_db = db_row[0] + except Exception as e: + self.report.report_failure("current_database", f"Error: {e}") + try: + logger.info("Checking current schema") + for db_row in db_engine.execute("select current_schema()"): + self.report.default_schema = db_row[0] + except Exception as e: + self.report.report_failure("current_schema", f"Error: {e}") def inspect_role_grants(self) -> Any: db_engine = self.get_metadata_engine() @@ -345,7 +366,7 @@ def _populate_external_lineage(self) -> None: ) # Handles the case for explicitly created external tables. # NOTE: Snowflake does not log this information to the access_history table. - external_tables_query: str = "show external tables" + external_tables_query: str = "show external tables in account" try: for db_row in engine.execute(external_tables_query): key = ( @@ -635,11 +656,7 @@ def get_workunits(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]: if not self.should_run_ingestion(): return - try: - self.inspect_version() - except Exception as e: - self.report.report_failure("version", f"Error: {e}") - return + self.inspect_session_metadata() self.inspect_role_grants() for wu in super().get_workunits(): diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py index 83ea47459678d0..7235fea3b3b709 100644 --- a/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py +++ b/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py @@ -29,7 +29,10 @@ class SnowflakeReport(BaseSnowflakeReport, SQLSourceReport): provision_role_success: bool = False # https://community.snowflake.com/s/topic/0TO0Z000000Unu5WAC/releases - saas_version: str = "" + saas_version: Optional[str] = None + default_warehouse: Optional[str] = None + default_db: Optional[str] = None + default_schema: Optional[str] = None role: str = "" check_role_grants: Optional[bool] = None role_grants: List[str] = field(default_factory=list)