openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

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

@@ -119,8 +119,8 @@ class OracleSource(StoredProcedureMixin, CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config = WorkflowSource.parse_obj(config_dict)
-        connection: OracleConnection = config.serviceConnection.__root__.config
+        config = WorkflowSource.model_validate(config_dict)
+        connection: OracleConnection = config.serviceConnection.root.config
         if not isinstance(connection, OracleConnection):
             raise InvalidSourceException(
                 f"Expected OracleConnection, but got {connection}"
@@ -161,20 +161,23 @@ class OracleSource(StoredProcedureMixin, CommonDbSourceService):
                 schema_definition = inspector.get_table_ddl(
                     self.connection, table_name, schema_name
                 )
+                return (
+                    str(schema_definition).strip()
+                    if schema_definition is not None
+                    else None
+                )
 
-            else:
-                definition_fn = inspector.get_view_definition
-                if table_type == TableType.MaterializedView:
-                    definition_fn = inspector.get_mview_definition
+            definition_fn = inspector.get_view_definition
+            if table_type == TableType.MaterializedView:
+                definition_fn = inspector.get_mview_definition
 
-                schema_definition = definition_fn(table_name, schema_name)
+            schema_definition = definition_fn(table_name, schema_name)
 
-            schema_definition = (
-                str(schema_definition.strip())
+            return (
+                str(schema_definition).strip()
                 if schema_definition is not None
                 else None
             )
-            return schema_definition
 
         except NotImplementedError:
             logger.warning("Schema definition not implemented")
@@ -221,13 +224,13 @@ class OracleSource(StoredProcedureMixin, CommonDbSourceService):
 
         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,
                 ),
-                owner=self.metadata.get_reference_by_name(
-                    name=stored_procedure.owner.lower()
+                owners=self.metadata.get_reference_by_name(
+                    name=stored_procedure.owner.lower(), is_owner=True
                 ),
                 databaseSchema=fqn.build(
                     metadata=self.metadata,

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

@@ -20,7 +20,7 @@ class OracleStoredProcedure(BaseModel):
 class FetchProcedure(BaseModel):
     """Oracle Fetch Stored Procedure Raw Model"""
 
-    owner: Optional[str]
+    owner: Optional[str] = None
     name: str
     line: int
     text: str

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/converter_orm.py

@@ -0,0 +1,52 @@
+#  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.
+
+"""
+Converter logic to transform an OpenMetadata Table Entity for Redshift
+to an SQLAlchemy ORM class.
+"""
+
+from typing import Dict, Set
+
+from sqlalchemy.sql.sqltypes import TypeEngine
+
+from metadata.generated.schema.entity.data.table import Column, DataType
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseServiceType,
+)
+from metadata.ingestion.source.database.postgres.types.money import PostgresMoney
+from metadata.profiler.orm.converter.common import CommonMapTypes
+from metadata.profiler.orm.converter.converter_registry import converter_registry
+
+_CUSTOM_TYPE_MAP = {DataType.MONEY: PostgresMoney}
+_CUSTOM_TYPE_MAP_reversed = {v: {k} for k, v in _CUSTOM_TYPE_MAP.items()}
+
+
+class PostgrestMapTypes(CommonMapTypes):
+    """Postgres type mapper, inherits from CommonMapTypes"""
+
+    def return_custom_type(self, col: Column, table_service_type):
+        return _CUSTOM_TYPE_MAP.get(
+            col.dataType, super().return_custom_type(col, table_service_type)
+        )
+
+    @staticmethod
+    def map_sqa_to_om_types() -> Dict[TypeEngine, Set[DataType]]:
+        """returns an ORM type"""
+        # pylint: disable=import-outside-toplevel
+
+        return {
+            **CommonMapTypes.map_sqa_to_om_types(),
+            **_CUSTOM_TYPE_MAP_reversed,
+        }
+
+
+converter_registry[DatabaseServiceType.Postgres] = PostgrestMapTypes

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,
@@ -63,16 +65,23 @@ from metadata.ingestion.source.database.postgres.utils import (
 )
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database
+from metadata.utils.importer import import_side_effects
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
     get_all_table_comments,
     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
 
+import_side_effects(
+    "metadata.ingestion.source.database.postgres.converter_orm",
+    "metadata.ingestion.source.database.postgres.metrics",
+)
+
 TableKey = namedtuple("TableKey", ["schema", "table_name"])
 
 logger = ingestion_logger()
@@ -139,18 +148,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 +203,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 +230,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 +279,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/metrics.py

@@ -0,0 +1,42 @@
+from sqlalchemy.ext.compiler import compiles
+
+from metadata.ingestion.source.database.postgres.types.money import PostgresMoney
+from metadata.profiler.metrics.static.mean import AvgFn
+from metadata.profiler.metrics.static.stddev import StdDevFn
+from metadata.profiler.orm.functions.median import MedianFn
+from metadata.profiler.orm.registry import Dialects
+
+
+@compiles(AvgFn, Dialects.Postgres)
+def avg(element, compiler, **kw):
+    """
+    Cast to decimal to get around potential integer overflow error
+    """
+    proc = compiler.process(element.clauses, **kw)
+    if isinstance(list(element.clauses)[0].type, PostgresMoney):
+        return f"{element.name}({PostgresMoney.compile_as_float(proc)})"
+    return f"{element.name}({proc})"
+
+
+@compiles(StdDevFn, Dialects.Postgres)
+def stddev(element, compiler, **kw):
+    """Returns stdv for clickhouse database and handle empty tables.
+    If table is empty, clickhouse returns NaN.
+    """
+    proc = compiler.process(element.clauses, **kw)
+    if isinstance(list(element.clauses)[0].type, PostgresMoney):
+        return f"STDDEV_POP({PostgresMoney.compile_as_float(proc)})"
+    return f"STDDEV_POP({proc})"
+
+
+@compiles(MedianFn, Dialects.Postgres)
+def median(elements, compiler, **kwargs):  # pylint: disable=unused-argument
+    col, _, percentile = [
+        compiler.process(element, **kwargs) for element in elements.clauses
+    ]
+    if isinstance(list(elements.clauses)[0], PostgresMoney):
+        return "percentile_cont(%.2f) WITHIN GROUP (ORDER BY %s ASC)" % (
+            percentile,
+            PostgresMoney.compile_as_float(col),
+        )
+    return MedianFn.default_fn(elements, compiler, **kwargs)

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/types/money.py

@@ -0,0 +1,37 @@
+#  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.
+
+"""
+Custom sqlalchemy type for Postgres MONEY type
+"""
+
+
+from sqlalchemy.dialects.postgresql import MONEY
+from sqlalchemy.sql.sqltypes import String
+
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class PostgresMoney(MONEY):
+    """Wrapper for Postgres MONEY type"""
+
+    impl = String
+    cache_ok = True
+
+    @property
+    def quantifiable(self):
+        return True
+
+    @staticmethod
+    def compile_as_float(expr: str) -> str:
+        return f"(({expr})::numeric::float8)"

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,15 @@ 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())
+
+    def get_filters(self) -> str:
+        if filter_condition := self.source_config.filterCondition:
+            filter_condition = filter_condition.replace("%", "%%")
+            return f"{self.filters} AND s.{filter_condition}"
+        return self.filters

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/connection.py

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from functools import partial
 from typing import Optional
 
 from sqlalchemy.engine import Engine
@@ -27,8 +28,15 @@ from metadata.ingestion.connections.builders import (
     get_connection_args_common,
     get_connection_url_common,
 )
-from metadata.ingestion.connections.test_connections import test_connection_db_common
+from metadata.ingestion.connections.test_connections import (
+    SourceConnectionException,
+    execute_inspector_func,
+    test_connection_engine_step,
+    test_connection_steps,
+    test_query,
+)
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.connections import kill_active_connections
 from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_DATABASE_NAMES,
     REDSHIFT_TEST_GET_QUERIES,
@@ -57,15 +65,35 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    queries = {
-        "GetQueries": REDSHIFT_TEST_GET_QUERIES,
-        "GetDatabases": REDSHIFT_GET_DATABASE_NAMES,
-        "GetPartitionTableDetails": REDSHIFT_TEST_PARTITION_DETAILS,
+
+    def test_get_queries_permissions(engine_: Engine):
+        """Check if we have the right permissions to list queries"""
+        with engine_.connect() as conn:
+            res = conn.execute(REDSHIFT_TEST_GET_QUERIES).fetchone()
+            if not all(res):
+                raise SourceConnectionException(
+                    f"We don't have the right permissions to list queries - {res}"
+                )
+
+    test_fn = {
+        "CheckAccess": partial(test_connection_engine_step, engine),
+        "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
+        "GetTables": partial(execute_inspector_func, engine, "get_table_names"),
+        "GetViews": partial(execute_inspector_func, engine, "get_view_names"),
+        "GetQueries": partial(test_get_queries_permissions, engine),
+        "GetDatabases": partial(
+            test_query, statement=REDSHIFT_GET_DATABASE_NAMES, engine=engine
+        ),
+        "GetPartitionTableDetails": partial(
+            test_query, statement=REDSHIFT_TEST_PARTITION_DETAILS, engine=engine
+        ),
     }
-    test_connection_db_common(
+
+    test_connection_steps(
         metadata=metadata,
-        engine=engine,
-        service_connection=service_connection,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
-        queries=queries,
     )
+
+    kill_active_connections(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