acryl-datahub 1.1.1rc3__py3-none-any.whl → 1.2.0__py3-none-any.whl

datahub/ingestion/source/sql/vertica.py

@@ -25,6 +25,10 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_reader import DataReader
+from datahub.ingestion.source.common.subtypes import (
+    DatasetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -41,7 +45,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.dataset import UpstreamLineage
 from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     DatasetLineageTypeClass,
     DatasetPropertiesClass,
     SubTypesClass,
@@ -113,10 +116,14 @@ class VerticaConfig(BasicSQLAlchemyConfig):
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default, can be disabled via configuration `include_view_lineage` and `include_projection_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.PROJECTIONS,
+    ],
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class VerticaSource(SQLAlchemySource):
@@ -493,11 +500,8 @@ class VerticaSource(SQLAlchemySource):
         if dpi_aspect:
             yield dpi_aspect
         yield MetadataChangeProposalWrapper(
-            entityType="dataset",
-            changeType=ChangeTypeClass.UPSERT,
             entityUrn=dataset_urn,
-            aspectName="subTypes",
-            aspect=SubTypesClass(typeNames=["Projections"]),
+            aspect=SubTypesClass(typeNames=[DatasetSubTypes.PROJECTIONS]),
         ).as_workunit()
 
         if self.config.domain:

datahub/ingestion/source/sql_queries.py

@@ -66,7 +66,7 @@ class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
         description="The default schema to use for unqualified table names",
         default=None,
     )
-    default_dialect: Optional[str] = Field(
+    override_dialect: Optional[str] = Field(
         description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
         default=None,
     )
@@ -181,7 +181,7 @@ class SqlQueriesSource(Source):
             schema_resolver=self.schema_resolver,
             default_db=self.config.default_db,
             default_schema=self.config.default_schema,
-            default_dialect=self.config.default_dialect,
+            override_dialect=self.config.override_dialect,
         )
         if result.debug_info.table_error:
             logger.info(f"Error parsing table lineage, {result.debug_info.table_error}")

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -179,7 +179,7 @@ class StatefulIngestionReport(SourceReport):
 
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class StatefulIngestionSourceBase(Source):

datahub/ingestion/source/superset.py

@@ -272,7 +272,7 @@ def get_filter_name(filter_obj):
 @config_class(SupersetConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 @capability(SourceCapability.DOMAINS, "Enabled by `domain` config to assign domain_key")
 @capability(SourceCapability.LINEAGE_COARSE, "Supported by default")
@@ -658,6 +658,7 @@ class SupersetSource(StatefulIngestionSourceBase):
         if datasource_id:
             dataset_info = self.get_dataset_info(datasource_id).get("result", {})
             dataset_column_info = dataset_info.get("columns", [])
+            dataset_metric_info = dataset_info.get("metrics", [])
 
             for column in dataset_column_info:
                 col_name = column.get("column_name", "")
@@ -671,6 +672,17 @@ class SupersetSource(StatefulIngestionSourceBase):
                     continue
 
                 dataset_columns.append((col_name, col_type, col_description))
+
+            for metric in dataset_metric_info:
+                metric_name = metric.get("metric_name", "")
+                metric_type = metric.get("metric_type", "")
+                metric_description = metric.get("description", "")
+
+                if metric_name == "" or metric_type == "":
+                    logger.info(f"could not construct metric lineage for {metric}")
+                    continue
+
+                dataset_columns.append((metric_name, metric_type, metric_description))
         else:
             # if no datasource id, cannot build cll, just return
             logger.warning(
@@ -972,19 +984,44 @@ class SupersetSource(StatefulIngestionSourceBase):
             schema_fields.append(field)
         return schema_fields
 
+    def gen_metric_schema_fields(
+        self, metric_data: List[Dict[str, Any]]
+    ) -> List[SchemaField]:
+        schema_fields: List[SchemaField] = []
+        for metric in metric_data:
+            metric_type = metric.get("metric_type", "")
+            data_type = resolve_sql_type(metric_type)
+            if data_type is None:
+                data_type = NullType()
+
+            field = SchemaField(
+                fieldPath=metric.get("metric_name", ""),
+                type=SchemaFieldDataType(data_type),
+                nativeDataType=metric_type or "",
+                description=metric.get("description", ""),
+                nullable=True,
+            )
+            schema_fields.append(field)
+        return schema_fields
+
     def gen_schema_metadata(
         self,
         dataset_response: dict,
     ) -> SchemaMetadata:
         dataset_response = dataset_response.get("result", {})
         column_data = dataset_response.get("columns", [])
+        metric_data = dataset_response.get("metrics", [])
+
+        column_fields = self.gen_schema_fields(column_data)
+        metric_fields = self.gen_metric_schema_fields(metric_data)
+
         schema_metadata = SchemaMetadata(
             schemaName=dataset_response.get("table_name", ""),
             platform=make_data_platform_urn(self.platform),
             version=0,
             hash="",
             platformSchema=MySqlDDL(tableSchema=""),
-            fields=self.gen_schema_fields(column_data),
+            fields=column_fields + metric_fields,
         )
         return schema_metadata
 
@@ -1049,6 +1086,8 @@ class SupersetSource(StatefulIngestionSourceBase):
         # To generate column level lineage, we can manually decode the metadata
         # to produce the ColumnLineageInfo
         columns = dataset_response.get("result", {}).get("columns", [])
+        metrics = dataset_response.get("result", {}).get("metrics", [])
+
         fine_grained_lineages: List[FineGrainedLineageClass] = []
 
         for column in columns:
@@ -1067,6 +1106,22 @@ class SupersetSource(StatefulIngestionSourceBase):
                 )
             )
 
+        for metric in metrics:
+            metric_name = metric.get("metric_name", "")
+            if not metric_name:
+                continue
+
+            downstream = [make_schema_field_urn(datasource_urn, metric_name)]
+            upstreams = [make_schema_field_urn(upstream_dataset, metric_name)]
+            fine_grained_lineages.append(
+                FineGrainedLineageClass(
+                    downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
+                    downstreams=downstream,
+                    upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+                    upstreams=upstreams,
+                )
+            )
+
         upstream_lineage = UpstreamLineageClass(
             upstreams=[
                 UpstreamClass(
@@ -1087,7 +1142,7 @@ class SupersetSource(StatefulIngestionSourceBase):
         datasource_urn = self.get_datasource_urn_from_id(
             dataset_response, self.platform
         )
-        dataset_url = f"{self.config.display_uri}{dataset_response.get('result', {}).get('url', '')}"
+        dataset_url = f"{self.config.display_uri}/explore/?datasource_type=table&datasource_id={dataset.id}"
 
         modified_actor = f"urn:li:corpuser:{self.owner_info.get((dataset_data.get('changed_by') or {}).get('id', -1), 'unknown')}"
         now = datetime.now().strftime("%I:%M%p on %B %d, %Y")

datahub/ingestion/source/tableau/tableau.py

@@ -80,6 +80,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     BIContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -148,7 +149,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
 )
 from datahub.metadata.schema_classes import (
     BrowsePathsClass,
-    ChangeTypeClass,
     ChartInfoClass,
     ChartUsageStatisticsClass,
     DashboardInfoClass,
@@ -528,6 +528,14 @@ class TableauConfig(
         default=False,
         description="Ingest details for tables external to (not embedded in) tableau as entities.",
     )
+    emit_all_published_datasources: bool = Field(
+        default=False,
+        description="Ingest all published data sources. When False (default), only ingest published data sources that belong to an ingested workbook.",
+    )
+    emit_all_embedded_datasources: bool = Field(
+        default=False,
+        description="Ingest all embedded data sources. When False (default), only ingest embedded data sources that belong to an ingested workbook.",
+    )
 
     env: str = Field(
         default=builder.DEFAULT_ENV,
@@ -861,16 +869,29 @@ def report_user_role(report: TableauSourceReport, server: Server) -> None:
 @platform_name("Tableau")
 @config_class(TableauConfig)
 @support_status(SupportStatus.CERTIFIED)
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.TABLEAU_PROJECT,
+        SourceCapabilityModifier.TABLEAU_SITE,
+        SourceCapabilityModifier.TABLEAU_WORKBOOK,
+    ],
+)
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Requires transformer", supported=False)
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 @capability(
     SourceCapability.USAGE_STATS,
     "Dashboard/Chart view counts, enabled using extract_usage_stats config",
+    subtype_modifier=[
+        SourceCapabilityModifier.DASHBOARD,
+        SourceCapabilityModifier.CHART,
+    ],
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Enabled by default when stateful ingestion is turned on.",
+    "Enabled by default via stateful ingestion.",
 )
 @capability(SourceCapability.OWNERSHIP, "Requires recipe configuration")
 @capability(SourceCapability.TAGS, "Requires recipe configuration")
@@ -879,6 +900,7 @@ def report_user_role(report: TableauSourceReport, server: Server) -> None:
     SourceCapability.LINEAGE_FINE,
     "Enabled by default, configure using `extract_column_level_lineage`",
 )
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class TableauSource(StatefulIngestionSourceBase, TestableSource):
     platform = "tableau"
 
@@ -2174,32 +2196,32 @@ class TableauSiteSource:
                     else []
                 )
 
-                # The Tableau SQL parser much worse than our sqlglot based parser,
-                # so relying on metadata parsed by Tableau from SQL queries can be
-                # less accurate. This option allows us to ignore Tableau's parser and
-                # only use our own.
-                if self.config.force_extraction_of_lineage_from_custom_sql_queries:
-                    logger.debug("Extracting TLL & CLL from custom sql (forced)")
+                tableau_table_list = csql.get(c.TABLES, [])
+                if self.config.force_extraction_of_lineage_from_custom_sql_queries or (
+                    not tableau_table_list
+                    and self.config.extract_lineage_from_unsupported_custom_sql_queries
+                ):
+                    if not tableau_table_list:
+                        # custom sql tables may contain unsupported sql, causing incomplete lineage
+                        # we extract the lineage from the raw queries
+                        logger.debug(
+                            "Parsing TLL & CLL from custom sql (tableau metadata incomplete)"
+                        )
+                    else:
+                        # The Tableau SQL parser is much worse than our sqlglot based parser,
+                        # so relying on metadata parsed by Tableau from SQL queries can be
+                        # less accurate. This option allows us to ignore Tableau's parser and
+                        # only use our own.
+                        logger.debug("Parsing TLL & CLL from custom sql (forced)")
+
                     yield from self._create_lineage_from_unsupported_csql(
                         csql_urn, csql, columns
                     )
                 else:
-                    tables = csql.get(c.TABLES, [])
-
-                    if tables:
-                        # lineage from custom sql -> datasets/tables #
-                        yield from self._create_lineage_to_upstream_tables(
-                            csql_urn, tables, datasource
-                        )
-                    elif (
-                        self.config.extract_lineage_from_unsupported_custom_sql_queries
-                    ):
-                        logger.debug("Extracting TLL & CLL from custom sql")
-                        # custom sql tables may contain unsupported sql, causing incomplete lineage
-                        # we extract the lineage from the raw queries
-                        yield from self._create_lineage_from_unsupported_csql(
-                            csql_urn, csql, columns
-                        )
+                    # lineage from custom sql -> datasets/tables #
+                    yield from self._create_lineage_to_upstream_tables(
+                        csql_urn, tableau_table_list, datasource
+                    )
 
             #  Schema Metadata
             schema_metadata = self.get_schema_metadata_for_custom_sql(columns)
@@ -2237,7 +2259,6 @@ class TableauSiteSource:
             yield self.get_metadata_change_event(dataset_snapshot)
             yield self.get_metadata_change_proposal(
                 dataset_snapshot.urn,
-                aspect_name=c.SUB_TYPES,
                 aspect=SubTypesClass(typeNames=[DatasetSubTypes.VIEW, c.CUSTOM_SQL]),
             )
 
@@ -2402,7 +2423,6 @@ class TableauSiteSource:
             upstream_lineage = UpstreamLineage(upstreams=upstream_tables)
             yield self.get_metadata_change_proposal(
                 csql_urn,
-                aspect_name=c.UPSTREAM_LINEAGE,
                 aspect=upstream_lineage,
             )
             self.report.num_tables_with_upstream_lineage += 1
@@ -2588,7 +2608,6 @@ class TableauSiteSource:
         )
         yield self.get_metadata_change_proposal(
             csql_urn,
-            aspect_name=c.UPSTREAM_LINEAGE,
             aspect=upstream_lineage,
         )
         self.report.num_tables_with_upstream_lineage += 1
@@ -2634,14 +2653,10 @@ class TableauSiteSource:
     def get_metadata_change_proposal(
         self,
         urn: str,
-        aspect_name: str,
         aspect: Union["UpstreamLineage", "SubTypesClass"],
     ) -> MetadataWorkUnit:
         return MetadataChangeProposalWrapper(
-            entityType=c.DATASET,
-            changeType=ChangeTypeClass.UPSERT,
             entityUrn=urn,
-            aspectName=aspect_name,
             aspect=aspect,
         ).as_workunit()
 
@@ -2749,7 +2764,6 @@ class TableauSiteSource:
                 )
                 yield self.get_metadata_change_proposal(
                     datasource_urn,
-                    aspect_name=c.UPSTREAM_LINEAGE,
                     aspect=upstream_lineage,
                 )
                 self.report.num_tables_with_upstream_lineage += 1
@@ -2768,7 +2782,6 @@ class TableauSiteSource:
         yield self.get_metadata_change_event(dataset_snapshot)
         yield self.get_metadata_change_proposal(
             dataset_snapshot.urn,
-            aspect_name=c.SUB_TYPES,
             aspect=SubTypesClass(
                 typeNames=(
                     ["Embedded Data Source"]
@@ -2854,7 +2867,11 @@ class TableauSiteSource:
         return datasource
 
     def emit_published_datasources(self) -> Iterable[MetadataWorkUnit]:
-        datasource_filter = {c.ID_WITH_IN: self.datasource_ids_being_used}
+        datasource_filter = (
+            {}
+            if self.config.emit_all_published_datasources
+            else {c.ID_WITH_IN: self.datasource_ids_being_used}
+        )
 
         for datasource in self.get_connection_objects(
             query=published_datasource_graphql_query,
@@ -3547,7 +3564,11 @@ class TableauSiteSource:
         return browse_paths
 
     def emit_embedded_datasources(self) -> Iterable[MetadataWorkUnit]:
-        datasource_filter = {c.ID_WITH_IN: self.embedded_datasource_ids_being_used}
+        datasource_filter = (
+            {}
+            if self.config.emit_all_embedded_datasources
+            else {c.ID_WITH_IN: self.embedded_datasource_ids_being_used}
+        )
 
         for datasource in self.get_connection_objects(
             query=embedded_datasource_graphql_query,
@@ -3659,7 +3680,7 @@ class TableauSiteSource:
                 container_key=project_key,
                 name=project_.name,
                 description=project_.description,
-                sub_types=[c.PROJECT],
+                sub_types=[BIContainerSubTypes.TABLEAU_PROJECT],
                 parent_container_key=parent_project_key,
             )
 
@@ -3677,7 +3698,7 @@ class TableauSiteSource:
         yield from gen_containers(
             container_key=self.gen_site_key(self.site_id),
             name=self.site.name or "Default",
-            sub_types=[c.SITE],
+            sub_types=[BIContainerSubTypes.TABLEAU_SITE],
         )
 
     def _fetch_groups(self):

datahub/ingestion/source/tableau/tableau_common.py

@@ -579,10 +579,12 @@ def get_platform(connection_type: str) -> str:
         platform = "oracle"
     elif connection_type in ("tbio", "teradata"):
         platform = "teradata"
-    elif connection_type in ("sqlserver"):
+    elif connection_type in ("sqlserver",):
         platform = "mssql"
-    elif connection_type in ("athena"):
+    elif connection_type in ("athena",):
         platform = "athena"
+    elif connection_type in ("googlebigquery",):
+        platform = "bigquery"
     elif connection_type.endswith("_jdbc"):
         # e.g. convert trino_jdbc -> trino
         platform = connection_type[: -len("_jdbc")]

datahub/ingestion/source/tableau/tableau_constant.py

@@ -50,7 +50,6 @@ TABLES = "tables"
 DESCRIPTION = "description"
 SQL = "SQL"
 QUERY = "query"
-SUB_TYPES = "subTypes"
 VIEW = "view"
 CUSTOM_SQL = "Custom SQL"
 REMOTE_TYPE = "remoteType"
@@ -58,7 +57,6 @@ UNKNOWN = "UNKNOWN"
 PUBLISHED_DATA_SOURCE = "PublishedDatasource"
 LUID = "luid"
 EMBEDDED_DATA_SOURCE = "EmbeddedDatasource"
-UPSTREAM_LINEAGE = "upstreamLineage"
 OWNER = "owner"
 USERNAME = "username"
 HAS_EXTRACTS = "hasExtracts"
@@ -78,8 +76,6 @@ CHART = "chart"
 DASHBOARD = "dashboard"
 DASHBOARDS_CONNECTION = "dashboardsConnection"
 EMBEDDED_DATA_SOURCES_CONNECTION = "embeddedDatasourcesConnection"
-PROJECT = "Project"
-SITE = "Site"
 IS_UNSUPPORTED_CUSTOM_SQL = "isUnsupportedCustomSql"
 SITE_PERMISSION = "sitePermission"
 ROLE_SITE_ADMIN_EXPLORER = "SiteAdministratorExplorer"

datahub/ingestion/source/unity/config.py

@@ -229,6 +229,11 @@ class UnityCatalogSourceConfig(
         description="Option to enable/disable ownership generation for metastores, catalogs, schemas, and tables.",
     )
 
+    include_tags: bool = pydantic.Field(
+        default=True,
+        description="Option to enable/disable column/table tag extraction.",
+    )
+
     _rename_table_ownership = pydantic_renamed_field(
         "include_table_ownership", "include_ownership"
     )

datahub/ingestion/source/unity/proxy.py

@@ -8,6 +8,8 @@ from datetime import datetime
 from typing import Any, Dict, Iterable, List, Optional, Union, cast
 from unittest.mock import patch
 
+import cachetools
+from cachetools import cached
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import (
     CatalogInfo,
@@ -25,8 +27,10 @@ from databricks.sdk.service.sql import (
     QueryStatus,
 )
 from databricks.sdk.service.workspace import ObjectType
+from databricks.sql import connect
 
 from datahub._version import nice_version_name
+from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.ingestion.source.unity.hive_metastore_proxy import HiveMetastoreProxy
 from datahub.ingestion.source.unity.proxy_profiling import (
@@ -108,6 +112,13 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         self.warehouse_id = warehouse_id or ""
         self.report = report
         self.hive_metastore_proxy = hive_metastore_proxy
+        self._sql_connection_params = {
+            "server_hostname": self._workspace_client.config.host.replace(
+                "https://", ""
+            ),
+            "http_path": f"/sql/1.0/warehouses/{self.warehouse_id}",
+            "access_token": self._workspace_client.config.token,
+        }
 
     def check_basic_connectivity(self) -> bool:
         return bool(self._workspace_client.catalogs.list(include_browse=True))
@@ -492,3 +503,110 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             executed_as_user_id=info.executed_as_user_id,
             executed_as_user_name=info.executed_as_user_name,
         )
+
+    def _execute_sql_query(self, query: str) -> List[List[str]]:
+        """Execute SQL query using databricks-sql connector for better performance"""
+        try:
+            with (
+                connect(**self._sql_connection_params) as connection,
+                connection.cursor() as cursor,
+            ):
+                cursor.execute(query)
+                return cursor.fetchall()
+
+        except Exception as e:
+            logger.warning(f"Failed to execute SQL query: {e}")
+            return []
+
+    @cached(cachetools.FIFOCache(maxsize=100))
+    def get_schema_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching schema tags for catalog: {catalog}")
+
+        query = f"SELECT * FROM {catalog}.information_schema.schema_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, schema_name, tag_name, tag_value = row
+            schema_key = f"{catalog_name}.{schema_name}"
+
+            if schema_key not in result_dict:
+                result_dict[schema_key] = []
+
+            result_dict[schema_key].append(
+                UnityCatalogTag(key=tag_name, value=tag_value)
+            )
+
+        return result_dict
+
+    @cached(cachetools.FIFOCache(maxsize=100))
+    def get_catalog_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching table tags for catalog: {catalog}")
+
+        query = f"SELECT * FROM {catalog}.information_schema.catalog_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, tag_name, tag_value = row
+
+            if catalog_name not in result_dict:
+                result_dict[catalog_name] = []
+
+            result_dict[catalog_name].append(
+                UnityCatalogTag(key=tag_name, value=tag_value)
+            )
+
+        return result_dict
+
+    @cached(cachetools.FIFOCache(maxsize=100))
+    def get_table_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching table tags for catalog: {catalog}")
+
+        query = f"SELECT * FROM {catalog}.information_schema.table_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, schema_name, table_name, tag_name, tag_value = row
+            table_key = f"{catalog_name}.{schema_name}.{table_name}"
+
+            if table_key not in result_dict:
+                result_dict[table_key] = []
+
+            result_dict[table_key].append(
+                UnityCatalogTag(key=tag_name, value=tag_value if tag_value else None)
+            )
+
+        return result_dict
+
+    @cached(cachetools.FIFOCache(maxsize=100))
+    def get_column_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching column tags for catalog: {catalog}")
+
+        query = f"SELECT * FROM {catalog}.information_schema.column_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, schema_name, table_name, column_name, tag_name, tag_value = (
+                row
+            )
+            column_key = f"{catalog_name}.{schema_name}.{table_name}.{column_name}"
+
+            if column_key not in result_dict:
+                result_dict[column_key] = []
+
+            result_dict[column_key].append(
+                UnityCatalogTag(key=tag_name, value=tag_value if tag_value else None)
+            )
+
+        return result_dict