acryl-datahub 1.2.0.1__py3-none-any.whl → 1.2.0.2__py3-none-any.whl

datahub/ingestion/source/s3/source.py

@@ -41,7 +41,11 @@ from datahub.ingestion.source.aws.s3_util import (
     get_key_prefix,
     strip_s3_prefix,
 )
-from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
+from datahub.ingestion.source.data_lake_common.data_lake_utils import (
+    ContainerWUCreator,
+    add_partition_columns_to_schema,
+)
 from datahub.ingestion.source.data_lake_common.object_store import (
     create_object_store_adapter,
 )
@@ -58,9 +62,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import TimeStamp
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
-    SchemaField,
     SchemaMetadata,
-    StringTypeClass,
 )
 from datahub.metadata.schema_classes import (
     DataPlatformInstanceClass,
@@ -70,7 +72,6 @@ from datahub.metadata.schema_classes import (
     OtherSchemaClass,
     PartitionsSummaryClass,
     PartitionSummaryClass,
-    SchemaFieldDataTypeClass,
     _Aspect,
 )
 from datahub.telemetry import stats, telemetry
@@ -196,7 +197,14 @@ class TableData:
 @platform_name("S3 / Local Files", id="s3")
 @config_class(DataLakeSourceConfig)
 @support_status(SupportStatus.INCUBATING)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.FOLDER,
+        SourceCapabilityModifier.S3_BUCKET,
+    ],
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(
     SourceCapability.SCHEMA_METADATA, "Can infer schema from supported file types"
@@ -474,7 +482,7 @@ class S3Source(StatefulIngestionSourceBase):
             fields = sorted(fields, key=lambda f: f.fieldPath)
 
         if self.source_config.add_partition_columns_to_schema and table_data.partitions:
-            self.add_partition_columns_to_schema(
+            add_partition_columns_to_schema(
                 fields=fields, path_spec=path_spec, full_path=table_data.full_path
             )
 
@@ -510,34 +518,6 @@ class S3Source(StatefulIngestionSourceBase):
         else:
             return None
 
-    def add_partition_columns_to_schema(
-        self, path_spec: PathSpec, full_path: str, fields: List[SchemaField]
-    ) -> None:
-        is_fieldpath_v2 = False
-        for field in fields:
-            if field.fieldPath.startswith("[version=2.0]"):
-                is_fieldpath_v2 = True
-                break
-        partition_keys = path_spec.get_partition_from_path(full_path)
-        if not partition_keys:
-            return None
-
-        for partition_key in partition_keys:
-            fields.append(
-                SchemaField(
-                    fieldPath=(
-                        f"{partition_key[0]}"
-                        if not is_fieldpath_v2
-                        else f"[version=2.0].[type=string].{partition_key[0]}"
-                    ),
-                    nativeDataType="string",
-                    type=SchemaFieldDataTypeClass(StringTypeClass()),
-                    isPartitioningKey=True,
-                    nullable=True,
-                    recursive=False,
-                )
-            )
-
     def get_table_profile(
         self, table_data: TableData, dataset_urn: str
     ) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/sql/athena.py

@@ -347,12 +347,18 @@ class Partitionitem:
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Supported for S3 tables",
-    subtype_modifier=[SourceCapabilityModifier.TABLE],
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
     "Supported for S3 tables",
-    subtype_modifier=[SourceCapabilityModifier.TABLE],
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
 )
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 class AthenaSource(SQLAlchemySource):

datahub/ingestion/source/sql/clickhouse.py

@@ -32,6 +32,7 @@ from datahub.ingestion.api.decorators import (
     support_status,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.sql_common import (
     SqlWorkUnit,
     logger,
@@ -383,6 +384,14 @@ clickhouse_datetime_format = "%Y-%m-%d %H:%M:%S"
     SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Enabled by default to get lineage for views via `include_view_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
+)
 class ClickHouseSource(TwoTierSQLAlchemySource):
     """
     This plugin extracts the following:

datahub/ingestion/source/sql/postgres.py

@@ -36,6 +36,15 @@ from datahub.ingestion.source.sql.sql_common import (
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
+from datahub.ingestion.source.sql.sql_utils import (
+    gen_database_key,
+    gen_schema_key,
+)
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    BaseProcedure,
+    generate_procedure_container_workunits,
+    generate_procedure_workunits,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     ArrayTypeClass,
     BytesTypeClass,
@@ -123,6 +132,15 @@ class PostgresConfig(BasePostgresConfig):
             "Note: this is not used if `database` or `sqlalchemy_uri` are provided."
         ),
     )
+    include_stored_procedures: bool = Field(
+        default=True,
+        description="Include ingest of stored procedures.",
+    )
+    procedure_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for stored procedures to filter in ingestion."
+        "Specify regex to match the entire procedure name in database.schema.procedure_name format. e.g. to match all procedures starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+    )
 
 
 @platform_name("Postgres")
@@ -135,7 +153,7 @@ class PostgresSource(SQLAlchemySource):
     """
     This plugin extracts the following:
 
-    - Metadata for databases, schemas, views, and tables
+    - Metadata for databases, schemas, views, tables, and stored procedures
     - Column types associated with each table
     - Also supports PostGIS extensions
     - Table, row, and column statistics via optional SQL profiling
@@ -291,3 +309,174 @@ class PostgresSource(SQLAlchemySource):
                     ] = row.table_size
         except Exception as e:
             logger.error(f"failed to fetch profile metadata: {e}")
+
+    def get_schema_level_workunits(
+        self,
+        inspector: Inspector,
+        schema: str,
+        database: str,
+    ) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        yield from super().get_schema_level_workunits(
+            inspector=inspector,
+            schema=schema,
+            database=database,
+        )
+
+        if self.config.include_stored_procedures:
+            try:
+                yield from self.loop_stored_procedures(inspector, schema, self.config)
+            except Exception as e:
+                self.report.failure(
+                    title="Failed to list stored procedures for schema",
+                    message="An error occurred while listing procedures for the schema.",
+                    context=f"{database}.{schema}",
+                    exc=e,
+                )
+
+    def loop_stored_procedures(
+        self,
+        inspector: Inspector,
+        schema: str,
+        config: PostgresConfig,
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Loop schema data for get stored procedures as dataJob-s.
+        """
+        db_name = self.get_db_name(inspector)
+
+        procedures = self.fetch_procedures_for_schema(inspector, schema, db_name)
+        if procedures:
+            yield from self._process_procedures(procedures, db_name, schema)
+
+    def fetch_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        try:
+            raw_procedures: List[BaseProcedure] = self.get_procedures_for_schema(
+                inspector, schema, db_name
+            )
+            procedures: List[BaseProcedure] = []
+            for procedure in raw_procedures:
+                procedure_qualified_name = self.get_identifier(
+                    schema=schema,
+                    entity=procedure.name,
+                    inspector=inspector,
+                )
+
+                if not self.config.procedure_pattern.allowed(procedure_qualified_name):
+                    self.report.report_dropped(procedure_qualified_name)
+                else:
+                    procedures.append(procedure)
+            return procedures
+        except Exception as e:
+            self.report.warning(
+                title="Failed to get procedures for schema",
+                message="An error occurred while fetching procedures for the schema.",
+                context=f"{db_name}.{schema}",
+                exc=e,
+            )
+            return []
+
+    def get_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        """
+        Get stored procedures for a specific schema.
+        """
+        base_procedures = []
+        with inspector.engine.connect() as conn:
+            procedures = conn.execute(
+                """
+                    SELECT
+                        p.proname AS name,
+                        l.lanname AS language,
+                        pg_get_function_arguments(p.oid) AS arguments,
+                        pg_get_functiondef(p.oid) AS definition,
+                        obj_description(p.oid, 'pg_proc') AS comment
+                    FROM
+                        pg_proc p
+                    JOIN
+                        pg_namespace n ON n.oid = p.pronamespace
+                    JOIN
+                        pg_language l ON l.oid = p.prolang
+                    WHERE
+                        p.prokind = 'p'
+                        AND n.nspname = '"""
+                + schema
+                + """';
+                    """
+            )
+
+            procedure_rows = list(procedures)
+            for row in procedure_rows:
+                base_procedures.append(
+                    BaseProcedure(
+                        name=row.name,
+                        language=row.language,
+                        argument_signature=row.arguments,
+                        return_type=None,
+                        procedure_definition=row.definition,
+                        created=None,
+                        last_altered=None,
+                        comment=row.comment,
+                        extra_properties=None,
+                    )
+                )
+            return base_procedures
+
+    def _process_procedures(
+        self,
+        procedures: List[BaseProcedure],
+        db_name: str,
+        schema: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if procedures:
+            yield from generate_procedure_container_workunits(
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+            )
+        for procedure in procedures:
+            yield from self._process_procedure(procedure, schema, db_name)
+
+    def _process_procedure(
+        self,
+        procedure: BaseProcedure,
+        schema: str,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        try:
+            yield from generate_procedure_workunits(
+                procedure=procedure,
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_resolver=self.get_schema_resolver(),
+            )
+        except Exception as e:
+            self.report.warning(
+                title="Failed to emit stored procedure",
+                message="An error occurred while emitting stored procedure",
+                context=procedure.name,
+                exc=e,
+            )

datahub/ingestion/source/sql_queries.py

@@ -4,7 +4,7 @@ import os
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import partial
-from typing import Iterable, List, Optional, Set
+from typing import Iterable, List, Optional, Union
 
 from pydantic import Field
 
@@ -14,9 +14,8 @@ from datahub.configuration.source_common import (
 )
 from datahub.emitter.mce_builder import (
     make_dataset_urn_with_platform_instance,
-    make_user_urn,
 )
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -35,8 +34,15 @@ from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
+from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
+from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
-from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
+from datahub.sql_parsing.sql_parsing_aggregator import (
+    KnownQueryLineageInfo,
+    ObservedQuery,
+    SqlAggregatorReport,
+    SqlParsingAggregator,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -72,23 +78,13 @@ class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
     )
 
 
-class SqlQueriesSourceReport(SourceReport):
-    num_queries_parsed: int = 0
-    num_table_parse_failures: int = 0
-    num_column_parse_failures: int = 0
+@dataclass
+class SqlQueriesSourceReport(SourceReport, IngestionStageReport):
+    num_entries_processed: int = 0
+    num_entries_failed: int = 0
+    num_queries_aggregator_failures: int = 0
 
-    def compute_stats(self) -> None:
-        super().compute_stats()
-        self.table_failure_rate = (
-            f"{self.num_table_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
-        self.column_failure_rate = (
-            f"{self.num_column_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
+    sql_aggregator: Optional[SqlAggregatorReport] = None
 
 
 @platform_name("SQL Queries")
@@ -107,15 +103,15 @@ class SqlQueriesSource(Source):
     - user (optional): string - The user who ran the query.
     This user value will be directly converted into a DataHub user urn.
     - operation_type (optional): string - Platform-specific operation type, used if the operation type can't be parsed.
+    - session_id (optional): string - Session identifier for temporary table resolution across queries.
     - downstream_tables (optional): string[] - Fallback list of tables that the query writes to,
      used if the query can't be parsed.
     - upstream_tables (optional): string[] - Fallback list of tables the query reads from,
      used if the query can't be parsed.
     """
 
-    urns: Optional[Set[str]]
-    schema_resolver: SchemaResolver
-    builder: SqlParsingBuilder
+    schema_resolver: Optional[SchemaResolver]
+    aggregator: SqlParsingAggregator
 
     def __init__(self, ctx: PipelineContext, config: SqlQueriesSourceConfig):
         if not ctx.graph:
@@ -128,22 +124,36 @@ class SqlQueriesSource(Source):
         self.config = config
         self.report = SqlQueriesSourceReport()
 
-        self.builder = SqlParsingBuilder(usage_config=self.config.usage)
-
         if self.config.use_schema_resolver:
+            # TODO: `initialize_schema_resolver_from_datahub` does a  bulk initialization by fetching all schemas
+            # for the given platform, platform instance, and env. Instead this should be configurable:
+            # bulk initialization vs lazy on-demand schema fetching.
             self.schema_resolver = self.graph.initialize_schema_resolver_from_datahub(
                 platform=self.config.platform,
                 platform_instance=self.config.platform_instance,
                 env=self.config.env,
             )
-            self.urns = self.schema_resolver.get_urns()
         else:
-            self.schema_resolver = self.graph._make_schema_resolver(
-                platform=self.config.platform,
-                platform_instance=self.config.platform_instance,
-                env=self.config.env,
-            )
-            self.urns = None
+            self.schema_resolver = None
+
+        self.aggregator = SqlParsingAggregator(
+            platform=self.config.platform,
+            platform_instance=self.config.platform_instance,
+            env=self.config.env,
+            schema_resolver=self.schema_resolver,
+            eager_graph_load=False,
+            generate_lineage=True,  # TODO: make this configurable
+            generate_queries=True,  # TODO: make this configurable
+            generate_query_subject_fields=True,  # TODO: make this configurable
+            generate_query_usage_statistics=True,  # This enables publishing SELECT query entities, otherwise only mutation queries are published
+            generate_usage_statistics=True,
+            generate_operations=True,  # TODO: make this configurable
+            usage_config=self.config.usage,
+            is_temp_table=None,
+            is_allowed_table=None,
+            format_queries=False,
+        )
+        self.report.sql_aggregator = self.aggregator.report
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "SqlQueriesSource":
@@ -156,68 +166,92 @@ class SqlQueriesSource(Source):
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [partial(auto_workunit_reporter, self.get_report())]
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         logger.info(f"Parsing queries from {os.path.basename(self.config.query_file)}")
+
+        with self.report.new_stage("Collecting queries from file"):
+            queries = list(self._parse_query_file())
+            logger.info(f"Collected {len(queries)} queries for processing")
+
+        with self.report.new_stage("Processing queries through SQL parsing aggregator"):
+            for query_entry in queries:
+                self._add_query_to_aggregator(query_entry)
+
+        with self.report.new_stage("Generating metadata work units"):
+            logger.info("Generating workunits from SQL parsing aggregator")
+            yield from self.aggregator.gen_metadata()
+
+    def _parse_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse the query file and yield QueryEntry objects."""
         with open(self.config.query_file) as f:
             for line in f:
                 try:
                     query_dict = json.loads(line, strict=False)
                     entry = QueryEntry.create(query_dict, config=self.config)
-                    yield from self._process_query(entry)
+                    self.report.num_entries_processed += 1
+                    if self.report.num_entries_processed % 1000 == 0:
+                        logger.info(
+                            f"Processed {self.report.num_entries_processed} query entries"
+                        )
+                    yield entry
                 except Exception as e:
-                    logger.warning("Error processing query", exc_info=True)
-                    self.report.report_warning("process-query", str(e))
-
-        logger.info("Generating workunits")
-        yield from self.builder.gen_workunits()
+                    self.report.num_entries_failed += 1
+                    self.report.warning(
+                        title="Error processing query",
+                        message="Query skipped due to parsing error",
+                        context=line.strip(),
+                        exc=e,
+                    )
 
-    def _process_query(self, entry: "QueryEntry") -> Iterable[MetadataWorkUnit]:
-        self.report.num_queries_parsed += 1
-        if self.report.num_queries_parsed % 1000 == 0:
-            logger.info(f"Parsed {self.report.num_queries_parsed} queries")
-
-        result = sqlglot_lineage(
-            sql=entry.query,
-            schema_resolver=self.schema_resolver,
-            default_db=self.config.default_db,
-            default_schema=self.config.default_schema,
-            override_dialect=self.config.override_dialect,
-        )
-        if result.debug_info.table_error:
-            logger.info(f"Error parsing table lineage, {result.debug_info.table_error}")
-            self.report.num_table_parse_failures += 1
-            for downstream_urn in set(entry.downstream_tables):
-                self.builder.add_lineage(
-                    downstream_urn=downstream_urn,
-                    upstream_urns=entry.upstream_tables,
-                    timestamp=entry.timestamp,
-                    user=entry.user,
+    def _add_query_to_aggregator(self, query_entry: "QueryEntry") -> None:
+        """Add a query to the SQL parsing aggregator."""
+        try:
+            # If we have explicit lineage, use it directly
+            if query_entry.upstream_tables or query_entry.downstream_tables:
+                logger.debug("Using explicit lineage from query file")
+                for downstream_table in query_entry.downstream_tables:
+                    known_lineage = KnownQueryLineageInfo(
+                        query_text=query_entry.query,
+                        downstream=downstream_table,
+                        upstreams=query_entry.upstream_tables,
+                        timestamp=query_entry.timestamp,
+                        session_id=query_entry.session_id,
+                    )
+                    self.aggregator.add_known_query_lineage(known_lineage)
+            else:
+                # No explicit lineage, rely on parsing
+                observed_query = ObservedQuery(
+                    query=query_entry.query,
+                    timestamp=query_entry.timestamp,
+                    user=query_entry.user,
+                    session_id=query_entry.session_id,
+                    default_db=self.config.default_db,
+                    default_schema=self.config.default_schema,
+                    override_dialect=self.config.override_dialect,
                 )
-            return
-        elif result.debug_info.column_error:
-            logger.debug(
-                f"Error parsing column lineage, {result.debug_info.column_error}"
+                self.aggregator.add_observed_query(observed_query)
+
+        except Exception as e:
+            self.report.num_queries_aggregator_failures += 1
+            self.report.warning(
+                title="Error adding query to aggregator",
+                message="Query skipped due to failure when adding query to SQL parsing aggregator",
+                context=query_entry.query,
+                exc=e,
             )
-            self.report.num_column_parse_failures += 1
-
-        yield from self.builder.process_sql_parsing_result(
-            result,
-            query=entry.query,
-            query_timestamp=entry.timestamp,
-            user=entry.user,
-            custom_operation_type=entry.operation_type,
-            include_urns=self.urns,
-        )
 
 
 @dataclass
 class QueryEntry:
     query: str
     timestamp: Optional[datetime]
-    user: Optional[str]
+    user: Optional[CorpUserUrn]
     operation_type: Optional[str]
     downstream_tables: List[str]
     upstream_tables: List[str]
+    session_id: Optional[str] = None
 
     @classmethod
     def create(
@@ -230,7 +264,7 @@ class QueryEntry:
                 if "timestamp" in entry_dict
                 else None
             ),
-            user=make_user_urn(entry_dict["user"]) if "user" in entry_dict else None,
+            user=CorpUserUrn(entry_dict["user"]) if "user" in entry_dict else None,
             operation_type=entry_dict.get("operation_type"),
             downstream_tables=[
                 make_dataset_urn_with_platform_instance(
@@ -250,4 +284,5 @@ class QueryEntry:
                 )
                 for table in entry_dict.get("upstream_tables", [])
             ],
+            session_id=entry_dict.get("session_id"),
         )

datahub/ingestion/source/unity/proxy.py

@@ -521,9 +521,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @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}")
+        logger.info(f"Fetching schema tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.schema_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.schema_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -544,9 +544,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @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}")
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.catalog_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.catalog_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -566,9 +566,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @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}")
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.table_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.table_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -589,9 +589,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @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}")
+        logger.info(f"Fetching column tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.column_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.column_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}