Add `close()` method to RestCatalog (#2403)
<!--
Thanks for opening a pull request!
-->
<!-- In the case this PR will resolve an issue, please replace
${GITHUB_ISSUE_ID} below with the actual Github issue id. -->
Part of #2399
# Rationale for this change
This PR introducs the implementation of `close()` method to the
`RestCatalog`.
And add corresponding test case which following the test pattern in
#2390
## Are these changes tested?
Yes
## Are there any user-facing changes?
<!-- In the case of user-facing changes, please add the changelog label.
-->diff --git a/pyiceberg/catalog/rest/__init__.py b/pyiceberg/catalog/rest/__init__.py
index afbaa7d..207b3c4 100644
--- a/pyiceberg/catalog/rest/__init__.py
+++ b/pyiceberg/catalog/rest/__init__.py
@@ -876,3 +876,10 @@
response.raise_for_status()
except HTTPError as exc:
_handle_non_200_response(exc, {404: NoSuchViewError})
+
+ def close(self) -> None:
+ """Close the catalog and release Session connection adapters.
+
+ This method closes mounted HttpAdapters' pooled connections and any active Proxy pooled connections.
+ """
+ self._session.close()
diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py
index 2619105..223c6d2 100644
--- a/tests/catalog/test_rest.py
+++ b/tests/catalog/test_rest.py
@@ -1845,3 +1845,76 @@
assert len(history) == 1
actual_headers = history[0].headers
assert actual_headers["Authorization"] == expected_auth_header
+
+
+class TestRestCatalogClose:
+ """Tests RestCatalog close functionality"""
+
+ EXPECTED_ADAPTERS = 2
+ EXPECTED_ADAPTERS_SIGV4 = 3
+
+ def test_catalog_close(self, rest_mock: Mocker) -> None:
+ rest_mock.get(
+ f"{TEST_URI}v1/config",
+ json={"defaults": {}, "overrides": {}},
+ status_code=200,
+ )
+
+ catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+ catalog.close()
+ # Verify session still exists after close the session pooled connections
+ assert hasattr(catalog, "_session")
+ assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS
+ # Second close should not raise any exception
+ catalog.close()
+
+ def test_rest_catalog_close_sigv4(self, rest_mock: Mocker) -> None:
+ catalog = None
+ rest_mock.get(
+ f"{TEST_URI}v1/config",
+ json={"defaults": {}, "overrides": {}},
+ status_code=200,
+ )
+
+ catalog = RestCatalog("rest", **{"uri": TEST_URI, "token": TEST_TOKEN, "rest.sigv4-enabled": "true"})
+ catalog.close()
+ assert hasattr(catalog, "_session")
+ assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS_SIGV4
+
+ def test_rest_catalog_context_manager_with_exception(self, rest_mock: Mocker) -> None:
+ """Test RestCatalog context manager properly closes with exceptions."""
+ catalog = None
+ rest_mock.get(
+ f"{TEST_URI}v1/config",
+ json={"defaults": {}, "overrides": {}},
+ status_code=200,
+ )
+
+ try:
+ with RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) as cat:
+ catalog = cat
+ raise ValueError("Test exception")
+ except ValueError:
+ pass
+
+ assert catalog is not None and hasattr(catalog, "_session")
+ assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS
+
+ def test_rest_catalog_context_manager_with_exception_sigv4(self, rest_mock: Mocker) -> None:
+ """Test RestCatalog context manager properly closes with exceptions."""
+ catalog = None
+ rest_mock.get(
+ f"{TEST_URI}v1/config",
+ json={"defaults": {}, "overrides": {}},
+ status_code=200,
+ )
+
+ try:
+ with RestCatalog("rest", **{"uri": TEST_URI, "token": TEST_TOKEN, "rest.sigv4-enabled": "true"}) as cat:
+ catalog = cat
+ raise ValueError("Test exception")
+ except ValueError:
+ pass
+
+ assert catalog is not None and hasattr(catalog, "_session")
+ assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS_SIGV4
