openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0rc2__py3-none-any.whl

metadata/utils/partition.py

@@ -11,19 +11,61 @@
 
 """Partition utility"""
 
-from typing import Optional
+from typing import List, Optional
 
 from metadata.generated.schema.entity.data.table import (
-    IntervalType,
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
     PartitionIntervalUnit,
     PartitionProfilerConfig,
     Table,
+    TablePartition,
+    TableProfilerConfig,
 )
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
 
 
+def validate_athena_injected_partitioning(
+    table_partitions: TablePartition,
+    table_profiler_config: Optional[TableProfilerConfig],
+    profiler_partitioning_config: Optional[PartitionProfilerConfig],
+) -> None:
+    """Validate Athena partitioning. Injected partition need to be defined
+    in the table profiler c onfig for the profiler to work correctly. We'll throw an
+    error if the partitioning is not defined in the table profiler config.
+
+    Attr:
+        entity (Table): entity table
+    """
+    error_msg = (
+        "Table profiler config is missing for table with injected partitioning. Please define "
+        "the partitioning in the table profiler config for column {column_name}. "
+        "For more information, visit "
+        "https://docs.open-metadata.org/v1.3.x/connectors/ingestion/workflows/profiler#profiler-options "
+    )
+
+    column_partitions: Optional[List[PartitionColumnDetails]] = table_partitions.columns
+    if not column_partitions:
+        raise RuntimeError("Table parition is set but no columns are defined.")
+
+    for column_partition in column_partitions:
+        if column_partition.intervalType == PartitionIntervalTypes.INJECTED:
+            if table_profiler_config is None or profiler_partitioning_config is None:
+                raise RuntimeError(
+                    error_msg.format(column_name=column_partition.columnName)
+                )
+
+            if (
+                profiler_partitioning_config.partitionColumnName
+                != column_partition.columnName
+            ):
+                raise RuntimeError(
+                    error_msg.format(column_name=column_partition.columnName)
+                )
+
+
 def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
     """Build PartitionProfilerConfig object from entity
 
@@ -32,59 +74,82 @@ def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
     Returns:
         PartitionProfilerConfig
     """
-    if (
-        hasattr(entity, "tableProfilerConfig")
-        and hasattr(entity.tableProfilerConfig, "partitioning")
-        and entity.tableProfilerConfig.partitioning
-    ):
-        return entity.tableProfilerConfig.partitioning
-
-    if (
-        hasattr(entity, "serviceType")
-        and entity.serviceType == DatabaseServiceType.BigQuery
-    ):
-        if hasattr(entity, "tablePartition") and entity.tablePartition:
-            if entity.tablePartition.intervalType == IntervalType.TIME_UNIT:
+    # Gather service type information
+    service_type = getattr(entity, "serviceType", None)
+
+    # Gather table partitioning information
+    table_partition = getattr(entity, "tablePartition", None)
+
+    # Profiler config
+    profiler_partitioning_config: Optional[PartitionProfilerConfig] = None
+    profiler_config: Optional[TableProfilerConfig] = getattr(
+        entity, "tableProfilerConfig", None
+    )
+    if profiler_config:
+        profiler_partitioning_config = getattr(profiler_config, "partitioning", None)
+
+    if table_partition and service_type == DatabaseServiceType.Athena:
+        # if table is an Athena table and it has been partitioned we need to validate injected partitioning
+        validate_athena_injected_partitioning(
+            table_partition, profiler_config, profiler_partitioning_config
+        )
+        return profiler_partitioning_config
+
+    if profiler_partitioning_config:
+        # if table has partitioning defined in the profiler config, return it
+        return profiler_partitioning_config
+
+    if service_type == DatabaseServiceType.BigQuery:
+        if table_partition:
+            column_partitions: Optional[
+                List[PartitionColumnDetails]
+            ] = entity.tablePartition.columns
+            if not column_partitions:
+                raise TypeError("table partition missing. Skipping table")
+
+            partiton = column_partitions[0]
+
+            if partiton.intervalType == PartitionIntervalTypes.TIME_UNIT:
                 return PartitionProfilerConfig(
                     enablePartitioning=True,
-                    partitionColumnName=entity.tablePartition.columns[0],
+                    partitionColumnName=partiton.columnName,
                     partitionIntervalUnit=PartitionIntervalUnit.DAY
-                    if entity.tablePartition.interval != "HOUR"
-                    else entity.tablePartition.interval,
+                    if partiton.interval != "HOUR"
+                    else partiton.interval,
                     partitionInterval=1,
-                    partitionIntervalType=entity.tablePartition.intervalType.value,
+                    partitionIntervalType=partiton.intervalType.value,
                     partitionValues=None,
                     partitionIntegerRangeStart=None,
                     partitionIntegerRangeEnd=None,
                 )
-            if entity.tablePartition.intervalType == IntervalType.INGESTION_TIME:
+            if partiton.intervalType == PartitionIntervalTypes.INGESTION_TIME:
                 return PartitionProfilerConfig(
                     enablePartitioning=True,
                     partitionColumnName="_PARTITIONDATE"
-                    if entity.tablePartition.interval == "DAY"
+                    if partiton.interval == "DAY"
                     else "_PARTITIONTIME",
                     partitionIntervalUnit=PartitionIntervalUnit.DAY
-                    if entity.tablePartition.interval != "HOUR"
-                    else entity.tablePartition.interval,
+                    if partiton.interval != "HOUR"
+                    else partiton.interval,
                     partitionInterval=1,
-                    partitionIntervalType=entity.tablePartition.intervalType.value,
+                    partitionIntervalType=partiton.intervalType.value,
                     partitionValues=None,
                     partitionIntegerRangeStart=None,
                     partitionIntegerRangeEnd=None,
                 )
-            if entity.tablePartition.intervalType == IntervalType.INTEGER_RANGE:
+            if partiton.intervalType == PartitionIntervalTypes.INTEGER_RANGE:
                 return PartitionProfilerConfig(
                     enablePartitioning=True,
-                    partitionColumnName=entity.tablePartition.columns[0],
+                    partitionColumnName=partiton.columnName,
                     partitionIntervalUnit=None,
                     partitionInterval=None,
-                    partitionIntervalType=entity.tablePartition.intervalType.value,
+                    partitionIntervalType=partiton.intervalType.value,
                     partitionValues=None,
                     partitionIntegerRangeStart=1,
                     partitionIntegerRangeEnd=10000,
                 )
             raise TypeError(
-                f"Unsupported partition type {entity.tablePartition.intervalType}. Skipping table"
+                f"Unsupported partition type {partiton.intervalType}. Skipping table"
             )
 
     return None

metadata/utils/sqlalchemy_utils.py

@@ -12,10 +12,15 @@
 """
 Module for sqlalchemy dialect utils
 """
-
+import traceback
 from typing import Dict, Optional, Tuple
 
 from sqlalchemy.engine import Engine, reflection
+from sqlalchemy.schema import CreateTable, MetaData
+
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 @reflection.cache
@@ -110,3 +115,57 @@ def get_display_datatype(
     if scale is not None and precision is not None:
         return f"{col_type}({str(precision)},{str(scale)})"
     return col_type
+
+
+def convert_numpy_to_list(data):
+    """
+    Recursively converts numpy arrays to lists in a nested data structure.
+    """
+    import numpy as np  # pylint: disable=import-outside-toplevel
+
+    if isinstance(data, np.ndarray):
+        return data.tolist()
+    if isinstance(data, list):
+        return [convert_numpy_to_list(item) for item in data]
+    if isinstance(data, dict):
+        return {key: convert_numpy_to_list(value) for key, value in data.items()}
+    return data
+
+
+@reflection.cache
+def get_all_table_ddls(
+    self, connection, query, schema_name, **kw
+):  # pylint: disable=unused-argument
+    """
+    Method to fetch ddl of all available tables
+    """
+    try:
+        self.all_table_ddls: Dict[Tuple[str, str], str] = {}
+        self.current_db: str = schema_name
+        meta = MetaData()
+        meta.reflect(bind=connection.engine, schema=schema_name)
+        for table in meta.sorted_tables or []:
+            self.all_table_ddls[(table.schema, table.name)] = str(CreateTable(table))
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Failed to get table ddls for {schema_name}: {exc}")
+
+
+def get_table_ddl_wrapper(
+    self, connection, query, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    if not hasattr(self, "all_table_ddls") or self.current_db != schema:
+        self.get_all_table_ddls(connection, query, schema)
+    return self.all_table_ddls.get((schema, table_name), "")
+
+
+def get_table_ddl(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    return get_table_ddl_wrapper(
+        self,
+        connection=connection,
+        query=None,
+        table_name=table_name,
+        schema=schema,
+    )

metadata/utils/ssl_manager.py

@@ -0,0 +1,181 @@
+#  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.
+
+
+"""
+Module to manage SSL certificates
+"""
+import os
+import tempfile
+from functools import singledispatch, singledispatchmethod
+from typing import Optional, Union, cast
+
+from pydantic import SecretStr
+
+from metadata.generated.schema.entity.services.connections.dashboard.qlikSenseConnection import (
+    QlikSenseConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
+    DorisConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.greenplumConnection import (
+    GreenplumConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
+    PostgresConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
+    RedshiftConnection,
+)
+from metadata.generated.schema.entity.services.connections.messaging.kafkaConnection import (
+    KafkaConnection,
+)
+from metadata.generated.schema.security.ssl import verifySSLConfig
+from metadata.ingestion.connections.builders import init_empty_connection_arguments
+
+
+class SSLManager:
+    "SSL Manager to manage SSL certificates for service connections"
+
+    def __init__(self, ca=None, key=None, cert=None):
+        self.temp_files = []
+        self.ca_file_path = None
+        self.cert_file_path = None
+        self.key_file_path = None
+        if ca:
+            self.ca_file_path = self.create_temp_file(ca)
+        if cert:
+            self.cert_file_path = self.create_temp_file(cert)
+        if key:
+            self.key_file_path = self.create_temp_file(key)
+
+    def create_temp_file(self, content: SecretStr):
+        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+            temp_file.write(content.get_secret_value().encode())
+            temp_file.close()
+        self.temp_files.append(temp_file.name)
+        return temp_file.name
+
+    def cleanup_temp_files(self):
+        for temp_file in self.temp_files:
+            try:
+                os.remove(temp_file)
+            except FileNotFoundError:
+                pass
+        self.temp_files = []
+
+    @singledispatchmethod
+    def setup_ssl(self, connection):
+        raise NotImplementedError(f"Connection {type(connection)} type not supported")
+
+    @setup_ssl.register(MysqlConnection)
+    @setup_ssl.register(DorisConnection)
+    def _(self, connection):
+        # Use the temporary file paths for SSL configuration
+        connection = cast(Union[MysqlConnection, DorisConnection], connection)
+        connection.connectionArguments = (
+            connection.connectionArguments or init_empty_connection_arguments()
+        )
+        ssl_args = connection.connectionArguments.__root__.get("ssl", {})
+        if connection.sslConfig.__root__.caCertificate:
+            ssl_args["ssl_ca"] = self.ca_file_path
+        if connection.sslConfig.__root__.sslCertificate:
+            ssl_args["ssl_cert"] = self.cert_file_path
+        if connection.sslConfig.__root__.sslKey:
+            ssl_args["ssl_key"] = self.key_file_path
+        connection.connectionArguments.__root__["ssl"] = ssl_args
+        return connection
+
+    @setup_ssl.register(PostgresConnection)
+    @setup_ssl.register(RedshiftConnection)
+    @setup_ssl.register(GreenplumConnection)
+    def _(self, connection):
+        connection = cast(
+            Union[PostgresConnection, RedshiftConnection, GreenplumConnection],
+            connection,
+        )
+
+        if not connection.connectionArguments:
+            connection.connectionArguments = init_empty_connection_arguments()
+        connection.connectionArguments.__root__["sslmode"] = connection.sslMode.value
+        if connection.sslMode in (
+            verifySSLConfig.SslMode.verify_ca,
+            verifySSLConfig.SslMode.verify_full,
+        ):
+            if self.ca_file_path:
+                connection.connectionArguments.__root__[
+                    "sslrootcert"
+                ] = self.ca_file_path
+            else:
+                raise ValueError(
+                    "CA certificate is required for SSL mode verify-ca or verify-full"
+                )
+        return connection
+
+    @setup_ssl.register(QlikSenseConnection)
+    def _(self, connection):
+        return {
+            "ca_certs": self.ca_file_path,
+            "certfile": self.cert_file_path,
+            "keyfile": self.key_file_path,
+            "check_hostname": connection.validateHostName,
+        }
+
+    @setup_ssl.register(KafkaConnection)
+    def _(self, connection):
+        connection = cast(KafkaConnection, connection)
+        connection.schemaRegistryConfig["ssl.ca.location"] = self.ca_file_path
+        connection.schemaRegistryConfig["ssl.key.location"] = self.key_file_path
+        connection.schemaRegistryConfig[
+            "ssl.certificate.location"
+        ] = self.cert_file_path
+        return connection
+
+
+@singledispatch
+def check_ssl_and_init(_):
+    return None
+
+
+@check_ssl_and_init.register(MysqlConnection)
+@check_ssl_and_init.register(DorisConnection)
+def _(connection):
+    service_connection = cast(Union[MysqlConnection, DorisConnection], connection)
+    ssl: Optional[verifySSLConfig.SslConfig] = service_connection.sslConfig
+    if ssl and (
+        ssl.__root__.caCertificate or ssl.__root__.sslCertificate or ssl.__root__.sslKey
+    ):
+        return SSLManager(
+            ca=ssl.__root__.caCertificate,
+            cert=ssl.__root__.sslCertificate,
+            key=ssl.__root__.sslKey,
+        )
+    return None
+
+
+@check_ssl_and_init.register(PostgresConnection)
+@check_ssl_and_init.register(RedshiftConnection)
+@check_ssl_and_init.register(GreenplumConnection)
+def _(connection):
+    connection = cast(
+        Union[PostgresConnection, RedshiftConnection, GreenplumConnection],
+        connection,
+    )
+    if connection.sslMode:
+        return SSLManager(
+            ca=connection.sslConfig.__root__.caCertificate
+            if connection.sslConfig
+            else None
+        )
+    return None

metadata/utils/tag_utils.py

@@ -14,13 +14,14 @@ Tag utils Module
 
 import functools
 import traceback
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, Union
 
 from metadata.generated.schema.api.classification.createClassification import (
     CreateClassificationRequest,
 )
 from metadata.generated.schema.api.classification.createTag import CreateTagRequest
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
@@ -82,29 +83,41 @@ def get_ometa_tag_and_classification(
 
 @functools.lru_cache(maxsize=512)
 def get_tag_label(
-    metadata: OpenMetadata, tag_name: str, classification_name: str
+    metadata: OpenMetadata,
+    tag_name: str,
+    classification_name: Optional[str],
+    tag_type: Union[Tag, GlossaryTerm] = Tag,
 ) -> Optional[TagLabel]:
     """
     Returns the tag label if the tag is created
     """
     try:
-        # Build the tag FQN
-        tag_fqn = fqn.build(
-            metadata,
-            Tag,
-            classification_name=classification_name,
-            tag_name=tag_name,
-        )
+        if tag_type == Tag:
+            # Build the tag FQN
+            tag_fqn = fqn.build(
+                metadata,
+                tag_type,
+                classification_name=classification_name,
+                tag_name=tag_name,
+            )
+            source = TagSource.Classification.value
+
+        if tag_type == GlossaryTerm:
+            tag_fqn = tag_name
+            source = TagSource.Glossary.value
 
         # Check if the tag exists
-        tag = metadata.get_by_name(entity=Tag, fqn=tag_fqn)
+        tag = metadata.get_by_name(entity=tag_type, fqn=tag_fqn)
         if tag:
             return TagLabel(
                 tagFQN=tag_fqn,
                 labelType=LabelType.Automated.value,
                 state=State.Suggested.value,
-                source=TagSource.Classification.value,
+                source=source,
             )
+
+        logger.warning(f"Tag does not exist: {tag_fqn}")
+
     except Exception as err:
         logger.debug(traceback.format_exc())
         logger.error(f"Error processing tag label: {err}")
@@ -114,8 +127,9 @@ def get_tag_label(
 def get_tag_labels(
     metadata: OpenMetadata,
     tags: List[str],
-    classification_name: str,
+    classification_name: Optional[str] = None,
     include_tags: bool = True,
+    tag_type: Union[Tag, GlossaryTerm] = Tag,
 ) -> Optional[List[TagLabel]]:
     """
     Method to create tag labels from the collected tags
@@ -125,7 +139,10 @@ def get_tag_labels(
         for tag in tags:
             try:
                 tag_label = get_tag_label(
-                    metadata, tag_name=tag, classification_name=classification_name
+                    metadata,
+                    tag_name=tag,
+                    classification_name=classification_name,
+                    tag_type=tag_type,
                 )
                 if tag_label:
                     tag_labels_list.append(tag_label)

metadata/utils/test_utils.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.
+"""
+Utility functions for testing
+"""
+from contextlib import contextmanager
+
+
+class MultipleException(Exception):
+    def __init__(self, exceptions):
+        self.exceptions = exceptions
+        super().__init__(f"Multiple exceptions occurred: {exceptions}")
+
+
+class ErrorHandler:
+    """
+    A context manager that accumulates errors and raises them at the end of the block.
+    Useful for cleaning up resources and ensuring that all errors are raised at the end of a test.
+    Example:
+    ```
+    from metadata.utils.test_utils import accumulate_errors
+    with accumulate_errors() as error_handler:
+        error_handler.try_execute(lambda : 1 / 0)
+        error_handler.try_execute(print, "Hello, World!")
+    ```
+
+    ```
+    > Hello, World!
+    > Traceback (most recent call last):
+    >  ...
+    > ZeroDivisionError: division by zero
+    ```
+    """
+
+    def __init__(self):
+        self.errors = []
+
+    def try_execute(self, func, *args, **kwargs):
+        try:
+            func(*args, **kwargs)
+        except Exception as exc:
+            self.errors.append(exc)
+
+    def raise_if_errors(self):
+        if len(self.errors) == 1:
+            raise self.errors[0]
+        if len(self.errors) > 1:
+            raise MultipleException(self.errors)
+
+
+@contextmanager
+def accumulate_errors():
+    error_handler = ErrorHandler()
+    try:
+        yield error_handler
+    finally:
+        error_handler.raise_if_errors()

metadata/workflow/application.py

@@ -68,7 +68,12 @@ class AppRunner(Step, ABC):
         """App logic to execute"""
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata) -> "Step":
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ) -> "Step":
         config = OpenMetadataApplicationConfig.parse_obj(config_dict)
         return cls(config=config, metadata=metadata)
 

metadata/workflow/base.py

@@ -102,6 +102,7 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         # We create the ometa client at the workflow level and pass it to the steps
         self.metadata_config = metadata_config
         self.metadata = create_ometa_client(metadata_config)
+        self.set_ingestion_pipeline_status(state=PipelineState.running)
 
         self.post_init()
 
@@ -180,7 +181,6 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         pipeline_state = PipelineState.success
         self.timer.trigger()
         try:
-            self.set_ingestion_pipeline_status(state=PipelineState.running)
             self.execute_internal()
 
             if SUCCESS_THRESHOLD_VALUE <= self.calculate_success() < 100:

metadata/workflow/ingestion.py

@@ -41,6 +41,7 @@ from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step, Summary
 from metadata.ingestion.api.steps import BulkSink, Processor, Sink, Source, Stage
 from metadata.ingestion.models.custom_types import ServiceWithConnectionType
+from metadata.profiler.api.models import ProfilerProcessorConfig
 from metadata.utils.class_helper import (
     get_service_class_from_service_type,
     get_service_type_from_source_type,
@@ -201,3 +202,21 @@ class IngestionWorkflow(BaseWorkflow, ABC):
                     f"Unknown error getting service connection for service name [{service_name}]"
                     f" using the secrets manager provider [{self.metadata.config.secretsManagerProvider}]: {exc}"
                 )
+
+    def validate(self):
+        try:
+            if (
+                not self.config.source.serviceConnection.__root__.config.supportsProfiler
+            ):
+                raise AttributeError()
+        except AttributeError:
+            if ProfilerProcessorConfig.parse_obj(
+                self.config.processor.dict().get("config")
+            ).ignoreValidation:
+                logger.debug(
+                    f"Profiler is not supported for the service connection: {self.config.source.serviceConnection}"
+                )
+                return
+            raise WorkflowExecutionError(
+                f"Profiler is not supported for the service connection: {self.config.source.serviceConnection}"
+            )

metadata/workflow/metadata.py

@@ -57,7 +57,15 @@ class MetadataWorkflow(IngestionWorkflow):
             )
         )
 
-        source: Source = source_class.create(self.config.source.dict(), self.metadata)
+        pipeline_name = (
+            self.ingestion_pipeline.fullyQualifiedName.__root__
+            if self.ingestion_pipeline
+            else None
+        )
+
+        source: Source = source_class.create(
+            self.config.source.dict(), self.metadata, pipeline_name
+        )
         logger.debug(f"Source type:{source_type},{source_class} configured")
         source.prepare()
         logger.debug(f"Source type:{source_type},{source_class}  prepared")