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

metadata/ingestion/source/database/external_table_lineage_mixin.py

@@ -0,0 +1,82 @@
+#  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.
+"""
+External Table Lineage Mixin
+"""
+
+import traceback
+from abc import ABC
+from typing import Iterable
+
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.type.entityLineage import EntitiesEdge
+from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.ingestion.api.models import Either
+from metadata.utils import fqn
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class ExternalTableLineageMixin(ABC):
+    """
+    This mixin class is for deriving lineage between external table and container source/
+    """
+
+    def yield_external_table_lineage(self) -> Iterable[AddLineageRequest]:
+        """
+        Yield external table lineage
+        """
+        for table_qualified_tuple, location in self.external_location_map.items() or []:
+            try:
+                location_entity = self.metadata.es_search_container_by_path(
+                    full_path=location
+                )
+                database_name, schema_name, table_name = table_qualified_tuple
+
+                table_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=Table,
+                    service_name=self.context.get().database_service,
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    table_name=table_name,
+                    skip_es_search=True,
+                )
+                table_entity = self.metadata.es_search_from_fqn(
+                    entity_type=Table,
+                    fqn_search_string=table_fqn,
+                )
+
+                if (
+                    location_entity
+                    and location_entity[0]
+                    and table_entity
+                    and table_entity[0]
+                ):
+                    yield Either(
+                        right=AddLineageRequest(
+                            edge=EntitiesEdge(
+                                fromEntity=EntityReference(
+                                    id=location_entity[0].id,
+                                    type="container",
+                                ),
+                                toEntity=EntityReference(
+                                    id=table_entity[0].id,
+                                    type="table",
+                                ),
+                            )
+                        )
+                    )
+            except Exception as exc:
+                logger.warning(f"Failed to yield external table lineage due to - {exc}")
+                logger.debug(traceback.format_exc())

metadata/ingestion/source/database/glue/metadata.py

@@ -81,7 +81,9 @@ class GlueSource(DatabaseServiceSource):
         self.test_connection()
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: GlueConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, GlueConnection):
@@ -97,7 +99,7 @@ class GlueSource(DatabaseServiceSource):
             yield DatabasePage(**page)
 
     def _get_glue_tables(self):
-        schema_name = self.context.database_schema
+        schema_name = self.context.get().database_schema
         paginator = self.glue.get_paginator("get_tables")
         paginator_response = paginator.paginate(DatabaseName=schema_name)
         for page in paginator_response:
@@ -124,7 +126,7 @@ class GlueSource(DatabaseServiceSource):
                         database_fqn = fqn.build(
                             self.metadata,
                             entity_type=Database,
-                            service_name=self.context.database_service,
+                            service_name=self.context.get().database_service,
                             database_name=schema.CatalogId,
                         )
                         if filter_by_database(
@@ -162,7 +164,7 @@ class GlueSource(DatabaseServiceSource):
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -176,8 +178,8 @@ class GlueSource(DatabaseServiceSource):
                     schema_fqn = fqn.build(
                         self.metadata,
                         entity_type=DatabaseSchema,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
                         schema_name=schema.Name,
                     )
                     if filter_by_schema(
@@ -211,11 +213,11 @@ class GlueSource(DatabaseServiceSource):
                 database=fqn.build(
                     metadata=self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                 ),
                 sourceUrl=self.get_source_url(
-                    database_name=self.context.database,
+                    database_name=self.context.get().database,
                     schema_name=schema_name,
                 ),
             )
@@ -230,7 +232,7 @@ class GlueSource(DatabaseServiceSource):
 
         :return: tables or views, depending on config
         """
-        schema_name = self.context.database_schema
+        schema_name = self.context.get().database_schema
 
         for page in self._get_glue_tables():
             for table in page.TableList:
@@ -240,9 +242,9 @@ class GlueSource(DatabaseServiceSource):
                     table_fqn = fqn.build(
                         self.metadata,
                         entity_type=Table,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
+                        schema_name=self.context.get().database_schema,
                         table_name=table_name,
                     )
                     if filter_by_table(
@@ -269,7 +271,7 @@ class GlueSource(DatabaseServiceSource):
                     elif table.TableType == "VIRTUAL_VIEW":
                         table_type = TableType.View
 
-                    self.context.table_data = table
+                    self.context.get().table_data = table
                     yield table_name, table_type
                 except Exception as exc:
                     self.status.failed(
@@ -288,7 +290,7 @@ class GlueSource(DatabaseServiceSource):
         Prepare a table request and pass it to the sink
         """
         table_name, table_type = table_name_and_type
-        table = self.context.table_data
+        table = self.context.get().table_data
         table_constraints = None
         try:
             columns = self.get_columns(table.StorageDescriptor)
@@ -302,14 +304,14 @@ class GlueSource(DatabaseServiceSource):
                 databaseSchema=fqn.build(
                     metadata=self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                 ),
                 sourceUrl=self.get_source_url(
                     table_name=table_name,
-                    schema_name=self.context.database_schema,
-                    database_name=self.context.database,
+                    schema_name=self.context.get().database_schema,
+                    database_name=self.context.get().database,
                 ),
             )
             yield Either(right=table_request)
@@ -350,7 +352,7 @@ class GlueSource(DatabaseServiceSource):
             yield self._get_column_object(column)
 
         # process table regular columns info
-        for column in self.context.table_data.PartitionKeys:
+        for column in self.context.get().table_data.PartitionKeys:
             yield self._get_column_object(column)
 
     def standardize_table_name(self, _: str, table: str) -> str:

metadata/ingestion/source/database/greenplum/connection.py

@@ -22,13 +22,11 @@ from metadata.generated.schema.entity.automations.workflow import (
 )
 from metadata.generated.schema.entity.services.connections.database.greenplumConnection import (
     GreenplumConnection,
-    SslMode,
 )
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
     get_connection_url_common,
-    init_empty_connection_arguments,
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -39,14 +37,6 @@ def get_connection(connection: GreenplumConnection) -> Engine:
     """
     Create connection
     """
-    if connection.sslMode:
-        if not connection.connectionArguments:
-            connection.connectionArguments = init_empty_connection_arguments()
-        connection.connectionArguments.__root__["sslmode"] = connection.sslMode.value
-        if connection.sslMode in (SslMode.verify_ca, SslMode.verify_full):
-            connection.connectionArguments.__root__[
-                "sslrootcert"
-            ] = connection.sslConfig.__root__.certificatePath
     return create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url_common,

metadata/ingestion/source/database/greenplum/metadata.py

@@ -21,7 +21,8 @@ from sqlalchemy.engine.reflection import Inspector
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.table import (
-    IntervalType,
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
     TablePartition,
     TableType,
 )
@@ -66,9 +67,9 @@ logger = ingestion_logger()
 
 
 INTERVAL_TYPE_MAP = {
-    "list": IntervalType.COLUMN_VALUE.value,
-    "hash": IntervalType.COLUMN_VALUE.value,
-    "range": IntervalType.TIME_UNIT.value,
+    "list": PartitionIntervalTypes.COLUMN_VALUE,
+    "hash": PartitionIntervalTypes.COLUMN_VALUE,
+    "range": PartitionIntervalTypes.TIME_UNIT,
 }
 
 RELKIND_MAP = {
@@ -117,7 +118,12 @@ class GreenplumSource(CommonDbSourceService, MultiDBSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadataConnection):
+    def create(
+        cls,
+        config_dict,
+        metadata: OpenMetadataConnection,
+        pipeline_name: Optional[str] = None,
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: GreenplumConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, GreenplumConnection):
@@ -163,7 +169,7 @@ class GreenplumSource(CommonDbSourceService, MultiDBSource):
                 database_fqn = fqn.build(
                     self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=new_database,
                 )
 
@@ -187,18 +193,27 @@ class GreenplumSource(CommonDbSourceService, MultiDBSource):
 
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
-    ) -> Tuple[bool, TablePartition]:
+    ) -> Tuple[bool, Optional[TablePartition]]:
         result = self.engine.execute(
             GREENPLUM_PARTITION_DETAILS.format(
                 table_name=table_name, schema_name=schema_name
             )
         ).all()
+
         if result:
             partition_details = TablePartition(
-                intervalType=INTERVAL_TYPE_MAP.get(
-                    result[0].partition_strategy, IntervalType.COLUMN_VALUE.value
-                ),
-                columns=[row.column_name for row in result if row.column_name],
+                columns=[
+                    PartitionColumnDetails(
+                        columnName=row.column_name,
+                        intervalType=INTERVAL_TYPE_MAP.get(
+                            result[0].partition_strategy,
+                            PartitionIntervalTypes.COLUMN_VALUE,
+                        ),
+                        interval=None,
+                    )
+                    for row in result
+                    if row.column_name
+                ]
             )
             return True, partition_details
         return False, None

metadata/ingestion/source/database/hive/metadata.py

@@ -12,10 +12,11 @@
 Hive source methods.
 """
 
-from typing import Tuple
+import traceback
+from typing import Optional, Tuple
 
 from pyhive.sqlalchemy_hive import HiveDialect
-from sqlalchemy.inspection import inspect
+from sqlalchemy.engine.reflection import Inspector
 
 from metadata.generated.schema.entity.services.connections.database.hiveConnection import (
     HiveConnection,
@@ -56,7 +57,9 @@ class HiveSource(CommonDbSourceService):
     service_connection: HiveConnection
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: HiveConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, HiveConnection):
@@ -93,4 +96,28 @@ class HiveSource(CommonDbSourceService):
             self.engine = get_metastore_connection(
                 self.service_connection.metastoreConnection
             )
-        self.inspector = inspect(self.engine)
+        self._connection_map = {}  # Lazy init as well
+        self._inspector_map = {}
+
+    def get_schema_definition(  # pylint: disable=unused-argument
+        self, table_type: str, table_name: str, schema_name: str, inspector: Inspector
+    ) -> Optional[str]:
+        """
+        Get the DDL statement or View Definition for a table
+        """
+        try:
+            schema_definition = inspector.get_view_definition(table_name, schema_name)
+            schema_definition = (
+                str(schema_definition).strip()
+                if schema_definition is not None
+                else None
+            )
+            return schema_definition
+
+        except NotImplementedError:
+            logger.warning("Schema definition not implemented")
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to fetch schema definition for {table_name}: {exc}")
+        return None

metadata/ingestion/source/database/iceberg/helper.py

@@ -20,7 +20,12 @@ import pyiceberg.partitioning
 import pyiceberg.table
 import pyiceberg.types
 
-from metadata.generated.schema.entity.data.table import Column, Constraint, DataType
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    Constraint,
+    DataType,
+    PartitionIntervalTypes,
+)
 
 
 def namespace_to_str(namespace: tuple[str]) -> str:
@@ -45,13 +50,45 @@ def get_table_name_as_str(table: pyiceberg.table.Table) -> str:
 def get_column_from_partition(
     columns: Tuple[pyiceberg.types.NestedField, ...],
     partition: pyiceberg.partitioning.PartitionField,
-) -> str:
+) -> Optional[str]:
     """Returns the Column Name belonging to a partition."""
     # A Partition in Iceberg has a Source Column to which a Transformation is applied.
     # We need to return the Source Column name.
-    return [
-        column.name for column in columns if column.field_id == partition.source_id
-    ][0]
+    return next(
+        (column.name for column in columns if column.field_id == partition.source_id),
+        None,
+    )
+
+
+def get_column_partition_type(
+    columns: Tuple[pyiceberg.types.NestedField, ...],
+    partition: pyiceberg.partitioning.PartitionField,
+) -> Optional[PartitionIntervalTypes]:
+    """Get the partition type for a given partition column."""
+    iceberg_interval_type_map = {
+        "INT": PartitionIntervalTypes.INTEGER_RANGE,
+        **dict.fromkeys(
+            ["TIME", "DATE", "TIMESTAMP", "TIMESTAMPTZ"],
+            PartitionIntervalTypes.TIME_UNIT,
+        ),
+    }
+
+    data_type = str(
+        next(
+            (
+                column.field_type
+                for column in columns
+                if column.field_id == partition.source_id
+            ),
+            "",
+        )
+    )
+    if not data_type.isalpha():
+        return None
+
+    return iceberg_interval_type_map.get(
+        data_type.upper(), PartitionIntervalTypes.COLUMN_VALUE
+    )
 
 
 def get_owner_from_table(

metadata/ingestion/source/database/iceberg/metadata.py

@@ -82,7 +82,9 @@ class IcebergSource(DatabaseServiceSource):
         self.test_connection()
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: IcebergConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, IcebergConnection):
@@ -110,7 +112,7 @@ class IcebergSource(DatabaseServiceSource):
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -125,8 +127,8 @@ class IcebergSource(DatabaseServiceSource):
                 schema_fqn = fqn.build(
                     self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                     schema_name=namespace_name,
                 )
                 if filter_by_schema(
@@ -160,8 +162,8 @@ class IcebergSource(DatabaseServiceSource):
                 database=fqn.build(
                     metadata=self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                 ),
             )
         )
@@ -171,7 +173,7 @@ class IcebergSource(DatabaseServiceSource):
         Prepares the table name to be sent to stage.
         Filtering happens here.
         """
-        namespace = self.context.database_schema
+        namespace = self.context.get().database_schema
 
         for table_identifier in self.iceberg.list_tables(namespace):
             try:
@@ -180,9 +182,9 @@ class IcebergSource(DatabaseServiceSource):
                 table_fqn = fqn.build(
                     self.metadata,
                     entity_type=Table,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                     table_name=table_name,
                 )
                 if filter_by_table(
@@ -197,7 +199,7 @@ class IcebergSource(DatabaseServiceSource):
                     )
                     continue
 
-                self.context.iceberg_table = table
+                self.context.get().iceberg_table = table
                 yield table_name, TableType.Regular
             except pyiceberg.exceptions.NoSuchPropertyException:
                 logger.warning(
@@ -224,7 +226,7 @@ class IcebergSource(DatabaseServiceSource):
 
     def get_owner_ref(self, table_name: str) -> Optional[EntityReference]:
         owner = get_owner_from_table(
-            self.context.iceberg_table, self.service_connection.ownershipProperty
+            self.context.get().iceberg_table, self.service_connection.ownershipProperty
         )
         try:
             if owner:
@@ -245,7 +247,7 @@ class IcebergSource(DatabaseServiceSource):
         Also, update the self.inspector value to the current db.
         """
         table_name, table_type = table_name_and_type
-        iceberg_table = self.context.iceberg_table
+        iceberg_table = self.context.get().iceberg_table
         try:
             owner = self.get_owner_ref(table_name)
             table = IcebergTable.from_pyiceberg(
@@ -261,9 +263,9 @@ class IcebergSource(DatabaseServiceSource):
                 databaseSchema=fqn.build(
                     metadata=self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                 ),
             )
             yield Either(right=table_request)

metadata/ingestion/source/database/iceberg/models.py

@@ -20,6 +20,7 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.entity.data.table import (
     Column,
+    PartitionColumnDetails,
     TablePartition,
     TableType,
 )
@@ -27,6 +28,7 @@ from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.source.database.iceberg.helper import (
     IcebergColumnParser,
     get_column_from_partition,
+    get_column_partition_type,
 )
 
 
@@ -57,10 +59,16 @@ class IcebergTable(BaseModel):
             columns=[IcebergColumnParser.parse(column) for column in iceberg_columns],
             tablePartition=TablePartition(
                 columns=[
-                    get_column_from_partition(iceberg_columns, partition)
+                    PartitionColumnDetails(
+                        columnName=get_column_from_partition(
+                            iceberg_columns, partition
+                        ),
+                        intervalType=get_column_partition_type(
+                            iceberg_columns, partition
+                        ),
+                        interval=None,
+                    )
                     for partition in table.spec().fields
-                ],
-                intervalType=None,
-                interval=None,
+                ]
             ),
         )

metadata/ingestion/source/database/impala/metadata.py

@@ -13,6 +13,7 @@ Impala source methods.
 """
 
 import re
+from typing import Optional
 
 from impala.sqlalchemy import ImpalaDialect, _impala_type_to_sqlalchemy_type
 from sqlalchemy import types, util
@@ -180,7 +181,9 @@ class ImpalaSource(CommonDbSourceService):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: ImpalaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, ImpalaConnection):