openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -20,7 +20,7 @@ import math
 import threading
 import traceback
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Dict, List, Optional
 
 from sqlalchemy import Column, inspect, text
@@ -350,7 +350,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                     crs.scalar()
                 )  # raise MultipleResultsFound if more than one row is returned
                 custom_metrics.append(
-                    CustomMetricProfile(name=metric.name.__root__, value=row)
+                    CustomMetricProfile(name=metric.name.root, value=row)
                 )
 
             except Exception as exc:
@@ -515,9 +515,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                         profile_results["columns"][column].update(
                             {
                                 "name": column,
-                                "timestamp": int(
-                                    datetime.now(tz=timezone.utc).timestamp() * 1000
-                                ),
+                                "timestamp": int(datetime.now().timestamp() * 1000),
                                 **profile,
                             }
                         )

metadata/profiler/metrics/composed/null_ratio.py

@@ -49,8 +49,9 @@ class NullRatio(ComposedMetric):
         Safely compute null ratio based on the profiler
         results of other Metrics
         """
-        res_count = res.get(Count.name())
-        res_null = res.get(NullCount.name())
-        if res_count + res_null == 0:
+
+        count = res.get(Count.name())
+        null_count = res.get(NullCount.name())
+        if count + null_count == 0:
             return None
-        return res_null / (res_null + res_count)
+        return null_count / (null_count + count)

metadata/profiler/metrics/static/mean.py

@@ -52,6 +52,15 @@ def _(element, compiler, **kw):
     return f"if(isNaN(avg({proc})), null, avg({proc}))"
 
 
+@compiles(avg, Dialects.Redshift)
+def _(element, compiler, **kw):
+    """
+    Cast to decimal to get around potential integer overflow error
+    """
+    proc = compiler.process(element.clauses, **kw)
+    return f"avg(CAST({proc} AS DECIMAL(38,0)))"
+
+
 @compiles(avg, Dialects.MSSQL)
 def _(element, compiler, **kw):
     """

metadata/profiler/metrics/system/queries/snowflake.py

@@ -101,7 +101,7 @@ def get_identifiers(
         )
         es_tables = search_table_entities(
             metadata=ometa_client,
-            service_name=db_service.fullyQualifiedName.__root__,
+            service_name=db_service.fullyQualifiedName.root,
             database=database_name,
             database_schema=schema_name,
             table=table_name,

metadata/profiler/metrics/window/percentille_mixin.py

@@ -1,4 +1,4 @@
-"""function calls shared accross all percentile metrics"""
+"""function calls shared across all percentile metrics"""
 
 from metadata.profiler.orm.functions.median import MedianFn
 

metadata/profiler/orm/converter/base.py

@@ -64,7 +64,9 @@ def check_if_should_quote_column_name(table_service_type) -> Optional[bool]:
     return None
 
 
-def build_orm_col(idx: int, col: Column, table_service_type) -> sqlalchemy.Column:
+def build_orm_col(
+    idx: int, col: Column, table_service_type, *, _quote=None
+) -> sqlalchemy.Column:
     """
     Cook the ORM column from our metadata instance
     information.
@@ -76,16 +78,22 @@ def build_orm_col(idx: int, col: Column, table_service_type) -> sqlalchemy.Colum
     As this is only used for INSERT/UPDATE/DELETE,
     there is no impact for our read-only purposes.
     """
+    if _quote is not None:
+        quote = _quote
+    else:
+        quote = check_if_should_quote_column_name(
+            table_service_type
+        ) or check_snowflake_case_sensitive(table_service_type, col.name.root)
+
     return sqlalchemy.Column(
-        name=str(col.name.__root__),
+        name=str(col.name.root),
         type_=converter_registry[table_service_type]().map_types(
             col, table_service_type
         ),
         primary_key=not bool(idx),  # The first col seen is used as PK
-        quote=check_if_should_quote_column_name(table_service_type)
-        or check_snowflake_case_sensitive(table_service_type, col.name.__root__),
+        quote=quote,
         key=str(
-            col.name.__root__
+            col.name.root
         ).lower(),  # Add lowercase column name as key for snowflake case sensitive columns
     )
 
@@ -107,16 +115,16 @@ def ometa_to_sqa_orm(
     )  # satisfy mypy
     cols = {
         (
-            col.name.__root__ + "_"
-            if col.name.__root__ in SQA_RESERVED_ATTRIBUTES
-            else col.name.__root__
+            col.name.root + "_"
+            if col.name.root in SQA_RESERVED_ATTRIBUTES
+            else col.name.root
         ): build_orm_col(idx, col, table.serviceType)
         for idx, col in enumerate(table.columns)
     }
 
     orm_database_name = get_orm_database(table, metadata)
     orm_schema_name = get_orm_schema(table, metadata)
-    orm_name = f"{orm_database_name}_{orm_schema_name}_{table.name.__root__}".replace(
+    orm_name = f"{orm_database_name}_{orm_schema_name}_{table.name.root}".replace(
         ".", "_"
     )
 
@@ -125,7 +133,7 @@ def ometa_to_sqa_orm(
         orm_name,  # Output class name
         (Base,),  # SQLAlchemy declarative base
         {
-            "__tablename__": str(table.name.__root__),
+            "__tablename__": str(table.name.root),
             "__table_args__": {
                 # SQLite does not support schemas
                 "schema": orm_schema_name
@@ -133,7 +141,7 @@ def ometa_to_sqa_orm(
                 else None,
                 "extend_existing": True,  # Recreates the table ORM object if it already exists. Useful for testing
                 "quote": check_snowflake_case_sensitive(
-                    table.serviceType, table.name.__root__
+                    table.serviceType, table.name.root
                 ),
             },
             **cols,
@@ -166,7 +174,7 @@ def get_orm_schema(table: Table, metadata: OpenMetadata) -> str:
         entity=DatabaseSchema, entity_id=table.databaseSchema.id
     )
 
-    return str(schema.name.__root__)
+    return str(schema.name.root)
 
 
 def get_orm_database(table: Table, metadata: OpenMetadata) -> str:
@@ -184,4 +192,4 @@ def get_orm_database(table: Table, metadata: OpenMetadata) -> str:
         entity=Database, entity_id=table.database.id
     )
 
-    return str(database.name.__root__)
+    return str(database.name.root)

metadata/profiler/orm/converter/common.py

@@ -12,7 +12,6 @@
 """
 Common Class For Profiler Converter.
 """
-
 from typing import Dict, Set
 
 import sqlalchemy
@@ -51,8 +50,8 @@ class CommonMapTypes:
         DataType.CHAR: sqlalchemy.CHAR,
         DataType.VARCHAR: sqlalchemy.VARCHAR,
         DataType.BOOLEAN: sqlalchemy.BOOLEAN,
-        DataType.BINARY: sqlalchemy.LargeBinary,
-        DataType.VARBINARY: sqlalchemy.VARBINARY,
+        DataType.BINARY: CustomTypes.BYTES.value,
+        DataType.VARBINARY: CustomTypes.BYTES.value,
         DataType.ARRAY: CustomTypes.ARRAY.value,
         DataType.BLOB: CustomTypes.BYTES.value,
         DataType.LONGBLOB: sqlalchemy.LargeBinary,
@@ -81,7 +80,7 @@ class CommonMapTypes:
         return self.return_custom_type(col, table_service_type)
 
     def return_custom_type(self, col: Column, _):
-        return self._TYPE_MAP.get(col.dataType)
+        return self._TYPE_MAP.get(col.dataType, CustomTypes.UNDETERMINED.value)
 
     @staticmethod
     def map_sqa_to_om_types() -> Dict[TypeEngine, Set[DataType]]:

metadata/profiler/orm/functions/datetime.py

@@ -188,7 +188,7 @@ def _(elements, compiler, **kwargs):  # pylint: disable=unused-argument
             "Bigquery does not support `month` or `year` interval for table partitioned on timestamp",
             "field types. You can set the `interval_unit to day or hour directly from OpenMetadata UI`."
             # pylint: disable=line-too-long
-            "Visit https://docs.open-metadata.org/connectors/ingestion/workflows/profiler#4-updating-profiler-setting-at-the-table-level for more details.",
+            "Visit https://docs.open-metadata.org/how-to-guides/data-quality-observability/profiler/workflow#4-updating-profiler-setting-at-the-table-level for more details.",
         )
 
     return (

metadata/profiler/orm/functions/sum.py

@@ -32,13 +32,20 @@ def _(element, compiler, **kw):
     return f"SUM(CAST({proc} AS BIGINT))"
 
 
+@compiles(SumFn, Dialects.Redshift)
+def _(element, compiler, **kw):
+    """Cast to Decimal to address overflow error from summing 32-bit int in most database dialects"""
+    proc = compiler.process(element.clauses, **kw)
+    return f"SUM(CAST({proc} AS Decimal(38,0)))"
+
+
 @compiles(SumFn, Dialects.Athena)
 @compiles(SumFn, Dialects.Trino)
 @compiles(SumFn, Dialects.Presto)
 def _(element, compiler, **kw):
     """Cast to DECIMAL to address cannot cast nan to bigint"""
     proc = compiler.process(element.clauses, **kw)
-    return f"SUM(TRY_CAST({proc} AS DECIMAL))"
+    return f"COALESCE(SUM(CAST({proc} AS DECIMAL)),0)"
 
 
 @compiles(SumFn, Dialects.BigQuery)

metadata/profiler/orm/functions/table_metric_computer.py

@@ -295,7 +295,7 @@ class BigQueryTableMetricComputer(BaseTableMetricComputer):
 
         where_clause = [
             Column("project_id")
-            == self.conn_config.credentials.gcpConfig.projectId.__root__,
+            == self.conn_config.credentials.gcpConfig.projectId.root,
             Column("table_schema") == self.schema_name,
             Column("table_name") == self.table_name,
         ]
@@ -324,21 +324,20 @@ class BigQueryTableMetricComputer(BaseTableMetricComputer):
         columns = [
             Column("row_count").label("rowCount"),
             Column("size_bytes").label("sizeInBytes"),
-            Column("creation_time").label("createDateTime"),
+            func.TIMESTAMP_MILLIS(Column("creation_time")).label(CREATE_DATETIME),
             *self._get_col_names_and_count(),
         ]
         where_clause = [
             Column("project_id")
-            == self.conn_config.credentials.gcpConfig.projectId.__root__,
+            == self.conn_config.credentials.gcpConfig.projectId.root,
             Column("dataset_id") == self.schema_name,
             Column("table_id") == self.table_name,
         ]
-
         query = self._build_query(
             columns,
             self._build_table(
                 "__TABLES__",
-                f"{self.conn_config.credentials.gcpConfig.projectId.__root__}.{self.schema_name}",
+                f"{self.conn_config.credentials.gcpConfig.projectId.root}.{self.schema_name}",
             ),
             where_clause,
         )

metadata/profiler/orm/registry.py

@@ -28,6 +28,7 @@ from metadata.profiler.orm.types.custom_hex_byte_string import HexByteString
 from metadata.profiler.orm.types.custom_image import CustomImage
 from metadata.profiler.orm.types.custom_ip import CustomIP
 from metadata.profiler.orm.types.custom_timestamp import CustomTimestamp
+from metadata.profiler.orm.types.undetermined_type import UndeterminedType
 from metadata.profiler.orm.types.uuid import UUIDString
 from metadata.profiler.registry import TypeRegistry
 
@@ -41,6 +42,7 @@ class CustomTypes(TypeRegistry):
     IMAGE = CustomImage
     IP = CustomIP
     SQADATETIMERANGE = CustomDateTimeRange
+    UNDETERMINED = UndeterminedType
 
 
 class Dialects(Enum):
@@ -99,6 +101,8 @@ NOT_COMPUTE = {
     DataType.JSON.value,
     CustomTypes.ARRAY.value.__name__,
     CustomTypes.SQADATETIMERANGE.value.__name__,
+    DataType.XML.value,
+    CustomTypes.UNDETERMINED.value.__name__,
 }
 FLOAT_SET = {sqlalchemy.types.DECIMAL, sqlalchemy.types.FLOAT}
 
@@ -148,6 +152,7 @@ def is_date_time(_type) -> bool:
         issubclass(_type.__class__, Date)
         or issubclass(_type.__class__, Time)
         or issubclass(_type.__class__, DateTime)
+        or issubclass(_type.__class__, CustomTimestamp)
     )
 
 

metadata/profiler/orm/types/custom_hex_byte_string.py

@@ -42,9 +42,11 @@ class HexByteString(TypeDecorator):
         Make sure the data is of correct type
         """
         if not isinstance(value, (bytes, bytearray)):
-            raise TypeError("HexByteString columns support only bytes values.")
+            raise TypeError(
+                f"HexByteString columns support only bytes values. Received {type(value).__name__}."
+            )
 
-    def process_result_value(self, value: str, dialect) -> Optional[str]:
+    def process_result_value(self, value: Optional[bytes], dialect) -> Optional[str]:
         """This is executed during result retrieval
 
         Args:
@@ -64,8 +66,8 @@ class HexByteString(TypeDecorator):
                 value = bytes_value.decode(encoding=detected_encoding)
                 return value
             except Exception as exc:
+                logger.debug("Failed to parse bytes valud as string: %s", exc)
                 logger.debug(traceback.format_exc())
-                logger.error(exc)
 
         return value.hex()
 

metadata/profiler/orm/types/undetermined_type.py

@@ -0,0 +1,38 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+# pylint: disable=abstract-method
+
+"""
+Undetermined types for cases where we dont have typ mappings
+"""
+from sqlalchemy.sql.sqltypes import String, TypeDecorator
+
+
+class UndeterminedType(TypeDecorator):
+    """A fallback type for undetermined types returned from the database"""
+
+    impl = String
+    cache_ok = True
+
+    @property
+    def python_type(self):
+        return str
+
+    def process_result_value(self, value, _):
+        """
+        We have no idea what is this type. So we just casr
+        """
+        return (
+            f"OPENMETADATA_UNDETERMIND[{str(value)}]"
+            if value
+            else "OPENMETADATA_UNDETERMIND[]"
+        )

metadata/profiler/processor/core.py

@@ -40,6 +40,7 @@ from metadata.generated.schema.settings.settings import Settings
 from metadata.generated.schema.tests.customMetric import (
     CustomMetric as CustomMetricEntity,
 )
+from metadata.generated.schema.type.basic import Timestamp
 from metadata.profiler.api.models import ProfilerResponse, ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import (
@@ -51,6 +52,7 @@ from metadata.profiler.metrics.core import (
     TMetric,
 )
 from metadata.profiler.metrics.static.row_count import RowCount
+from metadata.profiler.orm.functions.table_metric_computer import CREATE_DATETIME
 from metadata.profiler.orm.registry import NOT_COMPUTE
 from metadata.profiler.processor.metric_filter import MetricFilter
 from metadata.profiler.processor.sample_data_handler import upload_sample_data
@@ -105,7 +107,7 @@ class Profiler(Generic[TMetric]):
         self.include_columns = include_columns
         self.exclude_columns = exclude_columns
         self._metrics = metrics
-        self._profile_date = int(datetime.now(tz=timezone.utc).timestamp() * 1000)
+        self._profile_ts = Timestamp(int(datetime.now().timestamp() * 1000))
         self.profile_sample_config = self.profiler_interface.profile_sample_config
 
         self.metric_filter = MetricFilter(
@@ -150,8 +152,8 @@ class Profiler(Generic[TMetric]):
         return self._get_included_columns()
 
     @property
-    def profile_date(self) -> datetime:
-        return self._profile_date
+    def profile_ts(self) -> Timestamp:
+        return self._profile_ts
 
     @property
     def columns(self) -> List[Column]:
@@ -218,7 +220,7 @@ class Profiler(Generic[TMetric]):
                     return
 
         raise RuntimeError(
-            f"No profile data computed for {self.profiler_interface.table_entity.fullyQualifiedName.__root__}"
+            f"No profile data computed for {self.profiler_interface.table_entity.fullyQualifiedName.root}"
         )
 
     def get_custom_metrics(
@@ -240,7 +242,7 @@ class Profiler(Generic[TMetric]):
             (
                 clmn
                 for clmn in self.profiler_interface.table_entity.columns
-                if clmn.name.__root__ == column_name
+                if clmn.name.root == column_name
             ),
             None,
         )
@@ -486,7 +488,7 @@ class Profiler(Generic[TMetric]):
 
         if self.source_config.computeMetrics:
             logger.debug(
-                f"Computing profile metrics for {self.profiler_interface.table_entity.fullyQualifiedName.__root__}..."
+                f"Computing profile metrics for {self.profiler_interface.table_entity.fullyQualifiedName.root}..."
             )
             self.compute_metrics()
 
@@ -517,7 +519,7 @@ class Profiler(Generic[TMetric]):
         try:
             logger.debug(
                 "Fetching sample data for "
-                f"{self.profiler_interface.table_entity.fullyQualifiedName.__root__}..."  # type: ignore
+                f"{self.profiler_interface.table_entity.fullyQualifiedName.root}..."  # type: ignore
             )
             table_data = self.profiler_interface.fetch_sample_data(
                 self.table, self.columns
@@ -568,22 +570,26 @@ class Profiler(Generic[TMetric]):
                     **self.column_results.get(
                         col.name
                         if not isinstance(col.name, ColumnName)
-                        else col.name.__root__
+                        else col.name.root
                     )
                 )
                 for col in self.columns
                 if self.column_results.get(
-                    col.name
-                    if not isinstance(col.name, ColumnName)
-                    else col.name.__root__
+                    col.name if not isinstance(col.name, ColumnName) else col.name.root
                 )
             ]
 
+            raw_create_date: Optional[datetime] = self._table_results.get(
+                CREATE_DATETIME
+            )
+            if raw_create_date:
+                raw_create_date = raw_create_date.replace(tzinfo=timezone.utc)
+
             table_profile = TableProfile(
-                timestamp=self.profile_date,
+                timestamp=self.profile_ts,
                 columnCount=self._table_results.get("columnCount"),
                 rowCount=self._table_results.get(RowCount.name()),
-                createDateTime=self._table_results.get("createDateTime"),
+                createDateTime=raw_create_date,
                 sizeInByte=self._table_results.get("sizeInBytes"),
                 profileSample=(
                     self.profile_sample_config.profile_sample

metadata/profiler/processor/default.py

@@ -77,9 +77,12 @@ class DefaultProfiler(Profiler):
         include_columns: Optional[List[ColumnProfilerConfig]] = None,
         exclude_columns: Optional[List[str]] = None,
         global_profiler_configuration: Optional[ProfilerConfiguration] = None,
+        db_service=None,
     ):
         _metrics = get_default_metrics(
-            table=profiler_interface.table, ometa_client=profiler_interface.ometa_client
+            table=profiler_interface.table,
+            ometa_client=profiler_interface.ometa_client,
+            db_service=db_service,
         )
 
         super().__init__(

metadata/profiler/processor/processor.py

@@ -43,8 +43,8 @@ class ProfilerProcessor(Processor):
         super().__init__()
 
         self.config = config
-        self.profiler_config = ProfilerProcessorConfig.parse_obj(
-            self.config.processor.dict().get("config")
+        self.profiler_config = ProfilerProcessorConfig.model_validate(
+            self.config.processor.model_dump().get("config")
         )
         self.source_config: DatabaseServiceProfilerPipeline = cast(
             DatabaseServiceProfilerPipeline, self.config.source.sourceConfig.config
@@ -64,8 +64,8 @@ class ProfilerProcessor(Processor):
         except Exception as exc:
             self.status.failed(
                 StackTraceError(
-                    name=record.entity.fullyQualifiedName.__root__,
-                    error=f"Unexpected exception processing entity {record.entity.fullyQualifiedName.__root__}: {exc}",
+                    name=record.entity.fullyQualifiedName.root,
+                    error=f"Unexpected exception processing entity {record.entity.fullyQualifiedName.root}: {exc}",
                     stackTrace=traceback.format_exc(),
                 )
             )

metadata/profiler/processor/sample_data_handler.py

@@ -17,14 +17,13 @@ from datetime import datetime
 from functools import singledispatch
 from io import BytesIO
 
-from pydantic.json import ENCODERS_BY_TYPE
-
 from metadata.clients.aws_client import AWSClient
 from metadata.generated.schema.entity.data.table import Table, TableData
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
     DataStorageConfig,
 )
 from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.ingestion.models.custom_pydantic import ignore_type_decoder
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import profiler_logger
@@ -62,7 +61,7 @@ def _get_object_key(
         service_name=table.service.name,
         database_name=table.database.name,
         database_schema_name=table.databaseSchema.name,
-        table_name=table.name.__root__,
+        table_name=table.name.root,
     )
     if not overwrite_data:
         file_name = file_name.replace(
@@ -83,11 +82,12 @@ def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -
         sample_storage_config: DataStorageConfig = profiler_interface.storage_config
         if not sample_storage_config:
             return
-        ENCODERS_BY_TYPE[bytes] = lambda v: v.decode("utf-8", "ignore")
+        # Ignore any decoding error for byte data
+        ignore_type_decoder(bytes)
         deserialized_data = json.loads(data.json())
         df = pd.DataFrame(
             data=deserialized_data.get("rows", []),
-            columns=[i.__root__ for i in data.columns],
+            columns=[i.root for i in data.columns],
         )
         pq_buffer = BytesIO()
         df.to_parquet(pq_buffer)

metadata/profiler/processor/sampler/nosql/sampler.py

@@ -26,11 +26,13 @@ class NoSQLSampler(SamplerInterface):
         """
         records = self._rdn_sample_from_user_query()
         columns = [
-            SQALikeColumn(name=column.name.__root__, type=column.dataType)
+            SQALikeColumn(name=column.name.root, type=column.dataType)
             for column in self.table.columns
         ]
         rows, cols = self.transpose_records(records, columns)
-        return TableData(rows=rows, columns=[c.name for c in cols])
+        return TableData(
+            rows=[list(map(str, row)) for row in rows], columns=[c.name for c in cols]
+        )
 
     def random_sample(self):
         pass
@@ -40,14 +42,17 @@ class NoSQLSampler(SamplerInterface):
             return self._fetch_sample_data_from_user_query()
         return self._fetch_sample_data(columns)
 
-    def _fetch_sample_data(self, columns: List[SQALikeColumn]):
+    def _fetch_sample_data(self, columns: List[SQALikeColumn]) -> TableData:
         """
         returns sampled ometa dataframes
         """
         limit = self._get_limit()
         records = self.client.scan(self.table, self.table.columns, limit)
         rows, cols = self.transpose_records(records, columns)
-        return TableData(rows=rows, columns=[col.name for col in cols])
+        return TableData(
+            rows=[list(map(str, row)) for row in rows],
+            columns=[col.name for col in cols],
+        )
 
     def _get_limit(self) -> Optional[int]:
         num_rows = self.client.item_count(self.table)

metadata/profiler/processor/sampler/sampler_factory.py

@@ -15,6 +15,9 @@ Factory class for creating sampler objects
 
 from typing import Union
 
+from metadata.generated.schema.entity.services.connections.database.azureSQLConnection import (
+    AzureSQLConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
@@ -27,16 +30,25 @@ from metadata.generated.schema.entity.services.connections.database.dynamoDBConn
 from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
     MongoDBConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
+    SnowflakeConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
     TrinoConnection,
 )
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.profiler.processor.sampler.nosql.sampler import NoSQLSampler
 from metadata.profiler.processor.sampler.pandas.sampler import DatalakeSampler
+from metadata.profiler.processor.sampler.sqlalchemy.azuresql.sampler import (
+    AzureSQLSampler,
+)
 from metadata.profiler.processor.sampler.sqlalchemy.bigquery.sampler import (
     BigQuerySampler,
 )
 from metadata.profiler.processor.sampler.sqlalchemy.sampler import SQASampler
+from metadata.profiler.processor.sampler.sqlalchemy.snowflake.sampler import (
+    SnowflakeSampler,
+)
 from metadata.profiler.processor.sampler.sqlalchemy.trino.sampler import TrinoSampler
 
 
@@ -62,9 +74,27 @@ class SamplerFactory:
 
 
 sampler_factory_ = SamplerFactory()
-sampler_factory_.register(DatabaseConnection.__name__, SQASampler)
-sampler_factory_.register(BigQueryConnection.__name__, BigQuerySampler)
-sampler_factory_.register(DatalakeConnection.__name__, DatalakeSampler)
-sampler_factory_.register(TrinoConnection.__name__, TrinoSampler)
-sampler_factory_.register(MongoDBConnection.__name__, NoSQLSampler)
-sampler_factory_.register(DynamoDBConnection.__name__, NoSQLSampler)
+sampler_factory_.register(
+    source_type=DatabaseConnection.__name__, sampler_class=SQASampler
+)
+sampler_factory_.register(
+    source_type=BigQueryConnection.__name__, sampler_class=BigQuerySampler
+)
+sampler_factory_.register(
+    source_type=DatalakeConnection.__name__, sampler_class=DatalakeSampler
+)
+sampler_factory_.register(
+    source_type=TrinoConnection.__name__, sampler_class=TrinoSampler
+)
+sampler_factory_.register(
+    source_type=MongoDBConnection.__name__, sampler_class=NoSQLSampler
+)
+sampler_factory_.register(
+    source_type=SnowflakeConnection.__name__, sampler_class=SnowflakeSampler
+)
+sampler_factory_.register(
+    source_type=DynamoDBConnection.__name__, sampler_class=NoSQLSampler
+)
+sampler_factory_.register(
+    source_type=AzureSQLConnection.__name__, sampler_class=AzureSQLSampler
+)