openmetadata-ingestion 1.7.0.1__py3-none-any.whl → 1.7.1.0__py3-none-any.whl

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -12,18 +12,16 @@
 Helper module to handle data sampling
 for the profiler
 """
-import traceback
-from typing import Dict, List, Optional, Union
+from typing import Dict, Optional, Union
 
 from sqlalchemy import Column
 from sqlalchemy import Table as SqaTable
-from sqlalchemy import inspect, text
+from sqlalchemy import text
 from sqlalchemy.orm import Query
 
 from metadata.generated.schema.entity.data.table import (
     ProfileSampleType,
     Table,
-    TableData,
     TableType,
 )
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
@@ -130,62 +128,3 @@ class BigQuerySampler(SQASampler):
             )
 
         return super().get_sample_query(column=column)
-
-    def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:
-        """
-        Use the sampler to retrieve sample data rows as per limit given by user
-
-        Args:
-            columns (Optional[List]): List of columns to fetch
-        Returns:
-            TableData to be added to the Table Entity
-        """
-        if self.sample_query:
-            return self._fetch_sample_data_from_user_query()
-
-        # Add new RandomNumFn column
-        ds = self.get_dataset()
-        if not columns:
-            sqa_columns = [col for col in inspect(ds).c if col.name != RANDOM_LABEL]
-        else:
-            # 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
-                for col in inspect(ds).c
-                if col.name != RANDOM_LABEL and col.name in names
-            ]
-        with self.get_client() as client:
-            try:
-                query = (
-                    str(client.query(*sqa_columns)).replace(
-                        f"`{ds.__table_args__['schema']}`.`{ds.__tablename__}`",
-                        f"`{self.entity.database.name}`.`{ds.__table_args__['schema']}`.`{ds.__tablename__}`",
-                    )
-                    + f" LIMIT {self.sample_limit}"
-                )
-                sqa_sample = client.execute(text(query))
-            except Exception:
-                logger.debug(
-                    "Cannot fetch sample data with random sampling. Falling back to 100 rows."
-                )
-                logger.debug(traceback.format_exc())
-                ds_inspect = inspect(self.raw_dataset)
-                sqa_columns = list(ds_inspect.c)
-
-                schema = ds_inspect.class_.__table_args__["schema"]
-                table = ds_inspect.class_.__tablename__
-
-                query = (
-                    str(client.query(*sqa_columns)).replace(
-                        f"`{schema}`.`{table}`",
-                        f"`{self.entity.database.name}`.`{schema}`.`{table}`",
-                    )
-                    + f" LIMIT {self.sample_limit}"
-                )
-                sqa_sample = client.execute(text(query))
-            return TableData(
-                columns=[column.name for column in sqa_columns],
-                rows=[list(row) for row in sqa_sample],
-            )

metadata/sampler/sqlalchemy/sampler.py

@@ -13,7 +13,6 @@ Helper module to handle data sampling
 for the profiler
 """
 import hashlib
-import traceback
 from typing import List, Optional, Union, cast
 
 from sqlalchemy import Column, inspect, text
@@ -100,16 +99,16 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         """
         # only sample the column if we are computing a column metric to limit the amount of data scaned
         selectable = self.set_tablesample(self.raw_dataset.__table__)
+        client = self.get_client()
 
         entity = selectable if column is None else selectable.c.get(column.key)
-        with self.get_client() as client:
-            if label is not None:
-                query = client.query(entity, label)
-            else:
-                query = client.query(entity)
+        if label is not None:
+            query = client.query(entity, label)
+        else:
+            query = client.query(entity)
 
-            if self.partition_details:
-                query = self.get_partitioned_query(query)
+        if self.partition_details:
+            query = self.get_partitioned_query(query)
         return query
 
     def get_sampler_table_name(self) -> str:
@@ -125,19 +124,18 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
+        client = self.get_client()
         if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
             rnd = self._base_sample_query(
                 column,
                 (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
             ).cte(f"{self.get_sampler_table_name()}_rnd")
-            with self.get_client() as client:
-                session_query = client.query(rnd)
+            session_query = client.query(rnd)
             return session_query.where(
                 rnd.c.random <= self.sample_config.profileSample
             ).cte(f"{self.get_sampler_table_name()}_sample")
 
-        with self.get_client() as client:
-            table_query = client.query(self.raw_dataset)
+        table_query = client.query(self.raw_dataset)
         session_query = self._base_sample_query(
             column,
             (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL)
@@ -198,31 +196,20 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 for col in inspect(ds).c
                 if col.name != RANDOM_LABEL and col.name in names
             ]
+
         with self.get_client() as client:
-            try:
-                sqa_sample = (
-                    client.query(*sqa_columns)
-                    .select_from(ds)
-                    .limit(self.sample_limit)
-                    .all()
-                )
-            except Exception:
-                logger.debug(
-                    "Cannot fetch sample data with random sampling. Falling back to 100 rows."
-                )
-                logger.debug(traceback.format_exc())
-                sqa_columns = list(inspect(self.raw_dataset).c)
-                sqa_sample = (
-                    client.query(*sqa_columns)
-                    .select_from(self.raw_dataset)
-                    .limit(100)
-                    .all()
-                )
-            return TableData(
-                columns=[column.name for column in sqa_columns],
-                rows=[list(row) for row in sqa_sample],
+            sqa_sample = (
+                client.query(*sqa_columns)
+                .select_from(ds)
+                .limit(self.sample_limit)
+                .all()
             )
 
+        return TableData(
+            columns=[column.name for column in sqa_columns],
+            rows=[list(row) for row in sqa_sample],
+        )
+
     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.sample_query):
@@ -246,10 +233,11 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
         stmt = text(f"{self.sample_query}")
         stmt = stmt.columns(*list(inspect(self.raw_dataset).c))
-        with self.get_client() as client:
-            return client.query(stmt.subquery()).cte(
-                f"{self.get_sampler_table_name()}_user_sampled"
-            )
+        return (
+            self.get_client()
+            .query(stmt.subquery())
+            .cte(f"{self.get_sampler_table_name()}_user_sampled")
+        )
 
     def _partitioned_table(self) -> Query:
         """Return the Query object for partitioned tables"""
@@ -266,9 +254,13 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         )
         if query is not None:
             return query.filter(partition_filter)
-        with self.get_client() as client:
-            return client.query(self.raw_dataset).filter(partition_filter)
+        return self.get_client().query(self.raw_dataset).filter(partition_filter)
 
     def get_columns(self):
         """get columns from entity"""
         return list(inspect(self.raw_dataset).c)
+
+    def close(self):
+        """Close the connection"""
+        self.get_client().close()
+        self.connection.pool.dispose()

metadata/utils/datalake/datalake_utils.py

@@ -131,7 +131,13 @@ class DataFrameColumnParser:
             shuffle: whether to shuffle the dataframe list or not if sample is True. (default: False)
         """
         data_frame = cls._get_data_frame(data_frame, sample, shuffle)
-        if file_type == SupportedTypes.PARQUET:
+        if file_type in {
+            SupportedTypes.PARQUET,
+            SupportedTypes.PARQUET_PQ,
+            SupportedTypes.PARQUET_PQT,
+            SupportedTypes.PARQUET_PARQ,
+            SupportedTypes.PARQUET_SNAPPY,
+        }:
             parser = ParquetDataFrameColumnParser(data_frame)
         elif file_type in {
             SupportedTypes.JSON,
@@ -248,7 +254,7 @@ class GenericDataFrameColumnParser:
             data_frame (DataFrame)
             column_name (string)
         """
-        data_type = None
+        data_type = None  # default to string
         try:
             if data_frame[column_name].dtypes.name == "object" and any(
                 data_frame[column_name].dropna().values
@@ -310,7 +316,7 @@ class GenericDataFrameColumnParser:
                 f"Failed to distinguish data type for column {column_name}, Falling back to {data_type}, exc: {err}"
             )
             logger.debug(traceback.format_exc())
-        return data_type
+        return data_type or DataType.STRING
 
     @classmethod
     def unique_json_structure(cls, dicts: List[Dict]) -> Dict:

metadata/utils/fqn.py

@@ -151,7 +151,7 @@ def build(
 def _(
     metadata: Optional[OpenMetadata],
     *,
-    service_name: str,
+    service_name: Optional[str],
     database_name: Optional[str],
     schema_name: Optional[str],
     table_name: str,
@@ -161,7 +161,7 @@ def _(
     """
     Building logic for tables
     :param metadata: OMeta client
-    :param service_name: Service Name to filter
+    :param service_name: Service Name to filter or None
     :param database_name: DB name or None
     :param schema_name: Schema name or None
     :param table_name: Table name
@@ -180,8 +180,8 @@ def _(
             service_name=service_name,
         )
 
-    # if entity not found in ES proceed to build FQN with database_name and schema_name
-    if not entity and database_name and schema_name:
+    # if entity not found in ES proceed to build FQN with service_name, database_name and schema_name
+    if not entity and service_name and database_name and schema_name:
         fqn = _build(service_name, database_name, schema_name, table_name)
         return [fqn] if fetch_multiple_entities else fqn
     if entity and fetch_multiple_entities:

metadata/utils/service_spec/default.py

@@ -7,6 +7,9 @@ from typing import Optional
 from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import (
     SQATestSuiteInterface,
 )
+from metadata.data_quality.validations.runtime_param_setter.base_diff_params_setter import (
+    BaseTableParameter,
+)
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
     SQAProfilerInterface,
 )
@@ -19,3 +22,4 @@ 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)
+    data_diff: Optional[str] = get_class_path(BaseTableParameter)

metadata/utils/service_spec/service_spec.py

@@ -52,6 +52,7 @@ class BaseSpec(BaseModel):
     lineage_source_class: Optional[str] = None
     usage_source_class: Optional[str] = None
     sampler_class: Optional[str] = None
+    data_diff: Optional[str] = None
 
     @model_validator(mode="before")
     @classmethod

metadata/utils/sqa_utils.py

@@ -191,6 +191,21 @@ def get_query_filter_for_runner(kwargs: Dict) -> Optional[BinaryExpression]:
     return filter_
 
 
+def get_query_group_by_for_runner(kwargs: Dict) -> Optional[BinaryExpression]:
+    """Get query group by from kwargs. IMPORTANT, this will update the original dictionary
+    passed in the function argument.
+
+    Args:
+        kwargs (Dict): kwargs
+    """
+    if "query_group_by_" in kwargs:
+        group_by_ = kwargs.pop("query_group_by_")
+    else:
+        group_by_ = None
+
+    return group_by_
+
+
 def handle_array(
     query: Query, column: Column, table: Union[DeclarativeMeta, AliasedClass]
 ) -> Query:

metadata/workflow/base.py

@@ -219,7 +219,14 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
             if self.workflow_config.successThreshold <= self.calculate_success() < 100:
                 pipeline_state = PipelineState.partialSuccess
 
-        # Any unhandled exception breaking the workflow should update the status
+            # If there's any steps that should raise a failed status,
+            # raise it here to set the pipeline status as failed
+            try:
+                self.raise_from_status_internal()
+            except WorkflowExecutionError:
+                pipeline_state = PipelineState.failed
+
+        # Any unhandled exception should blow up the execution
         except Exception as err:
             pipeline_state = PipelineState.failed
             raise err

metadata/workflow/profiler.py

@@ -12,9 +12,6 @@
 Workflow definition for the profiler
 """
 
-from metadata.generated.schema.metadataIngestion.workflow import (
-    OpenMetadataWorkflowConfig,
-)
 from metadata.ingestion.api.steps import Processor, Sink
 from metadata.ingestion.connections.test_connections import (
     raise_test_connection_exception,
@@ -40,12 +37,6 @@ class ProfilerWorkflow(IngestionWorkflow):
     this workflow. No need to do anything here if this does not pass
     """
 
-    def __init__(self, config: OpenMetadataWorkflowConfig):
-        super().__init__(config)
-
-        # Validate that we can properly reach the source database
-        self.test_connection()
-
     def _get_source_class(self):
         if self.config.source.serviceName:
             self.import_source_class()
@@ -58,6 +49,15 @@ class ProfilerWorkflow(IngestionWorkflow):
         return OpenMetadataSourceExt
 
     def set_steps(self):
+        # TODO: Clean after https://github.com/open-metadata/OpenMetadata/issues/21259
+        # We are forcing the secret evaluation to "ignore" null secrets down the line
+        # Remove this when the issue above is fixed and empty secrets migrated
+        source_config_class = type(self.config.source.serviceConnection.root.config)
+        dumped_config = self.config.source.serviceConnection.root.config.model_dump()
+        self.config.source.serviceConnection.root.config = (
+            source_config_class.model_validate(dumped_config)
+        )
+
         # NOTE: Call test_connection to update host value before creating the source class
         self.test_connection()
 

metadata/workflow/workflow_status_mixin.py

@@ -17,7 +17,6 @@ from datetime import datetime
 from enum import Enum
 from typing import Optional, Tuple
 
-from metadata.config.common import WorkflowExecutionError
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     IngestionPipeline,
     PipelineState,
@@ -129,13 +128,8 @@ class WorkflowStatusMixin:
     def raise_from_status(self, raise_warnings=False):
         """
         Method to raise error if failed execution
-        and updating Ingestion Pipeline Status
         """
-        try:
-            self.raise_from_status_internal(raise_warnings)
-        except WorkflowExecutionError as err:
-            self.set_ingestion_pipeline_status(PipelineState.failed)
-            raise err
+        self.raise_from_status_internal(raise_warnings)
 
     def result_status(self) -> WorkflowResultStatus:
         """