openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/{profiler/processor/sampler → sampler}/sqlalchemy/sampler.py

@@ -16,29 +16,24 @@ import traceback
 from typing import List, Optional, Union, cast
 
 from sqlalchemy import Column, inspect, text
-from sqlalchemy.orm import DeclarativeMeta, Query, aliased
+from sqlalchemy.orm import DeclarativeMeta, Query
 from sqlalchemy.orm.util import AliasedClass
+from sqlalchemy.schema import Table
 from sqlalchemy.sql.sqltypes import Enum
 
 from metadata.generated.schema.entity.data.table import (
-    PartitionIntervalTypes,
     PartitionProfilerConfig,
     ProfileSampleType,
     TableData,
 )
+from metadata.ingestion.connections.session import create_and_bind_thread_safe_session
+from metadata.mixins.sqalchemy.sqa_mixin import SQAInterfaceMixin
 from metadata.profiler.orm.functions.modulo import ModuloFn
 from metadata.profiler.orm.functions.random_num import RandomNumFn
-from metadata.profiler.processor.handle_partition import partition_filter_handler
-from metadata.profiler.processor.sampler.sampler_interface import SamplerInterface
+from metadata.profiler.processor.handle_partition import build_partition_predicate
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.helpers import is_safe_sql_query
 from metadata.utils.logger import profiler_interface_registry_logger
-from metadata.utils.sqa_utils import (
-    build_query_filter,
-    dispatch_to_date_or_datetime,
-    get_integer_range_filter,
-    get_partition_col_type,
-    get_value_filter,
-)
 
 logger = profiler_interface_registry_logger()
 
@@ -61,12 +56,37 @@ def _object_value_for_elem(self, elem):
 Enum._object_value_for_elem = _object_value_for_elem  # pylint: disable=protected-access
 
 
-class SQASampler(SamplerInterface):
+class SQASampler(SamplerInterface, SQAInterfaceMixin):
     """
     Generates a sample of the data to not
     run the query in the whole table.
+
+    Args:
+        orm_table (Optional[DeclarativeMeta]): ORM Table
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._table = self.build_table_orm(
+            self.entity, self.service_connection_config, self.ometa_client
+        )
+
+    @property
+    def raw_dataset(self):
+        return self._table
+
+    def get_client(self):
+        """Build the SQA Client"""
+        session_factory = create_and_bind_thread_safe_session(self.connection)
+        return session_factory()
+
+    def set_tablesample(self, selectable: Table):
+        """Set the tablesample for the table. To be implemented by the child SQA sampler class
+        Args:
+            selectable (Table): a selectable table
+        """
+        return selectable
+
     def _base_sample_query(self, column: Optional[Column], label=None):
         """Base query for sampling
 
@@ -77,54 +97,60 @@ class SQASampler(SamplerInterface):
         Returns:
         """
         # only sample the column if we are computing a column metric to limit the amount of data scaned
-        entity = self.table if column is None else column
+        selectable = self.set_tablesample(self.raw_dataset.__table__)
+
+        entity = selectable if column is None else selectable.c.get(column.key)
         if label is not None:
-            return self.client.query(entity, label)
-        return self.client.query(entity)
+            query = self.client.query(entity, label)
+        else:
+            query = self.client.query(entity)
+
+        if self.partition_details:
+            query = self.get_partitioned_query(query)
+        return query
 
-    @partition_filter_handler(build_sample=True)
     def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
-        if self.profile_sample_type == ProfileSampleType.PERCENTAGE:
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
             rnd = self._base_sample_query(
                 column,
                 (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
-            ).cte(f"{self.table.__tablename__}_rnd")
+            ).cte(f"{self.raw_dataset.__tablename__}_rnd")
             session_query = self.client.query(rnd)
-            return session_query.where(rnd.c.random <= self.profile_sample).cte(
-                f"{self.table.__tablename__}_sample"
-            )
+            return session_query.where(
+                rnd.c.random <= self.sample_config.profile_sample
+            ).cte(f"{self.raw_dataset.__tablename__}_sample")
 
-        table_query = self.client.query(self.table)
+        table_query = self.client.query(self.raw_dataset)
         session_query = self._base_sample_query(
             column,
             (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL),
         )
         return (
             session_query.order_by(RANDOM_LABEL)
-            .limit(self.profile_sample)
-            .cte(f"{self.table.__tablename__}_rnd")
+            .limit(self.sample_config.profile_sample)
+            .cte(f"{self.raw_dataset.__tablename__}_rnd")
         )
 
-    def random_sample(self, ccolumn=None) -> Union[DeclarativeMeta, AliasedClass]:
+    def get_dataset(self, column=None, **__) -> Union[DeclarativeMeta, AliasedClass]:
         """
         Either return a sampled CTE of table, or
         the full table if no sampling is required.
         """
-        if self._profile_sample_query:
+        if self.sample_query:
             return self._rdn_sample_from_user_query()
 
-        if not self.profile_sample or int(self.profile_sample) == 100:
-            if self._partition_details:
-                return self._partitioned_table()
-
-            return self.table
+        if not self.sample_config.profile_sample or (
+            int(self.sample_config.profile_sample) == 100
+            and self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE
+        ):
+            if self.partition_details:
+                partitioned = self._partitioned_table()
+                return partitioned.cte(f"{self.raw_dataset.__tablename__}_partitioned")
 
-        # Add new RandomNumFn column
-        sampled = self.get_sample_query(column=ccolumn)
+            return self.raw_dataset
 
-        # Assign as an alias
-        return aliased(self.table, sampled)
+        return self.get_sample_query(column=column)
 
     def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:
         """
@@ -132,10 +158,10 @@ class SQASampler(SamplerInterface):
 
         Args:
             columns (Optional[List]): List of columns to fetch
-        Retunrs:
+        Returns:
             TableData to be added to the Table Entity
         """
-        if self._profile_sample_query:
+        if self.sample_query:
             return self._fetch_sample_data_from_user_query()
 
         # Add new RandomNumFn column
@@ -143,8 +169,8 @@ class SQASampler(SamplerInterface):
         if not columns:
             sqa_columns = [col for col in inspect(rnd).c if col.name != RANDOM_LABEL]
         else:
-            # we can't directly use columns as it is bound to self.table and not the rnd table.
-            # If we use it, it will result in a cross join between self.table and rnd table
+            # we can't directly use columns as it is bound to self.raw_dataset and not the rnd table.
+            # If we use it, it will result in a cross join between self.raw_dataset and rnd table
             names = [col.name for col in columns]
             sqa_columns = [
                 col
@@ -164,9 +190,12 @@ class SQASampler(SamplerInterface):
                 "Cannot fetch sample data with random sampling. Falling back to 100 rows."
             )
             logger.debug(traceback.format_exc())
-            sqa_columns = list(inspect(self.table).c)
+            sqa_columns = list(inspect(self.raw_dataset).c)
             sqa_sample = (
-                self.client.query(*sqa_columns).select_from(self.table).limit(100).all()
+                self.client.query(*sqa_columns)
+                .select_from(self.raw_dataset)
+                .limit(100)
+                .all()
             )
 
         return TableData(
@@ -176,12 +205,10 @@ class SQASampler(SamplerInterface):
 
     def _fetch_sample_data_from_user_query(self) -> TableData:
         """Returns a table data object using results from query execution"""
-        if not is_safe_sql_query(self._profile_sample_query):
-            raise RuntimeError(
-                f"SQL expression is not safe\n\n{self._profile_sample_query}"
-            )
+        if not is_safe_sql_query(self.sample_query):
+            raise RuntimeError(f"SQL expression is not safe\n\n{self.sample_query}")
 
-        rnd = self.client.execute(f"{self._profile_sample_query}")
+        rnd = self.client.execute(f"{self.sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
         except AttributeError:
@@ -193,65 +220,30 @@ class SQASampler(SamplerInterface):
 
     def _rdn_sample_from_user_query(self) -> Query:
         """Returns sql alchemy object to use when running profiling"""
-        if not is_safe_sql_query(self._profile_sample_query):
-            raise RuntimeError(
-                f"SQL expression is not safe\n\n{self._profile_sample_query}"
-            )
+        if not is_safe_sql_query(self.sample_query):
+            raise RuntimeError(f"SQL expression is not safe\n\n{self.sample_query}")
 
-        return self.client.query(self.table).from_statement(
-            text(f"{self._profile_sample_query}")
+        return self.client.query(self.raw_dataset).from_statement(
+            text(f"{self.sample_query}")
         )
 
     def _partitioned_table(self) -> Query:
         """Return the Query object for partitioned tables"""
-        return aliased(self.get_partitioned_query().subquery())
+        return self.get_partitioned_query()
 
-    def get_partitioned_query(self) -> Query:
+    def get_partitioned_query(self, query=None) -> Query:
         """Return the partitioned query"""
-        self._partition_details = cast(
-            PartitionProfilerConfig, self._partition_details
+        self.partition_details = cast(
+            PartitionProfilerConfig, self.partition_details
         )  # satisfying type checker
-        partition_field = self._partition_details.partitionColumnName
-
-        type_ = get_partition_col_type(
-            partition_field,
-            self.table.__table__.c,
+        partition_filter = build_partition_predicate(
+            self.partition_details,
+            self.raw_dataset.__table__.c,
         )
+        if query is not None:
+            return query.filter(partition_filter)
+        return self.client.query(self.raw_dataset).filter(partition_filter)
 
-        if (
-            self._partition_details.partitionIntervalType
-            == PartitionIntervalTypes.COLUMN_VALUE
-        ):
-            return self.client.query(self.table).filter(
-                get_value_filter(
-                    Column(partition_field),
-                    self._partition_details.partitionValues,
-                )
-            )
-        if (
-            self._partition_details.partitionIntervalType
-            == PartitionIntervalTypes.INTEGER_RANGE
-        ):
-            return self.client.query(self.table).filter(
-                get_integer_range_filter(
-                    Column(partition_field),
-                    self._partition_details.partitionIntegerRangeStart,
-                    self._partition_details.partitionIntegerRangeEnd,
-                )
-            )
-        return self.client.query(self.table).filter(
-            build_query_filter(
-                [
-                    (
-                        Column(partition_field),
-                        "ge",
-                        dispatch_to_date_or_datetime(
-                            self._partition_details.partitionInterval,
-                            text(self._partition_details.partitionIntervalUnit.value),
-                            type_,
-                        ),
-                    )
-                ],
-                False,
-            )
-        )
+    def get_columns(self):
+        """get columns from entity"""
+        return list(inspect(self.raw_dataset).c)

metadata/sampler/sqlalchemy/snowflake/sampler.py

@@ -0,0 +1,95 @@
+#  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.
+"""
+Helper module to handle data sampling
+for the profiler
+"""
+
+from typing import Dict, Optional, Union
+
+from sqlalchemy import Table, func, text
+from sqlalchemy.sql.selectable import CTE
+
+from metadata.generated.schema.entity.data.table import (
+    ProfileSampleType,
+    SamplingMethodType,
+)
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    DataStorageConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+
+
+class SnowflakeSampler(SQASampler):
+    """
+    Generates a sample of the data to not
+    run the query in the whole table.
+    """
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        sample_config: Optional[SampleConfig] = None,
+        partition_details: Optional[Dict] = None,
+        sample_query: Optional[str] = None,
+        storage_config: DataStorageConfig = None,
+        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+        **kwargs,
+    ):
+        super().__init__(
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            sample_config=sample_config,
+            partition_details=partition_details,
+            sample_query=sample_query,
+            storage_config=storage_config,
+            sample_data_count=sample_data_count,
+            **kwargs,
+        )
+        self.sampling_method_type = func.bernoulli
+        if (
+            sample_config
+            and sample_config.sampling_method_type == SamplingMethodType.SYSTEM
+        ):
+            self.sampling_method_type = func.system
+
+    def set_tablesample(self, selectable: Table):
+        """Set the TABLESAMPLE clause for Snowflake
+        Args:
+            selectable (Table): _description_
+        """
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
+            return selectable.tablesample(
+                self.sampling_method_type(self.sample_config.profile_sample or 100)
+            )
+
+        return selectable.tablesample(
+            func.ROW(text(f"{self.sample_config.profile_sample or 100} ROWS"))
+        )
+
+    def get_sample_query(self, *, column=None) -> CTE:
+        """get query for sample data"""
+        rnd = self._base_sample_query(column).cte(
+            f"{self.raw_dataset.__tablename__}_rnd"
+        )
+        query = self.client.query(rnd)
+        return query.cte(f"{self.raw_dataset.__tablename__}_sample")

metadata/{profiler/processor/sampler → sampler}/sqlalchemy/trino/sampler.py

@@ -16,7 +16,7 @@ from sqlalchemy import inspect, or_, text
 
 from metadata.profiler.orm.registry import FLOAT_SET
 from metadata.profiler.processor.handle_partition import RANDOM_LABEL
-from metadata.profiler.processor.sampler.sqlalchemy.sampler import SQASampler
+from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class TrinoSampler(SQASampler):
@@ -34,8 +34,10 @@ class TrinoSampler(SQASampler):
         super().__init__(*args, **kwargs)
 
     def _base_sample_query(self, column, label=None):
-        sqa_columns = [col for col in inspect(self.table).c if col.name != RANDOM_LABEL]
-        entity = self.table if column is None else column
+        sqa_columns = [
+            col for col in inspect(self.raw_dataset).c if col.name != RANDOM_LABEL
+        ]
+        entity = self.raw_dataset if column is None else column
         return self.client.query(entity, label).where(
             or_(
                 *[

metadata/utils/constants.py

@@ -27,6 +27,51 @@ from metadata.generated.schema.entity.data.topic import Topic
 from metadata.generated.schema.entity.domains.dataProduct import DataProduct
 from metadata.generated.schema.entity.domains.domain import Domain
 from metadata.generated.schema.entity.services.apiService import ApiService
+from metadata.generated.schema.entity.services.connections.database.bigTableConnection import (
+    BigtableType,
+)
+from metadata.generated.schema.entity.services.connections.database.couchbaseConnection import (
+    CouchbaseType,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeType,
+)
+from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
+    DeltaLakeType,
+)
+from metadata.generated.schema.entity.services.connections.database.domoDatabaseConnection import (
+    DomoDatabaseType,
+)
+from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
+    DorisType,
+)
+from metadata.generated.schema.entity.services.connections.database.druidConnection import (
+    DruidType,
+)
+from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
+    DynamoDBType,
+)
+from metadata.generated.schema.entity.services.connections.database.glueConnection import (
+    GlueType,
+)
+from metadata.generated.schema.entity.services.connections.database.icebergConnection import (
+    IcebergType,
+)
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBType,
+)
+from metadata.generated.schema.entity.services.connections.database.salesforceConnection import (
+    SalesforceType,
+)
+from metadata.generated.schema.entity.services.connections.database.sapErpConnection import (
+    SapErpType,
+)
+from metadata.generated.schema.entity.services.connections.database.sasConnection import (
+    SasType,
+)
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    DatabricksType,
+)
 from metadata.generated.schema.entity.services.dashboardService import DashboardService
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.services.messagingService import MessagingService
@@ -111,3 +156,21 @@ ENTITY_REFERENCE_TYPE_MAP = {
 }
 
 CUSTOM_CONNECTOR_PREFIX = "custom"
+
+NON_SQA_DATABASE_CONNECTIONS = (
+    DatalakeType.Datalake.value,
+    BigtableType.BigTable.value,
+    CouchbaseType.Couchbase.value,
+    DatabricksType.UnityCatalog.value,
+    DeltaLakeType.DeltaLake.value,
+    DomoDatabaseType.DomoDatabase.value,
+    DorisType.Doris.value,
+    DruidType.Druid.value,
+    DynamoDBType.DynamoDB.value,
+    GlueType.Glue.value,
+    IcebergType.Iceberg.value,
+    MongoDBType.MongoDB.value,
+    SalesforceType.Salesforce.value,
+    SapErpType.SapErp.value,
+    SasType.SAS.value,
+)

metadata/utils/helpers.py

@@ -360,8 +360,7 @@ def clean_uri(uri: Union[str, Url]) -> str:
     make it http://localhost:9000
     """
     # force a string of the given Uri if needed
-    if isinstance(uri, Url):
-        uri = str(uri)
+    uri = str(uri)
     return uri[:-1] if uri.endswith("/") else uri
 
 

metadata/utils/logger.py

@@ -52,6 +52,7 @@ class Loggers(Enum):
     OMETA = "OMetaAPI"
     CLI = "Metadata"
     PROFILER = "Profiler"
+    SAMPLER = "Sampler"
     PII = "PII"
     INGESTION = "Ingestion"
     UTILS = "Utils"
@@ -103,6 +104,14 @@ def profiler_logger():
     return logging.getLogger(Loggers.PROFILER.value)
 
 
+def sampler_logger():
+    """
+    Method to get the SAMPLER logger
+    """
+
+    return logging.getLogger(Loggers.SAMPLER.value)
+
+
 def pii_logger():
     """
     Method to get the PROFILER logger

metadata/utils/profiler_utils.py

@@ -20,6 +20,11 @@ from typing import Optional, Tuple
 import sqlparse
 from pydantic import BaseModel
 
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -100,3 +105,40 @@ def set_cache(cache: defaultdict, key: str, value):
             cache[key_] = value
             break
         cache = cache[key_]
+
+
+def get_context_entities(
+    entity: Table, metadata: OpenMetadata
+) -> Tuple[DatabaseSchema, Database, DatabaseService]:
+    """Based on the table, get all the parent entities"""
+    schema_entity = None
+    database_entity = None
+    db_service = None
+
+    if entity.databaseSchema:
+        schema_entity_list = metadata.es_search_from_fqn(
+            entity_type=DatabaseSchema,
+            fqn_search_string=entity.databaseSchema.fullyQualifiedName,
+            fields="databaseSchemaProfilerConfig",
+        )
+        if schema_entity_list:
+            schema_entity = schema_entity_list[0]
+
+    if entity.database:
+        database_entity_list = metadata.es_search_from_fqn(
+            entity_type=Database,
+            fqn_search_string=entity.database.fullyQualifiedName,
+            fields="databaseProfilerConfig",
+        )
+        if database_entity_list:
+            database_entity = database_entity_list[0]
+
+    if entity.service:
+        db_service_list = metadata.es_search_from_fqn(
+            entity_type=DatabaseService,
+            fqn_search_string=entity.service.fullyQualifiedName,
+        )
+        if db_service_list:
+            db_service = db_service_list[0]
+
+    return schema_entity, database_entity, db_service

metadata/utils/service_spec/default.py

@@ -4,12 +4,18 @@ Default service specs for services.
 
 from typing import Optional
 
+from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import (
+    SQATestSuiteInterface,
+)
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
     SQAProfilerInterface,
 )
+from metadata.sampler.sqlalchemy.sampler import SQASampler
 from metadata.utils.importer import get_class_path
 from metadata.utils.service_spec.service_spec import BaseSpec
 
 
 class DefaultDatabaseSpec(BaseSpec):
     profiler_class: Optional[str] = get_class_path(SQAProfilerInterface)
+    sampler_class: Optional[str] = get_class_path(SQASampler)
+    test_suite_class: Optional[str] = get_class_path(SQATestSuiteInterface)

metadata/utils/service_spec/service_spec.py

@@ -6,11 +6,15 @@ from typing import Optional, Type, cast
 
 from pydantic import model_validator
 
+from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface
 from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.profiler.interface.profiler_interface import ProfilerInterface
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.importer import (
     TYPE_SEPARATOR,
+    DynamicImportException,
     get_class_path,
     get_module_dir,
     import_from_module,
@@ -43,9 +47,11 @@ class BaseSpec(BaseModel):
     """
 
     profiler_class: Optional[str] = None
+    test_suite_class: Optional[str] = None
     metadata_source_class: str
     lineage_source_class: Optional[str] = None
     usage_source_class: Optional[str] = None
+    sampler_class: Optional[str] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -97,3 +103,44 @@ def import_source_class(
         Type[Source],
         import_from_module(spec.model_dump()[field]),
     )
+
+
+def import_profiler_class(
+    service_type: ServiceType, source_type: str
+) -> Type[ProfilerInterface]:
+    class_path = BaseSpec.get_for_source(service_type, source_type).profiler_class
+    return cast(Type[ProfilerInterface], import_from_module(class_path))
+
+
+def import_test_suite_class(
+    service_type: ServiceType,
+    source_type: str,
+    source_config_type: Optional[str] = None,
+) -> Type[TestSuiteInterface]:
+    try:
+        class_path = BaseSpec.get_for_source(service_type, source_type).test_suite_class
+    except DynamicImportException:
+        if source_config_type:
+            class_path = BaseSpec.get_for_source(
+                service_type, source_config_type.lower()
+            ).test_suite_class
+        else:
+            raise
+    return cast(Type[TestSuiteInterface], import_from_module(class_path))
+
+
+def import_sampler_class(
+    service_type: ServiceType,
+    source_type: str,
+    source_config_type: Optional[str] = None,
+) -> Type[SamplerInterface]:
+    try:
+        class_path = BaseSpec.get_for_source(service_type, source_type).sampler_class
+    except DynamicImportException:
+        if source_config_type:
+            class_path = BaseSpec.get_for_source(
+                service_type, source_config_type.lower()
+            ).sampler_class
+        else:
+            raise
+    return cast(Type[SamplerInterface], import_from_module(class_path))

metadata/utils/sqlalchemy_utils.py

@@ -180,7 +180,7 @@ def get_table_comment_results(
     """
     self.table_comment_result: Dict[Tuple[str, str], str] = {}
     self.current_db: str = database
-    result = connection.execute(query)
+    result = connection.execute(query).fetchall()
     self.table_comment_result[(table_name, schema)] = result