Skip to content

Commit 86b0fe8

Browse files
authored
Add close() method to RestCatalog (apache#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 apache#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 apache#2390 ## Are these changes tested? Yes ## Are there any user-facing changes? <!-- In the case of user-facing changes, please add the changelog label. -->
1 parent c3c314d commit 86b0fe8

File tree

2 files changed

+80
-0
lines changed

2 files changed

+80
-0
lines changed

pyiceberg/catalog/rest/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -876,3 +876,10 @@ def drop_view(self, identifier: Union[str]) -> None:
876876
response.raise_for_status()
877877
except HTTPError as exc:
878878
_handle_non_200_response(exc, {404: NoSuchViewError})
879+
880+
def close(self) -> None:
881+
"""Close the catalog and release Session connection adapters.
882+
883+
This method closes mounted HttpAdapters' pooled connections and any active Proxy pooled connections.
884+
"""
885+
self._session.close()

tests/catalog/test_rest.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1845,3 +1845,76 @@ def test_rest_catalog_with_google_credentials_path(
18451845
assert len(history) == 1
18461846
actual_headers = history[0].headers
18471847
assert actual_headers["Authorization"] == expected_auth_header
1848+
1849+
1850+
class TestRestCatalogClose:
1851+
"""Tests RestCatalog close functionality"""
1852+
1853+
EXPECTED_ADAPTERS = 2
1854+
EXPECTED_ADAPTERS_SIGV4 = 3
1855+
1856+
def test_catalog_close(self, rest_mock: Mocker) -> None:
1857+
rest_mock.get(
1858+
f"{TEST_URI}v1/config",
1859+
json={"defaults": {}, "overrides": {}},
1860+
status_code=200,
1861+
)
1862+
1863+
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
1864+
catalog.close()
1865+
# Verify session still exists after close the session pooled connections
1866+
assert hasattr(catalog, "_session")
1867+
assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS
1868+
# Second close should not raise any exception
1869+
catalog.close()
1870+
1871+
def test_rest_catalog_close_sigv4(self, rest_mock: Mocker) -> None:
1872+
catalog = None
1873+
rest_mock.get(
1874+
f"{TEST_URI}v1/config",
1875+
json={"defaults": {}, "overrides": {}},
1876+
status_code=200,
1877+
)
1878+
1879+
catalog = RestCatalog("rest", **{"uri": TEST_URI, "token": TEST_TOKEN, "rest.sigv4-enabled": "true"})
1880+
catalog.close()
1881+
assert hasattr(catalog, "_session")
1882+
assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS_SIGV4
1883+
1884+
def test_rest_catalog_context_manager_with_exception(self, rest_mock: Mocker) -> None:
1885+
"""Test RestCatalog context manager properly closes with exceptions."""
1886+
catalog = None
1887+
rest_mock.get(
1888+
f"{TEST_URI}v1/config",
1889+
json={"defaults": {}, "overrides": {}},
1890+
status_code=200,
1891+
)
1892+
1893+
try:
1894+
with RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) as cat:
1895+
catalog = cat
1896+
raise ValueError("Test exception")
1897+
except ValueError:
1898+
pass
1899+
1900+
assert catalog is not None and hasattr(catalog, "_session")
1901+
assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS
1902+
1903+
def test_rest_catalog_context_manager_with_exception_sigv4(self, rest_mock: Mocker) -> None:
1904+
"""Test RestCatalog context manager properly closes with exceptions."""
1905+
catalog = None
1906+
rest_mock.get(
1907+
f"{TEST_URI}v1/config",
1908+
json={"defaults": {}, "overrides": {}},
1909+
status_code=200,
1910+
)
1911+
1912+
try:
1913+
with RestCatalog("rest", **{"uri": TEST_URI, "token": TEST_TOKEN, "rest.sigv4-enabled": "true"}) as cat:
1914+
catalog = cat
1915+
raise ValueError("Test exception")
1916+
except ValueError:
1917+
pass
1918+
1919+
assert catalog is not None and hasattr(catalog, "_session")
1920+
assert len(catalog._session.adapters) == self.EXPECTED_ADAPTERS_SIGV4

0 commit comments

Comments
 (0)