fix(ingest): Closeable as a context manager

metadata-ingestion/src/datahub/emitter/rest_emitter.py

@@ -4,8 +4,7 @@
 import logging
 import os
 from json.decoder import JSONDecodeError
-from types import TracebackType
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import requests
 from requests.adapters import HTTPAdapter, Retry
@@ -16,6 +15,7 @@
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.request_helper import _make_curl_command
 from datahub.emitter.serialization_helper import pre_json_transform
+from datahub.ingestion.api.closeable import Closeable
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
     MetadataChangeProposal,
@@ -25,7 +25,7 @@
 logger = logging.getLogger(__name__)
 
 
-class DataHubRestEmitter:
+class DataHubRestEmitter(Closeable):
     DEFAULT_CONNECT_TIMEOUT_SEC = 30  # 30 seconds should be plenty to connect
     DEFAULT_READ_TIMEOUT_SEC = (
         30  # Any ingest call taking longer than 30 seconds should be abandoned
@@ -142,17 +142,6 @@ def __init__(
             timeout=(self._connect_timeout_sec, self._read_timeout_sec),
         )
 
-    def __enter__(self) -> "DataHubRestEmitter":
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self._session.close()
-
     def test_connection(self) -> dict:
         response = self._session.get(f"{self._gms_server}/config")
         if response.status_code == 200:
@@ -277,6 +266,9 @@ def __repr__(self) -> str:
             f"DataHubRestEmitter: configured to talk to {self._gms_server}{token_str}"
         )
 
+    def close(self) -> None:
+        self._session.close()
+
 
 class DatahubRestEmitter(DataHubRestEmitter):
     """This class exists as a pass-through for backwards compatibility"""

metadata-ingestion/src/datahub/ingestion/api/closeable.py

@@ -1,7 +1,18 @@
 from abc import abstractmethod
+from contextlib import AbstractContextManager
+from types import TracebackType
+from typing import Optional, Type
 
 
-class Closeable:
+class Closeable(AbstractContextManager):
     @abstractmethod
     def close(self) -> None:
         pass
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self.close()

metadata-ingestion/src/datahub/ingestion/graph/client.py

@@ -83,12 +83,6 @@ def __init__(self, config: Union[DatahubClientConfig, DataHubGraphConfig]) -> No
             self.server_id = "missing"
             logger.debug(f"Failed to get server id due to {e}")
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, type, value, tb):
-        return super().__exit__(type, value, tb)
-
     def _get_generic(self, url: str) -> Dict:
         try:
             response = self._session.get(url)

metadata-ingestion/tests/integration/kafka/docker-compose.yml

@@ -22,6 +22,7 @@ services:
 
   schema-registry:
     image: confluentinc/cp-schema-registry:7.2.0
+    restart: unless-stopped
     env_file: schema-registry.env
     container_name: test_schema_registry
     depends_on:
@@ -31,5 +32,4 @@ services:
       - "58081:8081"
 
 volumes:
-    test_zkdata:
-
+  test_zkdata: