openmetadata-ingestion 1.5.5.0__py3-none-any.whl → 1.5.7.0__py3-none-any.whl

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/metadata/alationsink/metadata.py

@@ -14,11 +14,17 @@ AlationSink source to extract metadata
 """
 
 import traceback
-from typing import Iterable, Optional
+from typing import Iterable, List, Optional
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
-from metadata.generated.schema.entity.data.table import Column, Constraint, Table
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    Constraint,
+    ConstraintType,
+    Table,
+    TableConstraint,
+)
 from metadata.generated.schema.entity.services.connections.metadata.alationSinkConnection import (
     AlationSinkConnection,
 )
@@ -108,7 +114,9 @@ class AlationsinkSource(Source):
                     ),
                 ),
                 db_username="Test",
-                title=model_str(om_database.name),
+                title=om_database.displayName
+                if om_database.displayName
+                else model_str(om_database.name),
                 description=model_str(om_database.description),
             )
         except Exception as exc:
@@ -129,7 +137,9 @@ class AlationsinkSource(Source):
                 key=fqn._build(  # pylint: disable=protected-access
                     str(alation_datasource_id), model_str(om_schema.name)
                 ),
-                title=model_str(om_schema.name),
+                title=om_schema.displayName
+                if om_schema.displayName
+                else model_str(om_schema.name),
                 description=model_str(om_schema.description),
             )
         except Exception as exc:
@@ -150,7 +160,9 @@ class AlationsinkSource(Source):
                 key=fqn._build(  # pylint: disable=protected-access
                     str(alation_datasource_id), schema_name, model_str(om_table.name)
                 ),
-                title=model_str(om_table.name),
+                title=om_table.displayName
+                if om_table.displayName
+                else model_str(om_table.name),
                 description=model_str(om_table.description),
                 table_type=TABLE_TYPE_MAPPER.get(om_table.tableType, "TABLE"),
                 sql=om_table.schemaDefinition,
@@ -162,14 +174,60 @@ class AlationsinkSource(Source):
             )
         return None
 
-    def _get_column_index(self, om_column: Column) -> Optional[ColumnIndex]:
+    def _update_foreign_key(
+        self,
+        alation_datasource_id: int,
+        om_column: Column,
+        table_constraints: Optional[List[TableConstraint]],
+        column_index: ColumnIndex,
+    ):
+        """
+        Method to update the foreign key metadata in columns index
+        """
+        try:
+            for table_constraint in table_constraints or []:
+                if table_constraint.constraintType == ConstraintType.FOREIGN_KEY:
+                    for i, constraint_column in enumerate(
+                        table_constraint.columns or []
+                    ):
+                        if constraint_column == model_str(om_column.name):
+                            column_index.isForeignKey = True
+                            # update the service name of OM with the alation datasource id in the column FQN
+                            splitted_col_fqn = fqn.split(
+                                model_str(table_constraint.referredColumns[i])
+                            )
+                            splitted_col_fqn[0] = str(alation_datasource_id)
+                            column_index.referencedColumnId = (
+                                fqn._build(  # pylint: disable=protected-access
+                                    *splitted_col_fqn
+                                )
+                            )
+                            break
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Failed to update foreign key for {model_str(om_column.name)}: {exc}"
+            )
+
+    def _get_column_index(
+        self,
+        alation_datasource_id: int,
+        om_column: Column,
+        table_constraints: Optional[List[TableConstraint]],
+    ) -> Optional[ColumnIndex]:
         """
         Method to get the alation column index
         """
         column_index = ColumnIndex()
         try:
+            # Attach the primary key
             if om_column.constraint == Constraint.PRIMARY_KEY:
-                return ColumnIndex(isPrimaryKey=True)
+                column_index.isPrimaryKey = True
+
+            # Attach the foreign key
+            self._update_foreign_key(
+                alation_datasource_id, om_column, table_constraints, column_index
+            )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
@@ -199,6 +257,7 @@ class AlationsinkSource(Source):
         schema_name: str,
         table_name: str,
         om_column: Column,
+        table_constraints: Optional[List[TableConstraint]],
     ) -> Optional[CreateColumnRequest]:
         """
         Method to form the CreateColumnRequest object
@@ -211,13 +270,19 @@ class AlationsinkSource(Source):
                     table_name,
                     model_str(om_column.name),
                 ),
-                column_type=om_column.dataType.value.lower(),
-                title=model_str(om_column.name),
+                column_type=om_column.dataTypeDisplay.lower()
+                if om_column.dataTypeDisplay
+                else om_column.dataType.value.lower(),
+                title=om_column.displayName
+                if om_column.displayName
+                else model_str(om_column.name),
                 description=model_str(om_column.description),
                 position=str(om_column.ordinalPosition)
                 if om_column.ordinalPosition
                 else None,
-                index=self._get_column_index(om_column),
+                index=self._get_column_index(
+                    alation_datasource_id, om_column, table_constraints
+                ),
                 nullable=self._check_nullable_column(om_column),
             )
         except Exception as exc:
@@ -241,6 +306,7 @@ class AlationsinkSource(Source):
                     schema_name=schema_name,
                     table_name=model_str(om_table.name),
                     om_column=om_column,
+                    table_constraints=om_table.tableConstraints,
                 )
                 if create_column_request:
                     create_requests.root.append(create_column_request)
@@ -266,6 +332,7 @@ class AlationsinkSource(Source):
                     entity=Table,
                     skip_on_failure=True,
                     params={"database": model_str(om_schema.fullyQualifiedName)},
+                    fields=["tableConstraints, columns"],
                 )
             )
             create_requests = CreateTableRequestList(root=[])

metadata/ingestion/source/pipeline/dbtcloud/metadata.py

@@ -314,7 +314,7 @@ class DbtcloudSource(PipelineServiceSource):
                 timestamp=Timestamp(
                     datetime_to_ts(
                         datetime.strptime(
-                            pipeline_details.created_at, "%Y-%m-%dT%H:%M:%S.%f%z"
+                            pipeline_details.created_at, "%Y-%m-%d %H:%M:%S.%f%z"
                         )
                         if pipeline_details.created_at
                         else None