openmetadata-ingestion 1.2.5.1__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

metadata/ingestion/source/storage/storage_service.py

@@ -12,9 +12,7 @@
 Base class for ingesting Object Storage services
 """
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Optional
-
-from pandas import DataFrame
+from typing import Any, Iterable, List, Optional, Set
 
 from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 from metadata.generated.schema.entity.data.container import Container
@@ -35,14 +33,16 @@ from metadata.generated.schema.metadataIngestion.storageServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
+from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
@@ -50,7 +50,11 @@ from metadata.ingestion.source.database.glue.models import Column
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.readers.dataframe.reader_factory import SupportedTypes
 from metadata.readers.models import ConfigSource
-from metadata.utils.datalake.datalake_utils import fetch_dataframe, get_columns
+from metadata.utils import fqn
+from metadata.utils.datalake.datalake_utils import (
+    DataFrameColumnParser,
+    fetch_dataframe,
+)
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.storage_metadata_config import (
     StorageMetadataConfigException,
@@ -64,7 +68,6 @@ OPENMETADATA_TEMPLATE_FILE_NAME = "openmetadata.json"
 
 
 class StorageServiceTopology(ServiceTopology):
-
     root = TopologyNode(
         producer="get_services",
         stages=[
@@ -78,6 +81,7 @@ class StorageServiceTopology(ServiceTopology):
             ),
         ],
         children=["container"],
+        post_process=["mark_containers_as_deleted"],
     )
 
     container = TopologyNode(
@@ -108,7 +112,8 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: StorageConnection.__fields__["config"].type_
 
     topology = StorageServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
+    container_source_state: Set = set()
 
     global_manifest: Optional[ManifestMetadataConfig]
 
@@ -135,6 +140,10 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             ManifestMetadataConfig
         ] = self.get_manifest_file()
 
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
     def get_manifest_file(self) -> Optional[ManifestMetadataConfig]:
         if self.source_config.storageMetadataConfigSource and not isinstance(
             self.source_config.storageMetadataConfigSource,
@@ -167,10 +176,43 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     def prepare(self):
         """By default, nothing needs to be taken care of when loading the source"""
 
+    def register_record(self, container_request: CreateContainerRequest) -> None:
+        """
+        Mark the container record as scanned and update
+        the storage_source_state
+        """
+        parent_container = (
+            self.metadata.get_by_id(
+                entity=Container, entity_id=container_request.parent.id
+            ).fullyQualifiedName.__root__
+            if container_request.parent
+            else None
+        )
+        container_fqn = fqn.build(
+            self.metadata,
+            entity_type=Container,
+            service_name=self.context.objectstore_service,
+            parent_container=parent_container,
+            container_name=container_request.name.__root__,
+        )
+
+        self.container_source_state.add(container_fqn)
+
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
         test_connection_fn(self.metadata, self.connection_obj, self.service_connection)
 
+    def mark_containers_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
+        """Method to mark the containers as deleted"""
+        if self.source_config.markDeletedContainers:
+            yield from delete_entity_from_source(
+                metadata=self.metadata,
+                entity_type=Container,
+                entity_source_state=self.container_source_state,
+                mark_deleted_entity=self.source_config.markDeletedContainers,
+                params={"service": self.context.objectstore_service},
+            )
+
     def yield_create_request_objectstore_service(self, config: WorkflowSource):
         yield Either(
             right=self.metadata.get_create_service_from_source(
@@ -229,10 +271,10 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             ),
         )
         columns = []
-        if isinstance(data_structure_details, DataFrame):
-            columns = get_columns(data_structure_details)
-        if isinstance(data_structure_details, list) and data_structure_details:
-            columns = get_columns(data_structure_details[0])
+        column_parser = DataFrameColumnParser.create(
+            data_structure_details, SupportedTypes(metadata_entry.structureFormat)
+        )
+        columns = column_parser.get_columns()
         return columns
 
     def _get_columns(

metadata/ingestion/stage/table_usage.py

@@ -22,10 +22,13 @@ from typing import Iterable, List, Optional, Tuple
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.type.queryParserData import ParsedData, QueryParserData
 from metadata.generated.schema.type.tableUsageCount import TableUsageCount
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Stage
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.constants import UTF_8
@@ -62,6 +65,10 @@ class TableUsageStage(Stage):
         init_staging_dir(self.config.filename)
         self.wrote_something = False
 
+    @property
+    def name(self) -> str:
+        return "Table Usage"
+
     @classmethod
     def create(cls, config_dict: dict, metadata: OpenMetadata):
         config = TableStageConfig.parse_obj(config_dict)
@@ -153,7 +160,7 @@ class TableUsageStage(Stage):
                 left=StackTraceError(
                     name=table,
                     error=f"Error in staging record [{exc}]",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
         yield Either(right=table)

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -26,6 +26,9 @@ from metadata.generated.schema.entity.services.connections.database.databricksCo
 from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
     SnowflakeType,
 )
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
 from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SESSION_TAG_QUERY,
@@ -84,17 +87,21 @@ class SQAInterfaceMixin:
             )
 
     def set_catalog(self, session) -> None:
-        """Set catalog for the session. Right now only databricks requires it
+        """Set catalog for the session. Right now only databricks and unity catalog requires it
 
         Args:
             session (Session): sqa session object
         """
-        if isinstance(self.service_connection_config, DatabricksConnection):
-            bind = session.get_bind()
-            bind.execute(
-                "USE CATALOG %(catalog)s;",
-                {"catalog": self.service_connection_config.catalog},
-            ).first()
+        if not isinstance(
+            self.service_connection_config,
+            (UnityCatalogConnection, DatabricksConnection),
+        ):
+            return
+        bind = session.get_bind()
+        bind.execute(
+            "USE CATALOG %(catalog)s;",
+            {"catalog": self.service_connection_config.catalog},
+        ).first()
 
     def close(self):
         """close session"""

metadata/parsers/protobuf_parser.py

@@ -20,12 +20,14 @@ import sys
 import traceback
 from enum import Enum
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import grpc_tools.protoc
 from pydantic import BaseModel
+from pydantic.main import ModelMetaclass
 
-from metadata.generated.schema.type.schema import FieldModel
+from metadata.generated.schema.entity.data.table import Column, DataType
+from metadata.generated.schema.type.schema import DataTypeTopic, FieldModel
 from metadata.utils.helpers import snake_to_camel
 from metadata.utils.logger import ingestion_logger
 
@@ -37,12 +39,12 @@ class ProtobufDataTypes(Enum):
     Enum for Protobuf Datatypes
     """
 
-    ERROR = 0
+    UNKNOWN = 0
     DOUBLE = 1
     FLOAT = 2
     INT = 3, 4, 5, 13, 17, 18
     FIXED = 6, 7, 15, 16
-    TYPE_BOOL = 8
+    BOOLEAN = 8
     STRING = 9
     UNION = 10
     RECORD = 11
@@ -165,7 +167,9 @@ class ProtobufParser:
             )
         return None
 
-    def parse_protobuf_schema(self) -> Optional[List[FieldModel]]:
+    def parse_protobuf_schema(
+        self, cls: ModelMetaclass = FieldModel
+    ) -> Optional[List[Union[FieldModel, Column]]]:
         """
         Method to parse the protobuf schema
         """
@@ -177,10 +181,12 @@ class ProtobufParser:
             )
 
             field_models = [
-                FieldModel(
+                cls(
                     name=instance.DESCRIPTOR.name,
                     dataType="RECORD",
-                    children=self.get_protobuf_fields(instance.DESCRIPTOR.fields),
+                    children=self.get_protobuf_fields(
+                        instance.DESCRIPTOR.fields, cls=cls
+                    ),
                 )
             ]
 
@@ -196,7 +202,17 @@ class ProtobufParser:
             )
         return None
 
-    def get_protobuf_fields(self, fields) -> Optional[List[FieldModel]]:
+    def _get_field_type(self, type_: int, cls: ModelMetaclass = FieldModel) -> str:
+        if type_ > 18:
+            return DataType.UNKNOWN.value
+        data_type = ProtobufDataTypes(type_).name
+        if cls == Column and data_type == DataTypeTopic.FIXED.value:
+            return DataType.INT.value
+        return data_type
+
+    def get_protobuf_fields(
+        self, fields, cls: ModelMetaclass = FieldModel
+    ) -> Optional[List[Union[FieldModel, Column]]]:
         """
         Recursively convert the parsed schema into required models
         """
@@ -205,10 +221,12 @@ class ProtobufParser:
         for field in fields:
             try:
                 field_models.append(
-                    FieldModel(
+                    cls(
                         name=field.name,
-                        dataType=ProtobufDataTypes(field.type).name,
-                        children=self.get_protobuf_fields(field.message_type.fields)
+                        dataType=self._get_field_type(field.type, cls=cls),
+                        children=self.get_protobuf_fields(
+                            field.message_type.fields, cls=cls
+                        )
                         if field.type == 11
                         else None,
                     )

metadata/pii/processor.py

@@ -16,6 +16,9 @@ import traceback
 from typing import List, Optional, cast
 
 from metadata.generated.schema.entity.data.table import Column, TableData
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
@@ -28,7 +31,7 @@ from metadata.generated.schema.type.tagLabel import (
     TagLabel,
     TagSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step
 from metadata.ingestion.api.steps import Processor
@@ -65,6 +68,10 @@ class PIIProcessor(Processor):
         self._ner_scanner = None
         self.confidence_threshold = self.source_config.confidence
 
+    @property
+    def name(self) -> str:
+        return "PII Processor"
+
     @property
     def ner_scanner(self) -> NERScanner:
         """Load the NER Scanner only if called"""
@@ -170,7 +177,7 @@ class PIIProcessor(Processor):
                     StackTraceError(
                         name=record.table.fullyQualifiedName.__root__,
                         error=f"Error computing PII tags for [{column}] - [{err}]",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 

metadata/profiler/api/models.py

@@ -15,7 +15,10 @@ Return types for Profiler workflow execution.
 We need to define this class as we end up having
 multiple profilers per table and columns.
 """
-from typing import List, Optional, Union
+from typing import List, Optional, Type, Union
+
+from sqlalchemy import Column
+from sqlalchemy.orm import DeclarativeMeta
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createTableProfile import (
@@ -31,9 +34,12 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
     SampleDataStorageConfig,
 )
+from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.ingestion.models.table_metadata import ColumnTag
+from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.processor.models import ProfilerDef
+from metadata.utils.sqa_like_column import SQALikeColumn
 
 
 class ColumnConfig(ConfigModel):
@@ -113,3 +119,15 @@ class ProfilerResponse(ConfigModel):
     def __str__(self):
         """Return the table name being processed"""
         return f"Table [{self.table.name.__root__}]"
+
+
+class ThreadPoolMetrics(ConfigModel):
+    """thread pool metric"""
+
+    metrics: Union[List[Union[Type[Metric], CustomMetric]], Type[Metric]]
+    metric_type: MetricTypes
+    column: Optional[Union[Column, SQALikeColumn]]
+    table: Union[Table, DeclarativeMeta]
+
+    class Config:
+        arbitrary_types_allowed = True

metadata/profiler/interface/pandas/profiler_interface.py

@@ -22,17 +22,22 @@ from typing import Dict, List, Optional
 
 from sqlalchemy import Column
 
-from metadata.generated.schema.entity.data.table import TableData
+from metadata.generated.schema.entity.data.table import CustomMetricProfile, TableData
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.mixins.pandas.pandas_mixin import PandasInterfaceMixin
+from metadata.profiler.api.models import ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR, SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.datalake.datalake_utils import fetch_col_types, fetch_dataframe
+from metadata.utils.datalake.datalake_utils import (
+    GenericDataFrameColumnParser,
+    fetch_dataframe,
+)
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
@@ -239,35 +244,67 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         """
         return None  # to be implemented
 
+    def _compute_custom_metrics(
+        self, metrics: List[CustomMetric], runner, *args, **kwargs
+    ):
+        """Compute custom metrics. For pandas source we expect expression
+        to be a boolean value. We'll return the length of the dataframe
+
+        Args:
+            metrics (List[Metrics]): list of customMetrics
+            runner (_type_): runner
+        """
+        if not metrics:
+            return None
+
+        custom_metrics = []
+
+        for metric in metrics:
+            try:
+                row = sum(
+                    len(df.query(metric.expression).index)
+                    for df in runner
+                    if len(df.query(metric.expression).index)
+                )
+                custom_metrics.append(
+                    CustomMetricProfile(name=metric.name.__root__, value=row)
+                )
+
+            except Exception as exc:
+                msg = f"Error trying to compute profile for custom metric: {exc}"
+                logger.debug(traceback.format_exc())
+                logger.warning(msg)
+        if custom_metrics:
+            return {"customMetrics": custom_metrics}
+        return None
+
     def compute_metrics(
         self,
-        metrics,
-        metric_type,
-        column,
-        table,
+        metric_func: ThreadPoolMetrics,
     ):
         """Run metrics in processor worker"""
-        logger.debug(f"Running profiler for {table}")
+        logger.debug(f"Running profiler for {metric_func.table}")
         try:
             row = None
             if self.complex_dataframe_sample:
-                row = self._get_metric_fn[metric_type.value](
-                    metrics,
+                row = self._get_metric_fn[metric_func.metric_type.value](
+                    metric_func.metrics,
                     self.complex_dataframe_sample,
-                    column=column,
+                    column=metric_func.column,
                 )
         except Exception as exc:
-            name = f"{column if column is not None else table}"
+            name = f"{metric_func.column if metric_func.column is not None else metric_func.table}"
             error = f"{name} metric_type.value: {exc}"
             logger.error(error)
             self.status.failed_profiler(error, traceback.format_exc())
             row = None
-        if column is not None:
-            column = column.name
-            self.status.scanned(f"{table.name.__root__}.{column}")
+        if metric_func.column is not None:
+            column = metric_func.column.name
+            self.status.scanned(f"{metric_func.table.name.__root__}.{column}")
         else:
-            self.status.scanned(table.name.__root__)
-        return row, column, metric_type.value
+            self.status.scanned(metric_func.table.name.__root__)
+            column = None
+        return row, column, metric_func.metric_type.value
 
     def fetch_sample_data(self, table, columns: SQALikeColumn) -> TableData:
         """Fetch sample data from database
@@ -329,7 +366,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
         profile_results = {"table": {}, "columns": defaultdict(dict)}
         metric_list = [
-            self.compute_metrics(*metric_func) for metric_func in metric_funcs
+            self.compute_metrics(metric_func) for metric_func in metric_funcs
         ]
         for metric_result in metric_list:
             profile, column, metric_type = metric_result
@@ -338,6 +375,8 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                     profile_results["table"].update(profile)
                 if metric_type == MetricTypes.System.value:
                     profile_results["system"] = profile
+                elif metric_type == MetricTypes.Custom.value and column is None:
+                    profile_results["table"].update(profile)
                 else:
                     if profile:
                         profile_results["columns"][column].update(
@@ -375,7 +414,9 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                 sqalike_columns.append(
                     SQALikeColumn(
                         column_name,
-                        fetch_col_types(self.complex_dataframe_sample[0], column_name),
+                        GenericDataFrameColumnParser.fetch_col_types(
+                            self.complex_dataframe_sample[0], column_name
+                        ),
                     )
                 )
             return sqalike_columns

metadata/profiler/interface/profiler_interface.py

@@ -42,10 +42,13 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseConnection,
     DatabaseService,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
-from metadata.ingestion.api.models import StackTraceError
+from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
@@ -75,7 +78,7 @@ class ProfilerProcessorStatus(Status):
             StackTraceError(
                 name=self.entity if self.entity else "",
                 error=error,
-                stack_trace=stack_trace,
+                stackTrace=stack_trace,
             )
         )
 
@@ -129,6 +132,7 @@ class ProfilerInterface(ABC):
             MetricTypes.Query.value: self._compute_query_metrics,
             MetricTypes.Window.value: self._compute_window_metrics,
             MetricTypes.System.value: self._compute_system_metrics,
+            MetricTypes.Custom.value: self._compute_custom_metrics,
         }
 
     @abstractmethod
@@ -458,6 +462,13 @@ class ProfilerInterface(ABC):
         """Get metrics"""
         raise NotImplementedError
 
+    @abstractmethod
+    def _compute_custom_metrics(
+        self, metrics: List[CustomMetric], runner, *args, **kwargs
+    ):
+        """Compute custom metrics"""
+        raise NotImplementedError
+
     @abstractmethod
     def get_all_metrics(self, metric_funcs) -> dict:
         """run profiler metrics"""

metadata/profiler/interface/profiler_interface_factory.py

@@ -18,9 +18,18 @@ from typing import cast
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    Db2Connection,
+)
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariaDBConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
     SingleStoreConnection,
 )
@@ -30,6 +39,9 @@ from metadata.generated.schema.entity.services.connections.database.snowflakeCon
 from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
     TrinoConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.profiler.interface.pandas.profiler_interface import (
     PandasProfilerInterface,
@@ -38,6 +50,15 @@ from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.interface.sqlalchemy.bigquery.profiler_interface import (
     BigQueryProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import (
+    DatabricksProfilerInterface,
+)
+from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
+    DB2ProfilerInterface,
+)
+from metadata.profiler.interface.sqlalchemy.mariadb.profiler_interface import (
+    MariaDBProfilerInterface,
+)
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
     SQAProfilerInterface,
 )
@@ -50,6 +71,9 @@ from metadata.profiler.interface.sqlalchemy.snowflake.profiler_interface import
 from metadata.profiler.interface.sqlalchemy.trino.profiler_interface import (
     TrinoProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
+    UnityCatalogProfilerInterface,
+)
 
 
 class ProfilerInterfaceFactory:
@@ -62,6 +86,17 @@ class ProfilerInterfaceFactory:
         """Register a new interface"""
         self._interface_type[interface_type] = interface_class
 
+    def register_many(self, interface_dict):
+        """
+        Registers multiple profiler interfaces at once.
+
+        Args:
+            interface_dict: A dictionary mapping connection class names (strings) to their
+            corresponding profiler interface classes.
+        """
+        for interface_type, interface_class in interface_dict.items():
+            self.register(interface_type, interface_class)
+
     def create(self, interface_type: str, *args, **kwargs):
         """Create interface object based on interface type"""
         interface_class = self._interface_type.get(interface_type)
@@ -72,17 +107,17 @@ class ProfilerInterfaceFactory:
 
 
 profiler_interface_factory = ProfilerInterfaceFactory()
-profiler_interface_factory.register(DatabaseConnection.__name__, SQAProfilerInterface)
-profiler_interface_factory.register(
-    BigQueryConnection.__name__, BigQueryProfilerInterface
-)
-profiler_interface_factory.register(
-    SingleStoreConnection.__name__, SingleStoreProfilerInterface
-)
-profiler_interface_factory.register(
-    DatalakeConnection.__name__, PandasProfilerInterface
-)
-profiler_interface_factory.register(
-    SnowflakeConnection.__name__, SnowflakeProfilerInterface
-)
-profiler_interface_factory.register(TrinoConnection.__name__, TrinoProfilerInterface)
+profilers = {
+    DatabaseConnection.__name__: SQAProfilerInterface,
+    BigQueryConnection.__name__: BigQueryProfilerInterface,
+    SingleStoreConnection.__name__: SingleStoreProfilerInterface,
+    DatalakeConnection.__name__: PandasProfilerInterface,
+    MariaDBConnection.__name__: MariaDBProfilerInterface,
+    SnowflakeConnection.__name__: SnowflakeProfilerInterface,
+    TrinoConnection.__name__: TrinoProfilerInterface,
+    UnityCatalogConnection.__name__: UnityCatalogProfilerInterface,
+    DatabricksConnection.__name__: DatabricksProfilerInterface,
+    Db2Connection.__name__: DB2ProfilerInterface,
+}
+
+profiler_interface_factory.register_many(profilers)

metadata/profiler/interface/sqlalchemy/bigquery/profiler_interface.py

@@ -32,6 +32,9 @@ class BigQueryProfilerInterface(SQAProfilerInterface):
         for key, value in columns.items():
             if not isinstance(value, STRUCT):
                 col = Column(f"{parent}.{key}", value)
+                # pylint: disable=protected-access
+                col._set_parent(self.table.__table__)
+                # pylint: enable=protected-access
                 columns_list.append(col)
             else:
                 col = self._get_struct_columns(