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

metadata/ingestion/source/search/search_service.py

@@ -41,6 +41,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.delete_entity import DeleteEntity
 from metadata.ingestion.models.search_index_data import OMetaIndexSampleData
 from metadata.ingestion.models.topology import (
@@ -192,7 +195,10 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
 
     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)
 
     def mark_search_indexes_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the search index as deleted"""

metadata/ingestion/source/storage/gcs/connection.py

@@ -22,6 +22,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.storage.gcsConnection import (
     GcsConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.generated.schema.security.credentials.gcpValues import (
     GcpCredentialsValues,
     SingleProjectId,
@@ -32,6 +35,7 @@ from metadata.ingestion.connections.test_connections import (
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.storage.gcs.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
 
@@ -136,7 +140,8 @@ def test_connection(
     client: GcsObjectStoreClient,
     service_connection: GcsConnection,
     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
@@ -151,9 +156,10 @@ def test_connection(
         "GetMetrics": tester.get_metrics,
     }
 
-    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/storage/s3/connection.py

@@ -27,8 +27,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.storage.s3Connection import (
     S3Connection,
 )
+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
 
 
 @dataclass
@@ -53,7 +57,8 @@ def test_connection(
     client: S3ObjectStoreClient,
     service_connection: S3Connection,
     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,9 +80,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,
     )

metadata/ingestion/source/storage/storage_service.py

@@ -40,6 +40,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.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.topology import (
@@ -236,7 +239,10 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
 
     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)
 
     def mark_containers_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the containers as deleted"""

metadata/utils/constants.py

@@ -39,6 +39,7 @@ from metadata.generated.schema.entity.teams.user import User
 
 DOT = "_DOT_"
 TEN_MIN = 10 * 60
+THREE_MIN = 3 * 60
 UTF_8 = "utf-8"
 CHUNKSIZE = 200000
 DEFAULT_DATABASE = "default"

metadata/utils/timeout.py

@@ -18,6 +18,7 @@ import inspect
 import os
 import platform
 import signal
+import threading
 from typing import Callable
 
 from metadata.utils.constants import TEN_MIN
@@ -47,7 +48,11 @@ def timeout(seconds: int = TEN_MIN) -> Callable:
     def decorator(fn):
         @functools.wraps(fn)
         def inner(*args, **kwargs):
-            if platform.system() != "Windows":  # SIGALRM not supported on Windows
+            # SIGALRM is not supported on Windows or sub-threads
+            if (
+                platform.system() != "Windows"
+                and threading.current_thread() == threading.main_thread()
+            ):
                 signal.signal(signal.SIGALRM, _handle_timeout)
                 signal.alarm(seconds)
                 try:

metadata/workflow/profiler.py

@@ -20,6 +20,9 @@ from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
 from metadata.ingestion.api.steps import Processor, Sink
+from metadata.ingestion.connections.test_connections import (
+    raise_test_connection_exception,
+)
 from metadata.ingestion.source.connections import get_test_connection_fn
 from metadata.pii.processor import PIIProcessor
 from metadata.profiler.processor.processor import ProfilerProcessor
@@ -75,12 +78,13 @@ class ProfilerWorkflow(IngestionWorkflow):
         else:
             self.steps = (profiler_processor, sink)
 
-    def test_connection(self):
+    def test_connection(self) -> None:
         service_config = self.config.source.serviceConnection.root.config
         conn = get_ssl_connection(service_config)
 
         test_connection_fn = get_test_connection_fn(service_config)
-        test_connection_fn(self.metadata, conn, service_config)
+        result = test_connection_fn(self.metadata, conn, service_config)
+        raise_test_connection_exception(result)
 
     def _get_sink(self) -> Sink:
         sink_type = self.config.sink.type