openmetadata-ingestion 1.5.6.0__py3-none-any.whl → 1.5.7.1__py3-none-any.whl

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

@@ -20,6 +20,7 @@ from pydantic import BaseModel
 class GlueSchema(BaseModel):
     CatalogId: Optional[str] = None
     Name: str
+    Description: Optional[str] = None
 
 
 class DatabasePage(BaseModel):

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

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from typing import Optional
 
 from sqlalchemy.engine import Engine
@@ -26,6 +27,7 @@ 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
@@ -41,8 +43,16 @@ def get_connection(connection: PinotDBConnection) -> Engine:
     """
     Create connection
     """
+    # TODO: Rename database field to DatabaseSchema
+    # Pinot does not support multi database concept
+    if connection.database is not None:
+        if not connection.connectionArguments:
+            connection.connectionArguments = init_empty_connection_arguments()
+        connection.connectionArguments.root["database"] = connection.database
+    connection_copy = deepcopy(connection)
+    connection_copy.database = None
     return create_generic_db_connection(
-        connection=connection,
+        connection=connection_copy,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )

metadata/ingestion/source/database/pinotdb/metadata.py

@@ -9,7 +9,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """PinotDb source module"""
-from typing import Optional
+from typing import Iterable, Optional
+
+from pinotdb import sqlalchemy as pinot_sqlalchemy
+from sqlalchemy import types
 
 from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import (
     PinotDBConnection,
@@ -22,6 +25,31 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.common_db_source import CommonDbSourceService
 
 
+def get_type_custom(data_type, field_size):
+    type_map = {
+        "int": types.BigInteger,
+        "long": types.BigInteger,
+        "float": types.Float,
+        "double": types.Numeric,
+        # BOOLEAN, is added after release 0.7.1.
+        # In release 0.7.1 and older releases, BOOLEAN is equivalent to STRING.
+        "boolean": types.Boolean,
+        "timestamp": types.TIMESTAMP,
+        "string": types.String,
+        "json": types.JSON,
+        "bytes": types.LargeBinary,
+        "big_decimal": types.DECIMAL,
+        # Complex types
+        "struct": types.BLOB,
+        "map": types.BLOB,
+        "array": types.ARRAY,
+    }
+    return type_map.get(data_type.lower())
+
+
+pinot_sqlalchemy.get_type = get_type_custom
+
+
 class PinotdbSource(CommonDbSourceService):
     """
     Implements the necessary methods to extract
@@ -39,3 +67,15 @@ class PinotdbSource(CommonDbSourceService):
                 f"Expected PinotdbConnection, but got {connection}"
             )
         return cls(config, metadata)
+
+    def get_database_names(self) -> Iterable[str]:
+        """
+        Default case with a single database.
+
+        It might come informed - or not - from the source.
+
+        Sources with multiple databases should overwrite this and
+        apply the necessary filters.
+        """
+        # TODO: Add databaseDisplayName field in PinotDBConnection
+        yield "default"

metadata/ingestion/source/database/postgres/metadata.py

@@ -13,14 +13,19 @@ Postgres source module
 """
 import traceback
 from collections import namedtuple
-from typing import Iterable, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from sqlalchemy import String as SqlAlchemyString
 from sqlalchemy import sql
 from sqlalchemy.dialects.postgresql.base import PGDialect, ischema_names
 from sqlalchemy.engine import Inspector
 
+from metadata.generated.schema.api.data.createStoredProcedure import (
+    CreateStoredProcedureRequest,
+)
 from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.data.table import (
     PartitionColumnDetails,
     PartitionIntervalTypes,
@@ -36,7 +41,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -46,10 +51,13 @@ from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.mssql.models import STORED_PROC_LANGUAGE_MAP
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
+from metadata.ingestion.source.database.postgres.models import PostgresStoredProcedure
 from metadata.ingestion.source.database.postgres.queries import (
     POSTGRES_GET_ALL_TABLE_PG_POLICY,
     POSTGRES_GET_DB_NAMES,
+    POSTGRES_GET_STORED_PROCEDURES,
     POSTGRES_GET_TABLE_NAMES,
     POSTGRES_PARTITION_DETAILS,
     POSTGRES_SCHEMA_COMMENTS,
@@ -63,6 +71,10 @@ from metadata.ingestion.source.database.postgres.utils import (
     get_table_owner,
     get_view_definition,
 )
+from metadata.ingestion.source.database.stored_procedures_mixin import (
+    QueryByProcedure,
+    StoredProcedureMixin,
+)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database
 from metadata.utils.importer import import_side_effects
@@ -142,7 +154,7 @@ Inspector.get_table_owner = get_etable_owner
 PGDialect.get_foreign_keys = get_foreign_keys
 
 
-class PostgresSource(CommonDbSourceService, MultiDBSource):
+class PostgresSource(CommonDbSourceService, MultiDBSource, StoredProcedureMixin):
     """
     Implements the necessary methods to extract
     Database metadata from Postgres Source
@@ -298,3 +310,62 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
                     stackTrace=traceback.format_exc(),
                 )
             )
+
+    def get_stored_procedures(self) -> Iterable[PostgresStoredProcedure]:
+        """List stored procedures"""
+        if self.source_config.includeStoredProcedures:
+            results = self.engine.execute(POSTGRES_GET_STORED_PROCEDURES).all()
+            for row in results:
+                try:
+                    stored_procedure = PostgresStoredProcedure.model_validate(
+                        dict(row._mapping)
+                    )
+                    yield stored_procedure
+                except Exception as exc:
+                    logger.error()
+                    self.status.failed(
+                        error=StackTraceError(
+                            name=dict(row).get("name", "UNKNOWN"),
+                            error=f"Error parsing Stored Procedure payload: {exc}",
+                            stackTrace=traceback.format_exc(),
+                        )
+                    )
+
+    def yield_stored_procedure(
+        self, stored_procedure
+    ) -> Iterable[Either[CreateStoredProcedureRequest]]:
+        """Prepare the stored procedure payload"""
+        try:
+            stored_procedure_request = CreateStoredProcedureRequest(
+                name=EntityName(stored_procedure.name),
+                description=None,
+                storedProcedureCode=StoredProcedureCode(
+                    language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
+                    code=stored_procedure.definition,
+                ),
+                databaseSchema=fqn.build(
+                    metadata=self.metadata,
+                    entity_type=DatabaseSchema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
+                ),
+            )
+            yield Either(right=stored_procedure_request)
+            self.register_record_stored_proc_request(stored_procedure_request)
+
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=stored_procedure.name,
+                    error=f"Error yielding Stored Procedure [{stored_procedure.name}] due to [{exc}]",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+        """
+        Return the dictionary associating stored procedures to the
+        queries they triggered
+        """
+        return {}

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

@@ -0,0 +1,25 @@
+#  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.
+"""
+Postgres models
+"""
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class PostgresStoredProcedure(BaseModel):
+    """Postgres stored procedure list query results"""
+
+    name: str = Field(alias="procedure_name")
+    schema: str = Field(alias="schema_name")
+    definition: str
+    language: Optional[str] = None

metadata/ingestion/source/database/postgres/queries.py

@@ -211,3 +211,99 @@ POSTGRES_FETCH_FK = """
         n.oid = c.relnamespace
     ORDER BY 1
 """
+
+POSTGRES_GET_JSON_FIELDS = """
+    WITH RECURSIVE json_hierarchy AS (
+        SELECT
+            key AS path,
+            json_typeof(value) AS type,
+            value,
+            json_build_object() AS properties,
+            key AS title
+        FROM
+            {table_name} tbd,
+            LATERAL json_each({column_name}::json)
+    ),
+    build_hierarchy AS (
+        SELECT
+            path,
+            type,
+            title,
+            CASE
+                WHEN type = 'object' THEN
+                    json_build_object(
+                        'title', title,
+                        'type', 'object',
+                        'properties', (
+                            SELECT json_object_agg(
+                                key,
+                                json_build_object(
+                                    'title', key,
+                                    'type', json_typeof(value),
+                                    'properties', (
+                                        CASE
+                                            WHEN json_typeof(value) = 'object' THEN
+                                                (
+                                                    SELECT json_object_agg(
+                                                        key,
+                                                        json_build_object(
+                                                            'title', key,
+                                                            'type', json_typeof(value),
+                                                            'properties',
+                                                            json_build_object()
+                                                        )
+                                                    )
+                                                    FROM json_each(value::json) AS sub_key_value
+                                                )
+                                            ELSE json_build_object()
+                                        END
+                                    )
+                                )
+                            )
+                            FROM json_each(value::json) AS key_value
+                        )
+                    )
+                WHEN type = 'array' THEN
+                    json_build_object(
+                        'title', title,
+                        'type', 'array',
+                        'properties', json_build_object()
+                    )
+                ELSE
+                    json_build_object(
+                        'title', title,
+                        'type', type
+                    )
+            END AS hierarchy
+        FROM
+            json_hierarchy
+    ),
+    aggregate_hierarchy AS (
+        select
+            json_build_object(
+            'title','{column_name}',
+            'type','object',
+            'properties',
+            json_object_agg(
+                path,
+                hierarchy
+            )) AS result
+        FROM
+            build_hierarchy
+    )
+    SELECT
+        result
+    FROM
+        aggregate_hierarchy;
+"""
+
+POSTGRES_GET_STORED_PROCEDURES = """
+    SELECT proname AS procedure_name,
+        nspname AS schema_name,
+        proargtypes AS argument_types,
+        prorettype::regtype AS return_type,
+        prosrc AS definition
+    FROM pg_proc
+    JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
+    WHERE prokind = 'p';
+"""

metadata/ingestion/source/database/salesforce/metadata.py

@@ -203,15 +203,18 @@ class SalesforceSource(DatabaseServiceSource):
                 )
             )
 
-    def get_table_description(self, table_name: str) -> Optional[str]:
+    def get_table_description(
+        self, table_name: str, object_label: Optional[str]
+    ) -> Optional[str]:
         """
         Method to get the table description for salesforce with Tooling API
         """
+        table_description = None
         try:
             result = self.client.toolingexecute(
                 f"query/?q=SELECT+Description+FROM+EntityDefinition+WHERE+QualifiedApiName='{table_name}'"
             )
-            return result["records"][0]["Description"]
+            table_description = result["records"][0]["Description"]
         except KeyError as err:
             logger.warning(
                 f"Unable to get required key from Tooling API response for table [{table_name}]: {err}"
@@ -225,7 +228,7 @@ class SalesforceSource(DatabaseServiceSource):
             logger.warning(
                 f"Unable to get description with Tooling API for table [{table_name}]: {exc}"
             )
-        return None
+        return table_description if table_description else object_label
 
     def yield_table(
         self, table_name_and_type: Tuple[str, TableType]
@@ -241,11 +244,13 @@ class SalesforceSource(DatabaseServiceSource):
                 f"sobjects/{table_name}/describe/",
                 params=None,
             )
-            columns = self.get_columns(salesforce_objects["fields"])
+            columns = self.get_columns(salesforce_objects.get("fields", []))
             table_request = CreateTableRequest(
                 name=EntityName(table_name),
                 tableType=table_type,
-                description=self.get_table_description(table_name),
+                description=self.get_table_description(
+                    table_name, salesforce_objects.get("label")
+                ),
                 columns=columns,
                 tableConstraints=table_constraints,
                 databaseSchema=FullyQualifiedEntityName(

metadata/ingestion/source/pipeline/gluepipeline/metadata.py

@@ -25,6 +25,7 @@ from metadata.generated.schema.entity.data.pipeline import (
     Task,
     TaskStatus,
 )
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.pipeline.gluePipelineConnection import (
     GluePipelineConnection,
 )
@@ -40,14 +41,25 @@ from metadata.generated.schema.type.basic import (
     SourceUrl,
     Timestamp,
 )
+from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.pipeline.gluepipeline.models import (
+    AmazonRedshift,
+    CatalogSource,
+    JDBCSource,
+    JobNodeResponse,
+    S3Source,
+    S3Target,
+)
 from metadata.ingestion.source.pipeline.pipeline_service import PipelineServiceSource
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.time_utils import convert_timestamp_to_milliseconds
+from metadata.utils.time_utils import datetime_to_timestamp
 
 logger = ingestion_logger()
 
@@ -63,6 +75,28 @@ STATUS_MAP = {
     "incomplete": StatusType.Failed,
     "pending": StatusType.Pending,
 }
+TABLE_MODEL_MAP = {
+    "AmazonRedshiftSource": AmazonRedshift,
+    "AmazonRedshiftTarget": AmazonRedshift,
+    "AthenaConnectorSource": JDBCSource,
+    "JDBCConnectorSource": JDBCSource,
+    "JDBCConnectorTarget": JDBCSource,
+    "DirectJDBCSource": CatalogSource,
+    "RedshiftSource": CatalogSource,
+    "RedshiftTarget": CatalogSource,
+    "DirectJDBC": CatalogSource,
+}
+STORAGE_MODEL_MAP = {
+    "S3CsvSource": S3Source,
+    "S3JsonSource": S3Source,
+    "S3ParquetSource": S3Source,
+    "S3HudiSource": S3Source,
+    "S3DeltaSource": S3Source,
+    "S3DirectTarget": S3Target,
+    "S3DeltaDirectTarget": S3Target,
+    "S3GlueParquetTarget": S3Target,
+    "S3HudiDirectTarget": S3Target,
+}
 
 
 class GluepipelineSource(PipelineServiceSource):
@@ -145,9 +179,88 @@ class GluepipelineSource(PipelineServiceSource):
                 downstream_tasks.append(self.task_id_mapping[edges["DestinationId"]])
         return downstream_tasks
 
+    def get_lineage_details(self, job) -> Optional[dict]:
+        """
+        Get the Lineage Details of the pipeline
+        """
+        lineage_details = {"sources": [], "targets": []}
+        try:
+            job_details = JobNodeResponse.model_validate(
+                self.glue.get_job(JobName=job)
+            ).Job
+            if job_details and job_details.config_nodes:
+                nodes = job_details.config_nodes
+                for _, node in nodes.items():
+                    for key, entity in node.items():
+                        table_model, storage_model = None, None
+                        if key in TABLE_MODEL_MAP:
+                            table_model = TABLE_MODEL_MAP[key].model_validate(entity)
+                        elif "Catalog" in key:
+                            table_model = CatalogSource.model_validate(entity)
+                        elif key in STORAGE_MODEL_MAP:
+                            storage_model = STORAGE_MODEL_MAP[key].model_validate(
+                                entity
+                            )
+                        if table_model:
+                            for db_service_name in self.get_db_service_names():
+                                table_entity = self.metadata.get_entity_reference(
+                                    entity=Table,
+                                    fqn=fqn.build(
+                                        metadata=self.metadata,
+                                        entity_type=Table,
+                                        table_name=table_model.table_name,
+                                        database_name=table_model.database_name,
+                                        schema_name=table_model.schema_name,
+                                        service_name=db_service_name,
+                                    ),
+                                )
+                                if table_entity:
+                                    if key.endswith("Source"):
+                                        lineage_details["sources"].append(table_entity)
+                                    else:
+                                        lineage_details["targets"].append(table_entity)
+                                    break
+                        if storage_model:
+                            for path in storage_model.Paths or [storage_model.Path]:
+                                container = self.metadata.es_search_container_by_path(
+                                    full_path=path
+                                )
+                                if container and container[0]:
+                                    storage_entity = EntityReference(
+                                        id=container[0].id,
+                                        type="container",
+                                        name=container[0].name.root,
+                                        fullyQualifiedName=container[
+                                            0
+                                        ].fullyQualifiedName.root,
+                                    )
+                                    if storage_entity:
+                                        if key.endswith("Source"):
+                                            lineage_details["sources"].append(
+                                                storage_entity
+                                            )
+                                        else:
+                                            lineage_details["targets"].append(
+                                                storage_entity
+                                            )
+                                        break
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Failed to get lineage details for job : {job} due to : {exc}"
+            )
+        return lineage_details
+
     def yield_pipeline_status(
         self, pipeline_details: Any
     ) -> Iterable[Either[OMetaPipelineStatus]]:
+        pipeline_fqn = fqn.build(
+            metadata=self.metadata,
+            entity_type=Pipeline,
+            service_name=self.context.get().pipeline_service,
+            pipeline_name=self.context.get().pipeline,
+        )
         for job in self.job_name_list:
             try:
                 runs = self.glue.get_job_runs(JobName=job)
@@ -161,13 +274,13 @@ class GluepipelineSource(PipelineServiceSource):
                                 attempt["JobRunState"].lower(), StatusType.Pending
                             ).value,
                             startTime=Timestamp(
-                                convert_timestamp_to_milliseconds(
-                                    attempt["StartedOn"].timestamp()
+                                datetime_to_timestamp(
+                                    attempt["StartedOn"], milliseconds=True
                                 )
                             ),
                             endTime=Timestamp(
-                                convert_timestamp_to_milliseconds(
-                                    attempt["CompletedOn"].timestamp()
+                                datetime_to_timestamp(
+                                    attempt["CompletedOn"], milliseconds=True
                                 )
                             ),
                         )
@@ -175,20 +288,14 @@ class GluepipelineSource(PipelineServiceSource):
                     pipeline_status = PipelineStatus(
                         taskStatus=task_status,
                         timestamp=Timestamp(
-                            convert_timestamp_to_milliseconds(
-                                attempt["StartedOn"].timestamp()
+                            datetime_to_timestamp(
+                                attempt["StartedOn"], milliseconds=True
                             )
                         ),
                         executionStatus=STATUS_MAP.get(
                             attempt["JobRunState"].lower(), StatusType.Pending
                         ).value,
                     )
-                    pipeline_fqn = fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Pipeline,
-                        service_name=self.context.get().pipeline_service,
-                        pipeline_name=self.context.get().pipeline,
-                    )
                     yield Either(
                         right=OMetaPipelineStatus(
                             pipeline_fqn=pipeline_fqn,
@@ -199,7 +306,7 @@ class GluepipelineSource(PipelineServiceSource):
                 yield Either(
                     left=StackTraceError(
                         name=pipeline_fqn,
-                        error=f"Failed to yield pipeline status: {exc}",
+                        error=f"Failed to yield pipeline status for job {job}: {exc}",
                         stackTrace=traceback.format_exc(),
                     )
                 )
@@ -210,3 +317,42 @@ class GluepipelineSource(PipelineServiceSource):
         """
         Get lineage between pipeline and data sources
         """
+        try:
+            pipeline_fqn = fqn.build(
+                metadata=self.metadata,
+                entity_type=Pipeline,
+                service_name=self.context.get().pipeline_service,
+                pipeline_name=self.context.get().pipeline,
+            )
+
+            pipeline_entity = self.metadata.get_by_name(
+                entity=Pipeline, fqn=pipeline_fqn
+            )
+
+            lineage_details = LineageDetails(
+                pipeline=EntityReference(id=pipeline_entity.id.root, type="pipeline"),
+                source=LineageSource.PipelineLineage,
+            )
+
+            for job in self.job_name_list:
+                lineage_enities = self.get_lineage_details(job)
+                for source in lineage_enities.get("sources"):
+                    for target in lineage_enities.get("targets"):
+                        yield Either(
+                            right=AddLineageRequest(
+                                edge=EntitiesEdge(
+                                    fromEntity=source,
+                                    toEntity=target,
+                                    lineageDetails=lineage_details,
+                                )
+                            )
+                        )
+
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=pipeline_details.get(NAME),
+                    error=f"Wild error ingesting pipeline lineage {pipeline_details} - {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )