openmetadata-ingestion 1.4.6.0rc1__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

metadata/ingestion/source/database/oracle/query_parser.py

@@ -39,8 +39,8 @@ class OracleQueryParserSource(QueryParserSource, ABC):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: OracleConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: OracleConnection = config.serviceConnection.root.config
         if not isinstance(connection, OracleConnection):
             raise InvalidSourceException(
                 f"Expected OracleConnection, but got {connection}"

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

@@ -32,8 +32,8 @@ class PinotdbSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: PinotDBConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: PinotDBConnection = config.serviceConnection.root.config
         if not isinstance(connection, PinotDBConnection):
             raise InvalidSourceException(
                 f"Expected PinotdbConnection, but got {connection}"

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

@@ -19,6 +19,7 @@ from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
     PostgresScheme,
 )
+from metadata.generated.schema.type.basic import DateTime
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.source.connections import get_connection
@@ -83,7 +84,7 @@ class PostgresLineageSource(PostgresQueryParserSource, LineageSource):
                         yield TableQuery(
                             query=row["query_text"],
                             userName=row["usename"],
-                            analysisDate=datetime.now(),
+                            analysisDate=DateTime(datetime.now()),
                             aborted=self.get_aborted_status(row),
                             databaseName=self.get_database_name(row),
                             serviceName=self.config.serviceName,

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

@@ -36,6 +36,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.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -51,6 +52,7 @@ from metadata.ingestion.source.database.postgres.queries import (
     POSTGRES_GET_DB_NAMES,
     POSTGRES_GET_TABLE_NAMES,
     POSTGRES_PARTITION_DETAILS,
+    POSTGRES_SCHEMA_COMMENTS,
 )
 from metadata.ingestion.source.database.postgres.utils import (
     get_column_info,
@@ -69,6 +71,7 @@ from metadata.utils.sqlalchemy_utils import (
     get_all_table_ddls,
     get_all_table_owners,
     get_all_view_definitions,
+    get_schema_descriptions,
     get_table_ddl,
 )
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
@@ -139,18 +142,33 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
     Database metadata from Postgres Source
     """
 
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        self.schema_desc_map = {}
+
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: PostgresConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: PostgresConnection = config.serviceConnection.root.config
         if not isinstance(connection, PostgresConnection):
             raise InvalidSourceException(
                 f"Expected PostgresConnection, but got {connection}"
             )
         return cls(config, metadata)
 
+    def get_schema_description(self, schema_name: str) -> Optional[str]:
+        """
+        Method to fetch the schema description
+        """
+        return self.schema_desc_map.get(schema_name)
+
+    def set_schema_description_map(self) -> None:
+        self.schema_desc_map = get_schema_descriptions(
+            self.engine, POSTGRES_SCHEMA_COMMENTS
+        )
+
     def query_table_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
@@ -179,9 +197,10 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         yield from self._execute_database_query(POSTGRES_GET_DB_NAMES)
 
     def get_database_names(self) -> Iterable[str]:
-        if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
-            configured_db = self.config.serviceConnection.__root__.config.database
+        if not self.config.serviceConnection.root.config.ingestAllDatabases:
+            configured_db = self.config.serviceConnection.root.config.database
             self.set_inspector(database_name=configured_db)
+            self.set_schema_description_map()
             yield configured_db
         else:
             for new_database in self.get_database_names_raw():
@@ -205,6 +224,7 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
 
                 try:
                     self.set_inspector(database_name=new_database)
+                    self.set_schema_description_map()
                     yield new_database
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
@@ -253,8 +273,10 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
                 row = list(res)
                 fqn_elements = [name for name in row[2:] if name]
                 yield from get_ometa_tag_and_classification(
-                    tag_fqn=fqn._build(  # pylint: disable=protected-access
-                        self.context.get().database_service, *fqn_elements
+                    tag_fqn=FullyQualifiedEntityName(
+                        fqn._build(  # pylint: disable=protected-access
+                            self.context.get().database_service, *fqn_elements
+                        )
                     ),
                     tags=[row[1]],
                     classification_name=self.service_connection.classificationName,

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

@@ -62,8 +62,8 @@ def _get_column_lineages(source_entity, target_entity):
             if source_column.name == target_column.name:
                 column_lineages.append(
                     ColumnLineage(
-                        fromColumns=[source_column.fullyQualifiedName.__root__],
-                        toColumn=target_column.fullyQualifiedName.__root__,
+                        fromColumns=[source_column.fullyQualifiedName.root],
+                        toColumn=target_column.fullyQualifiedName.root,
                     )
                 )
                 break

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

@@ -85,6 +85,13 @@ JOIN (SELECT pc.oid as object_id, pc.relname, pp.*
 WHERE it.table_schema='{schema_name}' AND it.table_catalog='{database_name}';
 """
 
+POSTGRES_SCHEMA_COMMENTS = """
+    SELECT n.nspname AS schema_name, 
+            d.description AS comment
+    FROM pg_catalog.pg_namespace n
+    LEFT JOIN pg_catalog.pg_description d ON d.objoid = n.oid AND d.objsubid = 0;
+"""
+
 POSTGRES_TABLE_COMMENTS = """
     SELECT n.nspname as schema,
             c.relname as table_name,

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

@@ -56,8 +56,8 @@ class PostgresQueryParserSource(QueryParserSource, ABC):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: PostgresConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: PostgresConnection = config.serviceConnection.root.config
         if not isinstance(connection, PostgresConnection):
             raise InvalidSourceException(
                 f"Expected PostgresConnection, but got {connection}"
@@ -81,7 +81,7 @@ class PostgresQueryParserSource(QueryParserSource, ABC):
             if self.config.sourceConfig.config.queryLogFilePath:
                 yield from super().yield_table_queries_from_logs()
             else:
-                database = self.config.serviceConnection.__root__.config.database
+                database = self.config.serviceConnection.root.config.database
                 if database:
                     self.engine: Engine = get_connection(self.service_connection)
                     yield from self.process_table_query()
@@ -90,7 +90,7 @@ class PostgresQueryParserSource(QueryParserSource, ABC):
                     for res in results:
                         row = list(res)
                         logger.info(f"Ingesting from database: {row[0]}")
-                        self.config.serviceConnection.__root__.config.database = row[0]
+                        self.config.serviceConnection.root.config.database = row[0]
                         self.engine = get_connection(self.service_connection)
                         yield from self.process_table_query()
 

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

@@ -15,6 +15,7 @@ import traceback
 from datetime import datetime
 from typing import Iterable
 
+from metadata.generated.schema.type.basic import DateTime
 from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.postgres.queries import POSTGRES_SQL_STATEMENT
@@ -39,9 +40,11 @@ class PostgresUsageSource(PostgresQueryParserSource, UsageSource):
         """
         Process Query
         """
+        query = None
         try:
+            query = self.get_sql_statement()
             with get_connection(self.service_connection).connect() as conn:
-                rows = conn.execute(self.get_sql_statement())
+                rows = conn.execute(query)
                 queries = []
                 for row in rows:
                     row = dict(row)
@@ -50,7 +53,7 @@ class PostgresUsageSource(PostgresQueryParserSource, UsageSource):
                             TableQuery(
                                 query=row["query_text"],
                                 userName=row["usename"],
-                                analysisDate=datetime.now(),
+                                analysisDate=DateTime(datetime.now()),
                                 aborted=self.get_aborted_status(row),
                                 databaseName=self.get_database_name(row),
                                 serviceName=self.config.serviceName,
@@ -64,5 +67,9 @@ class PostgresUsageSource(PostgresQueryParserSource, UsageSource):
             if queries:
                 yield TableQueries(queries=queries)
         except Exception as err:
+            if query:
+                logger.debug(
+                    f"###### USAGE QUERY #######\n{query}\n##########################"
+                )
             logger.error(f"Source usage processing error - {err}")
             logger.debug(traceback.format_exc())

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

@@ -63,12 +63,12 @@ def get_connection(connection: PrestoConnection) -> Engine:
         connection.connectionArguments or init_empty_connection_arguments()
     )
     if connection.protocol:
-        connection.connectionArguments.__root__["protocol"] = connection.protocol
+        connection.connectionArguments.root["protocol"] = connection.protocol
     if connection.verify:
         connection.connectionArguments = (
             connection.connectionArguments or init_empty_connection_arguments()
         )
-        connection.connectionArguments.__root__["requests_kwargs"] = {
+        connection.connectionArguments.root["requests_kwargs"] = {
             "verify": connection.verify
         }
 

metadata/ingestion/source/database/presto/metadata.py

@@ -121,8 +121,8 @@ class PrestoSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config = WorkflowSource.parse_obj(config_dict)
-        connection: PrestoConnection = config.serviceConnection.__root__.config
+        config = WorkflowSource.model_validate(config_dict)
+        connection: PrestoConnection = config.serviceConnection.root.config
         if not isinstance(connection, PrestoConnection):
             raise InvalidSourceException(
                 f"Expected PrestoConnection, but got {connection}"

metadata/ingestion/source/database/query/lineage.py

@@ -25,7 +25,7 @@ class QueryLogLineageSource(LineageSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
         return cls(config, metadata)
 
     def prepare(self):

metadata/ingestion/source/database/query/usage.py

@@ -11,7 +11,7 @@
 """
 Common Query Log Connector
 """
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Optional
 
 from metadata.generated.schema.metadataIngestion.workflow import (
@@ -24,13 +24,15 @@ from metadata.ingestion.source.database.usage_source import UsageSource
 class QueryLogUsageSource(UsageSource):
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         super().__init__(config, metadata)
-        self.analysis_date = datetime.today().strftime("%Y-%m-%d %H:%M:%S")
+        self.analysis_date = (
+            datetime.now(timezone.utc).date().strftime("%Y-%m-%d %H:%M:%S")
+        )
 
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
         return cls(config, metadata)
 
     def prepare(self):

metadata/ingestion/source/database/query_parser_source.py

@@ -53,7 +53,7 @@ class QueryParserSource(Source, ABC):
         super().__init__()
         self.config = config
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
         self.source_config = self.config.sourceConfig.config
         self.start, self.end = get_start_and_end(self.source_config.queryLogDuration)
         self.engine = (

metadata/ingestion/source/database/redshift/metadata.py

@@ -163,8 +163,8 @@ class RedshiftSource(
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: RedshiftConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: RedshiftConnection = config.serviceConnection.root.config
         if not isinstance(connection, RedshiftConnection):
             raise InvalidSourceException(
                 f"Expected RedshiftConnection, but got {connection}"
@@ -282,14 +282,14 @@ class RedshiftSource(
             )
 
     def get_database_names(self) -> Iterable[str]:
-        if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
+        if not self.config.serviceConnection.root.config.ingestAllDatabases:
             self.get_partition_details()
 
             self._set_incremental_table_processor(
-                self.config.serviceConnection.__root__.config.database
+                self.config.serviceConnection.root.config.database
             )
 
-            yield self.config.serviceConnection.__root__.config.database
+            yield self.config.serviceConnection.root.config.database
         else:
             for new_database in self.get_database_names_raw():
                 database_fqn = fqn.build(
@@ -385,7 +385,7 @@ class RedshiftSource(
                 )
             ).all()
             for row in results:
-                stored_procedure = RedshiftStoredProcedure.parse_obj(dict(row))
+                stored_procedure = RedshiftStoredProcedure.model_validate(dict(row))
                 yield stored_procedure
 
     @calculate_execution_time_generator()
@@ -396,7 +396,7 @@ class RedshiftSource(
 
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
-                name=EntityName(__root__=stored_procedure.name),
+                name=EntityName(stored_procedure.name),
                 storedProcedureCode=StoredProcedureCode(
                     language=Language.SQL,
                     code=stored_procedure.definition,

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

@@ -24,7 +24,7 @@ class RedshiftStoredProcedure(BaseModel):
     """Redshift stored procedure list query results"""
 
     name: str
-    owner: Optional[str]
+    owner: Optional[str] = None
     definition: str
 
 

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

@@ -237,6 +237,30 @@ REDSHIFT_TEST_PARTITION_DETAILS = "select * from SVV_TABLE_INFO limit 1"
 # hence we are appending "create view <schema>.<table> as " to select query
 # to generate the column level lineage
 REDSHIFT_GET_ALL_RELATIONS = """
+  WITH view_defs AS (
+      SELECT
+          c.oid,
+          pg_catalog.pg_get_viewdef(c.oid, true) AS view_definition,
+          n.nspname,
+          c.relname
+      FROM
+          pg_catalog.pg_class c
+          LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+      WHERE
+          c.relkind = 'v'
+    ),
+    adjusted_view_defs AS (
+        SELECT
+            oid,
+            CASE
+                WHEN view_definition LIKE '%WITH NO SCHEMA BINDING%' THEN
+                  REGEXP_REPLACE(view_definition, 'create view [^ ]+ as (.*WITH NO SCHEMA BINDING;?)', '\\1')
+                ELSE
+                  'CREATE VIEW ' || nspname || '.' || relname || ' AS ' || view_definition
+            END AS view_definition
+        FROM
+            view_defs
+    )
     SELECT
         c.relkind,
         n.oid as "schema_oid",
@@ -248,13 +272,13 @@ REDSHIFT_GET_ALL_RELATIONS = """
         AS "diststyle",
         c.relowner AS "owner_id",
         u.usename AS "owner_name",
-        TRIM(TRAILING ';' FROM
-        'create view ' || n.nspname || '.' || c.relname || ' as ' ||pg_catalog.pg_get_viewdef(c.oid, true))
+        avd.view_definition 
         AS "view_definition",
         pg_catalog.array_to_string(c.relacl, '\n') AS "privileges"
     FROM pg_catalog.pg_class c
             LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
             JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner
+            LEFT JOIN adjusted_view_defs avd ON avd.oid = c.oid
     WHERE c.relkind IN ('r', 'v', 'm', 'S', 'f')
         AND n.nspname !~ '^pg_' {schema_clause} {table_clause}
     UNION

metadata/ingestion/source/database/redshift/query_parser.py

@@ -41,8 +41,8 @@ class RedshiftQueryParserSource(QueryParserSource, ABC):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: RedshiftConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: RedshiftConnection = config.serviceConnection.root.config
         if not isinstance(connection, RedshiftConnection):
             raise InvalidSourceException(
                 f"Expected RedshiftConnection, but got {connection}"

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

@@ -37,9 +37,7 @@ def get_connection(connection: SalesforceConnection) -> Engine:
         security_token=connection.securityToken.get_secret_value(),
         domain=connection.salesforceDomain,
         version=connection.salesforceApiVersion,
-        **connection.connectionArguments.__root__
-        if connection.connectionArguments
-        else {},
+        **connection.connectionArguments.root if connection.connectionArguments else {},
     )
 
 

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

@@ -45,6 +45,7 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+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
@@ -76,7 +77,7 @@ class SalesforceSource(DatabaseServiceSource):
             self.config.sourceConfig.config
         )
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
         self.ssl_manager: SSLManager = check_ssl_and_init(self.service_connection)
         if self.ssl_manager:
             self.service_connection = self.ssl_manager.setup_ssl(
@@ -90,8 +91,8 @@ class SalesforceSource(DatabaseServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SalesforceConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SalesforceConnection = config.serviceConnection.root.config
         if not isinstance(connection, SalesforceConnection):
             raise InvalidSourceException(
                 f"Expected SalesforceConnection, but got {connection}"
@@ -139,12 +140,14 @@ class SalesforceSource(DatabaseServiceSource):
         """
         yield Either(
             right=CreateDatabaseSchemaRequest(
-                name=schema_name,
-                database=fqn.build(
-                    metadata=self.metadata,
-                    entity_type=Database,
-                    service_name=self.context.get().database_service,
-                    database_name=self.context.get().database,
+                name=EntityName(schema_name),
+                database=FullyQualifiedEntityName(
+                    fqn.build(
+                        metadata=self.metadata,
+                        entity_type=Database,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
+                    )
                 ),
             )
         )
@@ -225,7 +228,7 @@ class SalesforceSource(DatabaseServiceSource):
         return None
 
     def yield_table(
-        self, table_name_and_type: Tuple[str, str]
+        self, table_name_and_type: Tuple[str, TableType]
     ) -> Iterable[Either[CreateTableRequest]]:
         """
         From topology.
@@ -240,17 +243,19 @@ class SalesforceSource(DatabaseServiceSource):
             )
             columns = self.get_columns(salesforce_objects["fields"])
             table_request = CreateTableRequest(
-                name=table_name,
+                name=EntityName(table_name),
                 tableType=table_type,
                 description=self.get_table_description(table_name),
                 columns=columns,
                 tableConstraints=table_constraints,
-                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,
+                databaseSchema=FullyQualifiedEntityName(
+                    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,
+                    )
                 ),
                 sourceUrl=self.get_source_url(
                     table_name=table_name,