openmetadata-ingestion 1.8.9.0__py3-none-any.whl → 1.9.0.0__py3-none-any.whl

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

@@ -10,12 +10,16 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.exasolConnection import (
     ExasolConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
 )
-from metadata.ingestion.connections.test_connections import test_query
+from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -79,9 +83,16 @@ def test_connection(
     engine: Engine,
     service_connection: ExasolConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    test_query(engine, "SELECT 1;")
+    return test_connection_db_common(
+        metadata=metadata,
+        engine=engine,
+        service_connection=service_connection,
+        automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
+    )

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

@@ -0,0 +1,13 @@
+EXASOL_TEST_GET_QUERIES = """
+    SELECT
+      s.sql_text,
+      s.command_name,
+      se.user_name,
+      s.start_time,
+      s.stop_time,
+      s.duration
+    FROM EXA_DBA_AUDIT_SQL s
+    JOIN EXA_DBA_AUDIT_SESSIONS se
+    ON s.SESSION_ID = se.SESSION_ID
+    LIMIT 1
+    """

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

@@ -16,7 +16,6 @@ from typing import Optional
 
 from sqlalchemy.engine import Engine
 
-from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
@@ -27,7 +26,7 @@ from metadata.generated.schema.entity.services.connections.database.common.basic
     BasicAuth,
 )
 from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
-    MysqlConnection,
+    MysqlConnection as MySQLConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
@@ -47,36 +46,31 @@ from metadata.ingestion.source.database.mysql.queries import (
     MYSQL_TEST_GET_QUERIES_SLOW_LOGS,
 )
 from metadata.utils.constants import THREE_MIN
+from metadata.utils.credentials import get_azure_access_token
 
 
-class MySQLConnection(BaseConnection[MysqlConnection, Engine]):
-    def get_client(self) -> Engine:
+class MySQLConnection(BaseConnection[MySQLConnectionConfig, Engine]):
+    def _get_client(self) -> Engine:
         """
         Return the SQLAlchemy Engine for MySQL.
         """
         connection = self.service_connection
 
         if isinstance(connection.authType, AzureConfigurationSource):
-            if not connection.authType.azureConfig:
-                raise ValueError("Azure Config is missing")
-            azure_client = AzureClient(connection.authType.azureConfig).create_client()
-            if not connection.authType.azureConfig.scopes:
-                raise ValueError(
-                    "Azure Scopes are missing, please refer https://learn.microsoft.com/"
-                    "en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-micr"
-                    "osoft-entra-access-token and fetch the resource associated with it,"
-                    " for e.g. https://ossrdbms-aad.database.windows.net/.default"
-                )
-            access_token_obj = azure_client.get_token(
-                *connection.authType.azureConfig.scopes.split(",")
-            )
-            connection.authType = BasicAuth(password=access_token_obj.token)  # type: ignore
+            access_token = get_azure_access_token(connection.authType)
+            connection.authType = BasicAuth(password=access_token)  # type: ignore
         return create_generic_db_connection(
             connection=connection,
             get_connection_url_fn=get_connection_url_common,
             get_connection_args_fn=get_connection_args_common,
         )
 
+    def get_connection_dict(self) -> dict:
+        """
+        Return the connection dictionary for this service.
+        """
+        raise NotImplementedError("get_connection_dict is not implemented for MySQL")
+
     def test_connection(
         self,
         metadata: OpenMetadata,
@@ -94,7 +88,7 @@ class MySQLConnection(BaseConnection[MysqlConnection, Engine]):
         }
         return test_connection_db_schema_sources(
             metadata=metadata,
-            engine=self.get_client(),
+            engine=self.client,
             service_connection=self.service_connection,
             automation_workflow=automation_workflow,
             timeout_seconds=timeout_seconds,

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

@@ -56,6 +56,7 @@ from metadata.ingestion.source.database.oracle.queries import (
 from metadata.ingestion.source.database.oracle.utils import (
     _get_col_type,
     _get_constraint_data,
+    get_all_view_definitions,
     get_columns,
     get_mview_names,
     get_mview_names_dialect,
@@ -70,7 +71,6 @@ from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
     get_all_table_comments,
     get_all_table_ddls,
-    get_all_view_definitions,
     get_table_ddl,
 )
 

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

@@ -73,6 +73,25 @@ def get_view_definition(
     )
 
 
+@reflection.cache
+def get_all_view_definitions(self, connection, query):
+    """
+    Method to fetch view definition of all available views
+    """
+    self.all_view_definitions = {}
+    self.current_db: str = connection.engine.url.database  # type: ignore
+    result = connection.execute(query)
+    for view in result:
+        if hasattr(view, "view_def") and hasattr(view, "schema"):
+            self.all_view_definitions[
+                (view.view_name, view.schema)
+            ] = f"CREATE OR REPLACE VIEW {view.view_name} AS {view.view_def}"
+        elif hasattr(view, "VIEW_DEF") and hasattr(view, "SCHEMA"):
+            self.all_view_definitions[
+                (view.VIEW_NAME, view.SCHEMA)
+            ] = f"CREATE OR REPLACE VIEW {view.VIEW_NAME} AS {view.VIEW_DEF}"
+
+
 def _get_col_type(
     self, coltype, precision, scale, length, colname
 ):  # pylint: disable=too-many-branches

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

@@ -12,20 +12,21 @@
 """
 Source connection handler
 """
-
 from typing import Optional
 
 from sqlalchemy.engine import Engine
 
-from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
 from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
     BasicAuth,
 )
 from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
-    PostgresConnection,
+    PostgresConnection as PostgresConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
@@ -35,6 +36,7 @@ from metadata.ingestion.connections.builders import (
     get_connection_args_common,
     get_connection_url_common,
 )
+from metadata.ingestion.connections.connection import BaseConnection
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.postgres.queries import (
@@ -46,54 +48,55 @@ from metadata.ingestion.source.database.postgres.utils import (
     get_postgres_time_column_name,
 )
 from metadata.utils.constants import THREE_MIN
+from metadata.utils.credentials import get_azure_access_token
 
 
-def get_connection(connection: PostgresConnection) -> Engine:
-    """
-    Create connection
-    """
+class PostgresConnection(BaseConnection[PostgresConnectionConfig, Engine]):
+    def _get_client(self) -> Engine:
+        """
+        Return the SQLAlchemy Engine for PostgreSQL.
+        """
+        connection = self.service_connection
 
-    if hasattr(connection.authType, "azureConfig"):
-        azure_client = AzureClient(connection.authType.azureConfig).create_client()
-        if not connection.authType.azureConfig.scopes:
-            raise ValueError(
-                "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/postgresql/flexible-server/how-to-configure-sign-in-azure-ad-authentication#retrieve-the-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
-            )
-        access_token_obj = azure_client.get_token(
-            *connection.authType.azureConfig.scopes.split(",")
+        if isinstance(connection.authType, AzureConfigurationSource):
+            access_token = get_azure_access_token(connection.authType)
+            connection.authType = BasicAuth(password=access_token)  # type: ignore
+        return create_generic_db_connection(
+            connection=connection,
+            get_connection_url_fn=get_connection_url_common,
+            get_connection_args_fn=get_connection_args_common,
         )
-        connection.authType = BasicAuth(password=access_token_obj.token)
-    return create_generic_db_connection(
-        connection=connection,
-        get_connection_url_fn=get_connection_url_common,
-        get_connection_args_fn=get_connection_args_common,
-    )
-
 
-def test_connection(
-    metadata: OpenMetadata,
-    engine: Engine,
-    service_connection: PostgresConnection,
-    automation_workflow: Optional[AutomationWorkflow] = None,
-    timeout_seconds: Optional[int] = THREE_MIN,
-) -> TestConnectionResult:
-    """
-    Test connection. This can be executed either as part
-    of a metadata workflow or during an Automation Workflow
-    """
+    def get_connection_dict(self) -> dict:
+        """
+        Return the connection dictionary for this service.
+        """
+        raise NotImplementedError(
+            "get_connection_dict is not implemented for PostgreSQL"
+        )
 
-    queries = {
-        "GetQueries": POSTGRES_TEST_GET_QUERIES.format(
-            time_column_name=get_postgres_time_column_name(engine=engine),
-        ),
-        "GetDatabases": POSTGRES_GET_DATABASE,
-        "GetTags": POSTGRES_TEST_GET_TAGS,
-    }
-    return test_connection_db_common(
-        metadata=metadata,
-        engine=engine,
-        service_connection=service_connection,
-        automation_workflow=automation_workflow,
-        queries=queries,
-        timeout_seconds=timeout_seconds,
-    )
+    def test_connection(
+        self,
+        metadata: OpenMetadata,
+        automation_workflow: Optional[AutomationWorkflow] = None,
+        timeout_seconds: Optional[int] = THREE_MIN,
+    ) -> TestConnectionResult:
+        """
+        Test connection. This can be executed either as part
+        of a metadata workflow or during an Automation Workflow
+        """
+        queries = {
+            "GetQueries": POSTGRES_TEST_GET_QUERIES.format(
+                time_column_name=get_postgres_time_column_name(engine=self.client),
+            ),
+            "GetDatabases": POSTGRES_GET_DATABASE,
+            "GetTags": POSTGRES_TEST_GET_TAGS,
+        }
+        return test_connection_db_common(
+            metadata=metadata,
+            engine=self.client,
+            service_connection=self.service_connection,
+            automation_workflow=automation_workflow,
+            timeout_seconds=timeout_seconds,
+            queries=queries,
+        )

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

@@ -1,3 +1,4 @@
+from metadata.ingestion.source.database.postgres.connection import PostgresConnection
 from metadata.ingestion.source.database.postgres.lineage import PostgresLineageSource
 from metadata.ingestion.source.database.postgres.metadata import PostgresSource
 from metadata.ingestion.source.database.postgres.usage import PostgresUsageSource
@@ -7,4 +8,5 @@ ServiceSpec = DefaultDatabaseSpec(
     metadata_source_class=PostgresSource,
     lineage_source_class=PostgresLineageSource,
     usage_source_class=PostgresUsageSource,
+    connection_class=PostgresConnection,
 )