Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions pyiceberg/catalog/rest/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -876,3 +876,10 @@ def drop_view(self, identifier: Union[str]) -> None:
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()
73 changes: 73 additions & 0 deletions tests/catalog/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -1845,3 +1845,76 @@ def test_rest_catalog_with_google_credentials_path(
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