From aa1c855370dae3098fa1d3049aa0665930b7053c Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Tue, 16 Jul 2024 16:41:40 -0700 Subject: [PATCH] fix(ingest/snowflake): fix test connection Fixes a breakage from https://github.com/datahub-project/datahub/pull/10835 --- .../ingestion/source/snowflake/snowflake_v2.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py index d8eda98da422b9..11558aa74e9817 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py @@ -277,10 +277,12 @@ class SnowflakePrivilege: capabilities: List[SourceCapability] = [c.capability for c in SnowflakeV2Source.get_capabilities() if c.capability not in (SourceCapability.PLATFORM_INSTANCE, SourceCapability.DOMAINS, SourceCapability.DELETION_DETECTION)] # type: ignore cur = conn.query("select current_role()") - current_role = [row[0] for row in cur][0] + current_role = [row["CURRENT_ROLE()"] for row in cur][0] cur = conn.query("select current_secondary_roles()") - secondary_roles_str = json.loads([row[0] for row in cur][0])["roles"] + secondary_roles_str = json.loads( + [row["CURRENT_SECONDARY_ROLES()"] for row in cur][0] + )["roles"] secondary_roles = ( [] if secondary_roles_str == "" else secondary_roles_str.split(",") ) @@ -299,7 +301,9 @@ class SnowflakePrivilege: cur = conn.query(f'show grants to role "{role}"') for row in cur: privilege = SnowflakePrivilege( - privilege=row[1], object_type=row[2], object_name=row[3] + privilege=row["privilege"], + object_type=row["granted_on"], + object_name=row["name"], ) privileges.append(privilege) @@ -362,7 +366,7 @@ class SnowflakePrivilege: roles.append(privilege.object_name) cur = conn.query("select current_warehouse()") - current_warehouse = [row[0] for row in cur][0] + current_warehouse = [row["CURRENT_WAREHOUSE()"] for row in cur][0] default_failure_messages = { SourceCapability.SCHEMA_METADATA: "Either no tables exist or current role does not have permissions to access them",