openmetadata-ingestion 1.3.0.0.dev0__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

metadata/ingestion/source/messaging/kafka/connection.py

@@ -15,7 +15,7 @@ Source connection handler
 from dataclasses import dataclass
 from typing import Optional, Union
 
-from confluent_kafka.admin import AdminClient
+from confluent_kafka.admin import AdminClient, KafkaException
 from confluent_kafka.avro import AvroConsumer
 from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
 
@@ -35,6 +35,21 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+class InvalidKafkaCreds(Exception):
+    """
+    Class to indicate invalid kafka credentials exception
+    """
+
+
+class SchemaRegistryException(Exception):
+    """
+    Class to indicate invalid schema registry not initialized
+    """
+
+
+TIMEOUT_SECONDS = 10
+
+
 @dataclass
 class KafkaClient:
     def __init__(self, admin_client, schema_registry_client, consumer_client) -> None:
@@ -60,6 +75,14 @@ def get_connection(
         if connection.saslMechanism:
             connection.consumerConfig["sasl.mechanism"] = connection.saslMechanism.value
 
+        if (
+            connection.consumerConfig.get("security.protocol") is None
+            and connection.securityProtocol
+        ):
+            connection.consumerConfig[
+                "security.protocol"
+            ] = connection.securityProtocol.value
+
     if connection.basicAuthUserInfo:
         connection.schemaRegistryConfig = connection.schemaRegistryConfig or {}
         connection.schemaRegistryConfig[
@@ -109,9 +132,27 @@ def test_connection(
     """
 
     def custom_executor():
-        _ = client.admin_client.list_topics().topics
-
-    test_fn = {"GetTopics": custom_executor}
+        try:
+            client.admin_client.list_topics(timeout=TIMEOUT_SECONDS).topics
+        except KafkaException as err:
+            raise InvalidKafkaCreds(
+                f"Failed to fetch topics due to: {err}. "
+                "Please validate credentials and check if you are using correct security protocol"
+            )
+
+    def schema_registry_test():
+        if client.schema_registry_client:
+            client.schema_registry_client.get_subjects()
+        else:
+            raise SchemaRegistryException(
+                "Schema Registry not initialized, please provide schema registry "
+                "credentials in case you want topic schema and sample data to be ingested"
+            )
+
+    test_fn = {
+        "GetTopics": custom_executor,
+        "CheckSchemaRegistry": schema_registry_test,
+    }
 
     test_connection_steps(
         metadata=metadata,

metadata/ingestion/source/mlmodel/sagemaker/metadata.py

@@ -24,6 +24,9 @@ from metadata.generated.schema.entity.data.mlmodel import (
 from metadata.generated.schema.entity.services.connections.mlmodel.sageMakerConnection import (
     SageMakerConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
@@ -128,14 +131,23 @@ class SagemakerSource(MlModelServiceSource):
         """
         Prepare the Request model
         """
-        mlmodel_request = CreateMlModelRequest(
-            name=model.name,
-            algorithm=self._get_algorithm(),  # Setting this to a constant
-            mlStore=self._get_ml_store(model.name),
-            service=self.context.mlmodel_service,
-        )
-        yield mlmodel_request
-        self.register_record(mlmodel_request=mlmodel_request)
+        try:
+            mlmodel_request = CreateMlModelRequest(
+                name=model.name,
+                algorithm=self._get_algorithm(),  # Setting this to a constant
+                mlStore=self._get_ml_store(model.name),
+                service=self.context.mlmodel_service,
+            )
+            yield Either(right=mlmodel_request)
+            self.register_record(mlmodel_request=mlmodel_request)
+        except Exception as exc:  # pylint: disable=broad-except
+            yield Either(
+                left=StackTraceError(
+                    name=model.name,
+                    error=f"Error creating mlmodel: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
 
     def _get_ml_store(  # pylint: disable=arguments-differ
         self,

metadata/ingestion/source/pipeline/spline/metadata.py

@@ -131,7 +131,6 @@ class SplineSource(PipelineServiceSource):
         return None
 
     def _get_table_from_datasource_name(self, datasource: str) -> Optional[Table]:
-
         if (
             not datasource
             and not datasource.startswith("dbfs")

metadata/ingestion/source/search/elasticsearch/metadata.py

@@ -111,7 +111,6 @@ class ElasticsearchSource(SearchServiceSource):
         Method to Get Sample Data of Search Index Entity
         """
         if self.source_config.includeSampleData and self.context.search_index:
-
             sample_data = self.client.search(
                 index=self.context.search_index,
                 q=WILDCARD_SEARCH,

metadata/ingestion/source/storage/storage_service.py

@@ -14,8 +14,6 @@ Base class for ingesting Object Storage services
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Set
 
-from pandas import DataFrame
-
 from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.services.storageService import (
@@ -53,7 +51,10 @@ from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.readers.dataframe.reader_factory import SupportedTypes
 from metadata.readers.models import ConfigSource
 from metadata.utils import fqn
-from metadata.utils.datalake.datalake_utils import fetch_dataframe, get_columns
+from metadata.utils.datalake.datalake_utils import (
+    DataFrameColumnParser,
+    fetch_dataframe,
+)
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.storage_metadata_config import (
     StorageMetadataConfigException,
@@ -67,7 +68,6 @@ OPENMETADATA_TEMPLATE_FILE_NAME = "openmetadata.json"
 
 
 class StorageServiceTopology(ServiceTopology):
-
     root = TopologyNode(
         producer="get_services",
         stages=[
@@ -271,10 +271,10 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             ),
         )
         columns = []
-        if isinstance(data_structure_details, DataFrame):
-            columns = get_columns(data_structure_details)
-        if isinstance(data_structure_details, list) and data_structure_details:
-            columns = get_columns(data_structure_details[0])
+        column_parser = DataFrameColumnParser.create(
+            data_structure_details, SupportedTypes(metadata_entry.structureFormat)
+        )
+        columns = column_parser.get_columns()
         return columns
 
     def _get_columns(

metadata/profiler/interface/pandas/profiler_interface.py

@@ -34,7 +34,10 @@ from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR, SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.datalake.datalake_utils import fetch_col_types, fetch_dataframe
+from metadata.utils.datalake.datalake_utils import (
+    GenericDataFrameColumnParser,
+    fetch_dataframe,
+)
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
@@ -411,7 +414,9 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                 sqalike_columns.append(
                     SQALikeColumn(
                         column_name,
-                        fetch_col_types(self.complex_dataframe_sample[0], column_name),
+                        GenericDataFrameColumnParser.fetch_col_types(
+                            self.complex_dataframe_sample[0], column_name
+                        ),
                     )
                 )
             return sqalike_columns

metadata/profiler/interface/profiler_interface_factory.py

@@ -24,6 +24,9 @@ from metadata.generated.schema.entity.services.connections.database.databricksCo
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    Db2Connection,
+)
 from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
     MariaDBConnection,
 )
@@ -50,6 +53,9 @@ from metadata.profiler.interface.sqlalchemy.bigquery.profiler_interface import (
 from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import (
     DatabricksProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
+    DB2ProfilerInterface,
+)
 from metadata.profiler.interface.sqlalchemy.mariadb.profiler_interface import (
     MariaDBProfilerInterface,
 )
@@ -111,6 +117,7 @@ profilers = {
     TrinoConnection.__name__: TrinoProfilerInterface,
     UnityCatalogConnection.__name__: UnityCatalogProfilerInterface,
     DatabricksConnection.__name__: DatabricksProfilerInterface,
+    Db2Connection.__name__: DB2ProfilerInterface,
 }
 
 profiler_interface_factory.register_many(profilers)

metadata/profiler/interface/sqlalchemy/db2/profiler_interface.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.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+)
+from metadata.utils.logger import profiler_interface_registry_logger
+
+logger = profiler_interface_registry_logger()
+
+
+class DB2ProfilerInterface(SQAProfilerInterface):
+    """
+    Interface to interact with registry supporting
+    sqlalchemy.
+    """
+
+    def _programming_error_static_metric(self, runner, column, exc, session, metrics):
+        # pylint: disable=protected-access
+        if exc.orig and "overflow" in exc.orig._message:
+            logger.info(
+                f"Computing metrics without sum for {runner.table.__tablename__}.{column.name}"
+            )
+            return self._compute_static_metrics_wo_sum(metrics, runner, session, column)
+        return None

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -23,7 +23,7 @@ from datetime import datetime, timezone
 from typing import Dict, List, Optional
 
 from sqlalchemy import Column, inspect, text
-from sqlalchemy.exc import ProgrammingError, ResourceClosedError
+from sqlalchemy.exc import DBAPIError, ProgrammingError, ResourceClosedError
 from sqlalchemy.orm import scoped_session
 
 from metadata.generated.schema.entity.data.table import CustomMetricProfile, TableData
@@ -230,7 +230,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 ],
             )
             return dict(row)
-        except ProgrammingError as exc:
+        except (ProgrammingError, DBAPIError) as exc:
             return self._programming_error_static_metric(
                 runner, column, exc, session, metrics
             )
@@ -579,9 +579,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         """
         Override Programming Error for Static Metrics
         """
-        logger.error(
-            f"Skipping metrics due to {exc} for {runner.table.__tablename__}.{column.name}"
-        )
+        raise exc
 
     def get_columns(self):
         """get columns from entity"""

metadata/profiler/metrics/hybrid/histogram.py

@@ -30,6 +30,7 @@ from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
 
+
 # pylint: disable=too-many-locals
 class Histogram(HybridMetric):
     """

metadata/profiler/orm/functions/conn_test.py

@@ -48,6 +48,7 @@ def _(*_, **__):
 
 @compiles(ConnTestFn, Dialects.Db2)
 @compiles(ConnTestFn, Dialects.IbmDbSa)
+@compiles(ConnTestFn, Dialects.Ibmi)
 def _(*_, **__):
     return "SELECT 42 FROM SYSIBM.SYSDUMMY1;"
 

metadata/profiler/orm/functions/sum.py

@@ -40,6 +40,7 @@ def _(element, compiler, **kw):
 
 
 @compiles(SumFn, Dialects.BigQuery)
+@compiles(SumFn, Dialects.Postgres)
 def _(element, compiler, **kw):
     """Handle case where column type is INTEGER but SUM returns a NUMBER"""
     proc = compiler.process(element.clauses, **kw)

metadata/profiler/orm/functions/table_metric_computer.py

@@ -126,7 +126,6 @@ class AbstractTableMetricComputer(ABC):
         table: Table,
         where_clause: Optional[List[ColumnOperators]] = None,
     ):
-
         query = select(*columns).select_from(table)
         if where_clause:
             query = query.where(*where_clause)

metadata/profiler/orm/registry.py

@@ -64,6 +64,7 @@ class Dialects(Enum):
     Hive = "hive"
     Impala = "impala"
     IbmDbSa = "ibm_db_sa"
+    Ibmi = "ibmi"
     MariaDB = "mariadb"
     MSSQL = "mssql"
     MySQL = "mysql"

metadata/profiler/processor/core.py

@@ -53,7 +53,7 @@ from metadata.profiler.metrics.static.row_count import RowCount
 from metadata.profiler.orm.registry import NOT_COMPUTE
 from metadata.profiler.processor.sample_data_handler import upload_sample_data
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.helpers import calculate_execution_time
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -550,7 +550,7 @@ class Profiler(Generic[TMetric]):
 
         return table_profile
 
-    @calculate_execution_time
+    @calculate_execution_time(store=False)
     def generate_sample_data(self) -> Optional[TableData]:
         """Fetch and ingest sample data
 

metadata/profiler/processor/processor.py

@@ -40,7 +40,6 @@ class ProfilerProcessor(Processor):
     """
 
     def __init__(self, config: OpenMetadataWorkflowConfig):
-
         super().__init__()
 
         self.config = config
@@ -56,7 +55,6 @@ class ProfilerProcessor(Processor):
         return "Profiler"
 
     def _run(self, record: ProfilerSourceAndEntity) -> Either[ProfilerResponse]:
-
         profiler_runner: Profiler = record.profiler_source.get_profiler_runner(
             record.entity, self.profiler_config
         )

metadata/profiler/processor/sample_data_handler.py

@@ -11,7 +11,7 @@
 """
 Profiler Processor Step
 """
-
+import json
 import traceback
 from datetime import datetime
 from functools import singledispatch
@@ -48,7 +48,11 @@ def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -
         sample_storage_config = profiler_interface.storage_config
         if not sample_storage_config:
             return
-        df = pd.DataFrame(data=data.rows, columns=[i.__root__ for i in data.columns])
+        deserialized_data = json.loads(data.json())
+        df = pd.DataFrame(
+            data=deserialized_data.get("rows", []),
+            columns=[i.__root__ for i in data.columns],
+        )
         pq_buffer = BytesIO()
         df.to_parquet(pq_buffer)
         object_key = _get_object_key(

metadata/profiler/source/base/profiler_source.py

@@ -217,7 +217,6 @@ class ProfilerSource(ProfilerSourceInterface):
     def _get_context_entities(
         self, entity: Table
     ) -> Tuple[DatabaseSchema, Database, DatabaseService]:
-
         schema_entity = None
         database_entity = None
         db_service = None

metadata/profiler/source/bigquery/type_mapper.py

@@ -29,7 +29,6 @@ def bigquery_type_mapper(_type_map: dict, col: Column):
     from sqlalchemy_bigquery import STRUCT
 
     def build_struct(_type_map: dict, col: Column):
-
         structs = []
         for child in col.children:
             if child.dataType != DataType.STRUCT:

metadata/readers/file/api_reader.py

@@ -35,7 +35,6 @@ class ApiReader(Reader, ABC):
     """
 
     def __init__(self, credentials: ReadersCredentials):
-
         self._auth_headers = None
         self.credentials = credentials