openmetadata-ingestion 1.5.7.1.dev2__py3-none-any.whl → 1.5.8.0.dev0__py3-none-any.whl

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

@@ -27,6 +27,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
     AthenaConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
@@ -37,6 +40,7 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection_url(connection: AthenaConnection) -> str:
@@ -94,7 +98,8 @@ def test_connection(
     engine: Engine,
     service_connection: AthenaConnection,
     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
@@ -121,9 +126,10 @@ def test_connection(
         "GetViews": custom_executor_for_view,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -28,6 +28,9 @@ from metadata.generated.schema.entity.services.connections.database.azureSQLConn
 from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
     MssqlConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
@@ -35,6 +38,7 @@ from metadata.ingestion.connections.builders import (
 )
 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
 
 
 def get_connection_url(connection: Union[AzureSQLConnection, MssqlConnection]) -> str:
@@ -105,14 +109,16 @@ def test_connection(
     engine: Engine,
     service_connection: AzureSQLConnection,
     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_connection_db_common(
+    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/bigquery/connection.py

@@ -26,6 +26,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.generated.schema.security.credentials.gcpCredentials import (
     GcpCredentialsPath,
 )
@@ -46,6 +49,7 @@ from metadata.ingestion.connections.test_connections import (
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.bigquery.queries import BIGQUERY_TEST_STATEMENT
+from metadata.utils.constants import THREE_MIN
 from metadata.utils.credentials import set_google_credentials
 from metadata.utils.logger import ingestion_logger
 
@@ -108,7 +112,8 @@ def test_connection(
     engine: Engine,
     service_connection: BigQueryConnection,
     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
@@ -162,11 +167,12 @@ def test_connection(
             ),
         }
 
-        test_connection_steps(
+        return test_connection_steps(
             metadata=metadata,
             test_fn=test_fn,
             service_type=service_connection.type.value,
             automation_workflow=automation_workflow,
+            timeout_seconds=timeout_seconds,
         )
 
-    test_connection_inner(engine)
+    return test_connection_inner(engine)

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

@@ -19,6 +19,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.bigTableConnection import (
     BigTableConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.generated.schema.security.credentials.gcpValues import (
     GcpCredentialsValues,
     SingleProjectId,
@@ -29,6 +32,7 @@ from metadata.ingestion.connections.test_connections import (
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.bigtable.client import MultiProjectClient
+from metadata.utils.constants import THREE_MIN
 from metadata.utils.credentials import set_google_credentials
 from metadata.utils.logger import ingestion_logger
 
@@ -95,7 +99,8 @@ def test_connection(
     client: MultiProjectClient,
     service_connection: BigTableConnection,
     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
@@ -108,9 +113,10 @@ def test_connection(
         "GetRows": tester.get_row,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -23,6 +23,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.clickhouseConnection import (
     ClickhouseConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
@@ -35,6 +38,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.clickhouse.queries import (
     CLICKHOUSE_SQL_STATEMENT_TEST,
 )
+from metadata.utils.constants import THREE_MIN
 
 HTTPS_PROTOCOL = "https"
 
@@ -67,7 +71,8 @@ def test_connection(
     engine: Engine,
     service_connection: ClickhouseConnection,
     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
@@ -75,10 +80,11 @@ def test_connection(
 
     queries = {"GetQueries": CLICKHOUSE_SQL_STATEMENT_TEST}
 
-    test_connection_db_common(
+    return test_connection_db_common(
         metadata=metadata,
         engine=engine,
         service_connection=service_connection,
         automation_workflow=automation_workflow,
         queries=queries,
+        timeout_seconds=timeout_seconds,
     )

metadata/ingestion/source/database/common_db_source.py

@@ -478,6 +478,12 @@ class CommonDbSourceService(
         """Not Implemented"""
         yield from []
 
+    def get_location_path(self, table_name: str, schema_name: str) -> Optional[str]:
+        """
+        Method to fetch the location path of the table
+        by default there will be no location path
+        """
+
     @calculate_execution_time_generator()
     def yield_table(
         self, table_name_and_type: Tuple[str, TableType]
@@ -553,6 +559,9 @@ class CommonDbSourceService(
                     table_type=table_type,
                 ),
                 owners=self.get_owner_ref(table_name=table_name),
+                locationPath=self.get_location_path(
+                    table_name=table_name, schema_name=schema_name
+                ),
             )
 
             is_partitioned, partition_details = self.get_table_partition_details(

metadata/ingestion/source/database/common_nosql_source.py

@@ -232,19 +232,9 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         schema_name = self.context.get().database_schema
         try:
             data = self.get_table_columns_dict(schema_name, table_name)
-            try:
-                df = pd.DataFrame.from_records(list(data))
-                column_parser = DataFrameColumnParser.create(df)
-                columns = column_parser.get_columns()
-            except Exception as exc:
-                logger.error(
-                    f"Failed to parse columns for table "
-                    f"[{self.context.get().database}.{schema_name}.{table_name}]: {exc}"
-                )
-                logger.error(traceback.format_exc())
-                logger.error(f"Type of data: {type(data)}")
-                logger.error(f"Data: {list(data)[:100]}")
-                columns = []
+            df = pd.DataFrame.from_records(list(data))
+            column_parser = DataFrameColumnParser.create(df)
+            columns = column_parser.get_columns()
             table_request = CreateTableRequest(
                 name=EntityName(table_name),
                 tableType=table_type,
@@ -277,10 +267,7 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
             yield Either(
                 left=StackTraceError(
                     name=table_name,
-                    error=(
-                        f"Unexpected exception to yield table"
-                        f"[{self.context.get().database}.{schema_name}.{table_name}]: {exc}"
-                    ),
+                    error=f"Unexpected exception to yield table [{table_name}]: {exc}",
                     stackTrace=traceback.format_exc(),
                 )
             )

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

@@ -23,8 +23,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.couchbaseConnection import (
     CouchbaseConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: CouchbaseConnection):
@@ -49,7 +53,8 @@ def test_connection(
     client: Any,
     service_connection: CouchbaseConnection,
     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
@@ -80,9 +85,10 @@ def test_connection(
         "GetCollections": partial(test_get_collections, client, holder),
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

metadata/ingestion/source/database/couchbase/metadata.py

@@ -11,13 +11,10 @@
 """
 Couchbase source methods.
 """
-import re
+
 import traceback
 from typing import Dict, Iterable, List, Optional
 
-import requests
-from requests.auth import HTTPBasicAuth
-
 from metadata.generated.schema.entity.services.connections.database.couchbaseConnection import (
     CouchbaseConnection,
 )
@@ -26,17 +23,15 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.common_nosql_source import CommonNoSQLSource
-from metadata.ingestion.source.database.couchbase.models import IndexObject as Index
-from metadata.ingestion.source.database.couchbase.queries import (
-    COUCHBASE_GET_INDEX_KEYS,
+from metadata.ingestion.source.database.common_nosql_source import (
+    SAMPLE_SIZE,
+    CommonNoSQLSource,
 )
+from metadata.ingestion.source.database.couchbase.queries import COUCHBASE_GET_DATA
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
-DEFAULT_SCHEMA_NAME = "_default"
-
 
 class CouchbaseSource(CommonNoSQLSource):
     """
@@ -44,12 +39,9 @@ class CouchbaseSource(CommonNoSQLSource):
     Database metadata from Dynamo Source
     """
 
-    service_connection: CouchbaseConnection
-
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         super().__init__(config, metadata)
         self.couchbase = self.connection_obj
-        self.index_condition_map = {}
 
     @classmethod
     def create(
@@ -66,12 +58,10 @@ class CouchbaseSource(CommonNoSQLSource):
     def get_database_names(self) -> Iterable[str]:
         try:
             if self.service_connection.bucket:
-                self.index_condition_map.clear()
                 yield self.service_connection.__dict__.get("bucket")
             else:
                 buckets = self.couchbase.buckets()
                 for bucket_name in buckets.get_all_buckets():
-                    self.index_condition_map.clear()
                     yield bucket_name.name
         except Exception as exp:
             logger.debug(f"Failed to fetch bucket name: {exp}")
@@ -111,127 +101,21 @@ class CouchbaseSource(CommonNoSQLSource):
             logger.debug(traceback.format_exc())
         return []
 
-    def _is_valid_key(self, key: str) -> bool:
-        return bool(re.fullmatch(r"[a-zA-Z0-9._`\%]+", key))
-
-    def get_index_condition(self, schema_name: str) -> str:
-        """
-        Method to prepare query condition based on index
-        """
-        bucket_name = self.context.get().database
-        if self.index_condition_map.get((bucket_name, schema_name)):
-            return self.index_condition_map.get((bucket_name, schema_name))
-
-        index_condition = set()
-
-        if schema_name == DEFAULT_SCHEMA_NAME:
-            condition = f"keyspace_id = '{bucket_name}' "
-        else:
-            condition = f"bucket_id = '{bucket_name}' AND scope_id = '{schema_name}'"
-
-        query = COUCHBASE_GET_INDEX_KEYS.format(condition=condition)
-        result = self.couchbase.query(query)
-        for row in result.rows():
-            index_obj = Index(**dict(row))
-            if index_obj.indexes:
-                if index_obj.indexes.is_primary:
-                    self.index_condition_map[(bucket_name, schema_name)] = ""
-                    return ""
-                for key in index_obj.indexes.index_key or []:
-                    if self._is_valid_key(key):
-                        condition = ""
-                        if index_obj.indexes.condition:
-                            condition = f"AND {index_obj.indexes.condition}"
-                        index_condition.add(f"({key} is not missing {condition})")
-        if index_condition:
-            self.index_condition_map[
-                (bucket_name, schema_name)
-            ] = "WHERE " + " OR ".join(index_condition)
-            return self.index_condition_map[(bucket_name, schema_name)]
-
-        self.index_condition_map[(bucket_name, schema_name)] = ""
-        return ""
-
-    def _fetch_document_ids(self, schema_name: str, table_name: str) -> List[str]:
-        # Parameters for pagination
-        page_size = 100  # Number of documents to fetch per page
-        document_ids = []
-
-        try:
-            page = 0
-            while len(document_ids) < 1000:
-                # REST API endpoint with pagination
-                couchbase_port = 8091
-                database_name = self.context.get().database
-                connection_scheme = "http"
-                url = (
-                    f"{connection_scheme}://{self.service_connection.hostport}:{couchbase_port}"
-                    f"/pools/default/buckets/{database_name}/scopes/{schema_name}/collections/{table_name}/docs"
-                )
-                params = {"skip": page * page_size, "limit": page_size}
-
-                # Fetch the next page
-                response = requests.get(
-                    url,
-                    auth=HTTPBasicAuth(
-                        self.service_connection.username,
-                        self.service_connection.password.get_secret_value(),
-                    ),
-                    params=params,
-                )
-                response.raise_for_status()
-                data = response.json()
-
-                # Add document IDs from the current page
-                ids = [doc["id"] for doc in data["rows"]]
-                document_ids.extend(ids)
-
-                # If fewer documents than page_size were returned, we’re at the end
-                if len(ids) < page_size:
-                    break
-
-                # Increment the page number
-                page += 1
-
-            # Limit the list to the first 1,000 IDs
-            document_ids = document_ids[:1000]
-
-            return document_ids
-
-        except Exception as exc:
-            logger.debug(f"Error fetching document IDs: [{table_name}]: {exc}")
-            logger.debug(traceback.format_exc())
-
-        logger.info("Document IDs not found")
-        return []
-
     def get_table_columns_dict(self, schema_name: str, table_name: str) -> List[Dict]:
         """
         Method to get actual data available within table
         need to be overridden by sources
         """
-        from couchbase.exceptions import QueryIndexNotFoundException
-
         try:
             database_name = self.context.get().database
-            bucket = self.couchbase.bucket(database_name)
-            scope = bucket.scope(schema_name)
-            collection = scope.collection(table_name)
-            documents = collection.get_multi(
-                self._fetch_document_ids(schema_name, table_name)
+            query_coln = COUCHBASE_GET_DATA.format(
+                database_name=database_name,
+                schema_name=schema_name,
+                table_name=table_name,
+                sample_size=SAMPLE_SIZE,
             )
-            return list(
-                map(
-                    lambda x: x.value if isinstance(x.value, dict) else {},
-                    documents.results.values(),
-                )
-            )
-        except QueryIndexNotFoundException as exp:
-            logger.warning(
-                f"Fetching columns failed for [`{database_name}`.`{schema_name}`.`{table_name}`],"
-                " check if the index is created for the table or data exists in the table"
-            )
-            logger.debug(traceback.format_exc())
+            query_iter = self.couchbase.query(query_coln)
+            return list(query_iter.rows())
         except Exception as exp:
             logger.debug(f"Failed to list column names for table [{table_name}]: {exp}")
             logger.debug(traceback.format_exc())

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

@@ -14,11 +14,6 @@ SQL Queries used during ingestion
 
 import textwrap
 
-COUCHBASE_GET_INDEX_KEYS = textwrap.dedent(
-    """ select * from system:indexes where {condition}"""
-)
-
-
 COUCHBASE_GET_DATA = textwrap.dedent(
-    """ select * from `{database_name}`.`{schema_name}`.`{table_name}` {condition} limit {sample_size} """
+    """ select crc.* from `{database_name}`.`{schema_name}`.`{table_name}` crc limit {sample_size} """
 )

metadata/ingestion/source/database/database_service.py

@@ -58,6 +58,9 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
+from metadata.ingestion.connections.test_connections import (
+    raise_test_connection_exception,
+)
 from metadata.ingestion.models.life_cycle import OMetaLifeCycleData
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.topology import (
@@ -609,4 +612,7 @@ class DatabaseServiceSource(
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
-        test_connection_fn(self.metadata, self.connection_obj, self.service_connection)
+        result = test_connection_fn(
+            self.metadata, self.connection_obj, self.service_connection
+        )
+        raise_test_connection_exception(result)

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

@@ -25,6 +25,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
@@ -39,6 +42,7 @@ from metadata.ingestion.source.database.databricks.client import DatabricksClien
 from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
 )
+from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -71,7 +75,8 @@ def test_connection(
     connection: Engine,
     service_connection: DatabricksConnection,
     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
@@ -104,10 +109,10 @@ def test_connection(
         "GetQueries": client.test_query_api_access,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
-        timeout_seconds=service_connection.connectionTimeout,
+        timeout_seconds=service_connection.connectionTimeout or timeout_seconds,
     )

metadata/ingestion/source/database/databricks/metadata.py

@@ -684,11 +684,13 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                     table_name=table_name,
                 )
             )
+            found_location = False
             for result in list(cursor):
                 data = result.values()
                 if data[0] and data[0].strip() == "Comment":
                     description = data[1] if data and data[1] else None
                 elif data[0] and data[0].strip() == "Location":
+                    found_location = True
                     self.external_location_map[
                         (self.context.get().database, schema_name, table_name)
                     ] = (
@@ -696,6 +698,10 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                         if data and data[1] and not data[1].startswith("dbfs")
                         else None
                     )
+            if not found_location:
+                logger.debug(
+                    f"Location not found for table {self.context.get().database}.{schema_name}.{table_name}"
+                )
 
         # Catch any exception without breaking the ingestion
         except Exception as exc:  # pylint: disable=broad-except
@@ -705,6 +711,19 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
             )
         return description
 
+    def get_location_path(self, table_name: str, schema_name: str) -> Optional[str]:
+        """
+        Method to fetch the location path of the table
+        """
+        location = self.external_location_map.get(
+            (self.context.get().database, schema_name, table_name)
+        )
+        if not location:
+            logger.debug(
+                f"Location not found in map for table {self.context.get().database}.{schema_name}.{table_name}"
+            )
+        return location
+
     def _filter_owner_name(self, owner_name: str) -> str:
         """remove unnecessary keyword from name"""
         pattern = r"\(Unknown\)"

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

@@ -31,6 +31,9 @@ from metadata.generated.schema.entity.services.connections.database.datalake.s3C
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.datalake.clients.azure_blob import (
@@ -38,6 +41,7 @@ from metadata.ingestion.source.database.datalake.clients.azure_blob import (
 )
 from metadata.ingestion.source.database.datalake.clients.gcs import DatalakeGcsClient
 from metadata.ingestion.source.database.datalake.clients.s3 import DatalakeS3Client
+from metadata.utils.constants import THREE_MIN
 
 
 # Only import specific datalake dependencies if necessary
@@ -91,7 +95,8 @@ def test_connection(
     connection: DatalakeClient,
     service_connection: DatalakeConnection,
     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
@@ -102,9 +107,10 @@ def test_connection(
         ),
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )