openmetadata-ingestion 1.5.13.0__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/class_helper.py

@@ -22,6 +22,9 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
     PipelineType,
 )
 from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.apiServiceMetadataPipeline import (
+    ApiServiceMetadataPipeline,
+)
 from metadata.generated.schema.metadataIngestion.dashboardServiceMetadataPipeline import (
     DashboardServiceMetadataPipeline,
 )
@@ -65,6 +68,7 @@ from metadata.generated.schema.metadataIngestion.testSuitePipeline import (
 from metadata.generated.schema.metadataIngestion.workflow import SourceConfig
 
 SERVICE_TYPE_REF = {
+    ServiceType.Api.value: "apiService",
     ServiceType.Database.value: "databaseService",
     ServiceType.Dashboard.value: "dashboardService",
     ServiceType.Pipeline.value: "pipelineService",
@@ -76,6 +80,7 @@ SERVICE_TYPE_REF = {
 }
 
 SOURCE_CONFIG_TYPE_INGESTION = {
+    ApiServiceMetadataPipeline.__name__: PipelineType.metadata,
     DatabaseServiceMetadataPipeline.__name__: PipelineType.metadata,
     DatabaseServiceQueryUsagePipeline.__name__: PipelineType.usage,
     DatabaseServiceQueryLineagePipeline.__name__: PipelineType.lineage,

metadata/utils/collaborative_super.py

@@ -0,0 +1,28 @@
+"""
+This module provides utilties for building a collaborative constructor library.
+"""
+
+from abc import ABC
+
+
+class Root:
+    """Root class for any class that needs to implement a colllaborative constructor but
+    might end up at the end of the inheritance chain. Since python's object.__init__ is not
+    a collaborative constructor, we need to have a root class that has a collaborative constructor.
+    """
+
+    __terminal__ = {object, ABC}
+
+    def __init__(self, *args, **kwargs):
+        """Collaborative constructor"""
+        super_class = None
+        for cls, super_class in zip(
+            self.__class__.mro()[:-1], self.__class__.mro()[1:]
+        ):
+            if cls is Root:
+                break
+        for cls in self.__terminal__:
+            if super_class is cls:
+                break
+        else:
+            super().__init__(*args, **kwargs)

metadata/utils/constants.py

@@ -26,6 +26,52 @@ from metadata.generated.schema.entity.data.table import Table
 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
@@ -77,6 +123,7 @@ ENTITY_REFERENCE_CLASS_MAP = {
     # Service Entities
     "databaseService": DatabaseService,
     "messagingService": MessagingService,
+    "apiService": ApiService,
     "dashboardService": DashboardService,
     "pipelineService": PipelineService,
     "storageService": StorageService,
@@ -109,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/constraints.py

@@ -0,0 +1,65 @@
+#  Copyright 2024 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.
+
+"""
+Define constraints helper methods useful for the metadata ingestion
+"""
+
+from typing import Dict, List
+
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    ConstraintType,
+    RelationshipType,
+)
+from metadata.ingestion.ometa.utils import model_str
+
+
+def _is_column_unique(column: Dict, columns: List[Column]) -> bool:
+    """
+    Method to check if the column in unique in the table
+    """
+    if column and len(column) > 0:
+        constrained_column = column[0]
+        for col in columns or []:
+            if model_str(col.name) == constrained_column:
+                if col.constraint and col.constraint.value in {
+                    ConstraintType.UNIQUE.value,
+                    ConstraintType.PRIMARY_KEY.value,
+                }:
+                    return True
+                break
+    return False
+
+
+def get_relationship_type(
+    column: Dict, referred_table_columns: List[Column], columns: List[Column]
+) -> str:
+    """
+    Determine the type of relationship (one-to-one, one-to-many, etc.)
+    """
+    # Check if the column is unique in the current table
+    is_unique_in_current_table = _is_column_unique(
+        column.get("constrained_columns"), columns
+    )
+
+    # Check if the referred column is unique in the referred table
+    is_unique_in_referred_table = _is_column_unique(
+        column.get("referred_columns"), referred_table_columns
+    )
+
+    if is_unique_in_current_table and is_unique_in_referred_table:
+        return RelationshipType.ONE_TO_ONE
+    if is_unique_in_current_table:
+        return RelationshipType.ONE_TO_MANY
+    if is_unique_in_referred_table:
+        return RelationshipType.MANY_TO_ONE
+    return RelationshipType.MANY_TO_MANY

metadata/utils/dict.py

@@ -0,0 +1,8 @@
+"""
+A custom dictionary class that extends functionality.
+"""
+
+
+class ExtendedDict(dict):
+    def lower_case_keys(self):
+        return {k.lower(): v for k, v in self.items()}