From 4d11268895c9e3b44d192c60e8eb84efe1b9104e Mon Sep 17 00:00:00 2001
From: TJ Murphy <1796+teej@users.noreply.github.com>
Date: Fri, 8 Nov 2024 00:14:56 -0800
Subject: [PATCH] tag tests

---
 .../data_provider/test_list_resource.py       |  8 ++
 titan/data_provider.py                        | 79 +++++++++----------
 .../business_critical.yml                     |  6 ++
 tools/test_account_configs/enterprise.yml     |  6 ++
 4 files changed, 56 insertions(+), 43 deletions(-)

diff --git a/tests/integration/data_provider/test_list_resource.py b/tests/integration/data_provider/test_list_resource.py
index 5070e90..b9b1db8 100644
--- a/tests/integration/data_provider/test_list_resource.py
+++ b/tests/integration/data_provider/test_list_resource.py
@@ -88,3 +88,11 @@ def test_list_resource(cursor, list_resources_database, resource, marked_for_cle
     list_resources = data_provider.list_resource(cursor, resource_label_for_type(resource.resource_type))
     assert len(list_resources) > 0
     assert resource.fqn in list_resources
+
+
+@pytest.mark.enterprise
+def test_list_tag_references(cursor):
+    data_provider.fetch_session.cache_clear()
+    reset_cache()
+    tag_references = data_provider.list_tag_references(cursor)
+    assert len(tag_references) > 0
diff --git a/titan/data_provider.py b/titan/data_provider.py
index dcbf605..0020531 100644
--- a/titan/data_provider.py
+++ b/titan/data_provider.py
@@ -2613,50 +2613,43 @@ def list_tables(session: SnowflakeConnection) -> list[FQN]:
     return tables
 
 
-# def list_tag_references(session: SnowflakeConnection) -> list[FQN]:
-#     try:
-#         show_result = execute(session, "SHOW TAGS IN ACCOUNT")
-#         tag_references: list[FQN] = []
-#         for tag in show_result:
-#             if tag["database_name"] in SYSTEM_DATABASES or tag["schema_name"] == "INFORMATION_SCHEMA":
-#                 continue
-
-#             tag_refs = execute(
-#                 session,
-#                 f"""
-#                     SELECT *
-#                     FROM table(snowflake.account_usage.tag_references_with_lineage(
-#                         '{tag['database_name']}.{tag['schema_name']}.{tag['name']}'
-#                     ))
-#                 """,
-#             )
-
-#             for ref in tag_refs:
-#                 if ref["OBJECT_DELETED"] is not None:
-#                     continue
-#                 print(ref)
-#             # raise
-#         return tag_references
-
-#         #     tags.append(
-#         #         FQN(
-#         #             database=resource_name_from_snowflake_metadata(row["database_name"]),
-#         #             schema=resource_name_from_snowflake_metadata(tag["schema_name"]),
-#         #             name=resource_name_from_snowflake_metadata(tag["name"]),
-#         #         )
-#         #     )
-#         # return tags
-#     except ProgrammingError as err:
-#         if err.errno == UNSUPPORTED_FEATURE:
-#             return []
-#         else:
-#             raise
+def list_tag_references(session: SnowflakeConnection) -> list[FQN]:
+    try:
+        show_result = execute(session, "SHOW TAGS IN ACCOUNT")
+        tag_references: list[FQN] = []
+        for tag in show_result:
+            if tag["database_name"] in SYSTEM_DATABASES or tag["schema_name"] == "INFORMATION_SCHEMA":
+                continue
 
-#     tag_map = {}
-#     for tag_ref in tag_refs:
-#         tag_name = f"{tag_ref['TAG_DATABASE']}.{tag_ref['TAG_SCHEMA']}.{tag_ref['TAG_NAME']}"
-#         tag_map[tag_name] = tag_ref["TAG_VALUE"]
-#     return tag_map
+            tag_refs = execute(
+                session,
+                f"""
+                    SELECT *
+                    FROM table(snowflake.account_usage.tag_references_with_lineage(
+                        '{tag['database_name']}.{tag['schema_name']}.{tag['name']}'
+                    ))
+                """,
+            )
+
+            for ref in tag_refs:
+                if ref["OBJECT_DELETED"] is not None:
+                    continue
+
+                tag_references.append(
+                    FQN(
+                        database=resource_name_from_snowflake_metadata(ref["TAG_DATABASE"]),
+                        schema=resource_name_from_snowflake_metadata(ref["TAG_SCHEMA"]),
+                        name=resource_name_from_snowflake_metadata(ref["TAG_NAME"]),
+                    )
+                )
+
+        return tag_references
+
+    except ProgrammingError as err:
+        if err.errno == UNSUPPORTED_FEATURE:
+            return []
+        else:
+            raise
 
 
 def list_tags(session: SnowflakeConnection) -> list[FQN]:
diff --git a/tools/test_account_configs/business_critical.yml b/tools/test_account_configs/business_critical.yml
index ffe1f5f..f271c46 100644
--- a/tools/test_account_configs/business_critical.yml
+++ b/tools/test_account_configs/business_critical.yml
@@ -8,6 +8,12 @@ grants:
   - GRANT APPLY ROW ACCESS POLICY ON ACCOUNT TO ROLE EVERY_PRIVILEGE
   - GRANT APPLY TAG ON ACCOUNT TO ROLE EVERY_PRIVILEGE
 
+schemas:
+  - name: tagged_schema
+    database: static_database
+    tags:
+      static_database.public.static_tag: static_tag_value
+
 tags:
   - name: static_tag
     database: static_database
diff --git a/tools/test_account_configs/enterprise.yml b/tools/test_account_configs/enterprise.yml
index ffe1f5f..f271c46 100644
--- a/tools/test_account_configs/enterprise.yml
+++ b/tools/test_account_configs/enterprise.yml
@@ -8,6 +8,12 @@ grants:
   - GRANT APPLY ROW ACCESS POLICY ON ACCOUNT TO ROLE EVERY_PRIVILEGE
   - GRANT APPLY TAG ON ACCOUNT TO ROLE EVERY_PRIVILEGE
 
+schemas:
+  - name: tagged_schema
+    database: static_database
+    tags:
+      static_database.public.static_tag: static_tag_value
+
 tags:
   - name: static_tag
     database: static_database