diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 60e1780f0..330514484 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -122,7 +122,48 @@ def _view_description( class DBConnection(EnforceOverrides): - """An active LanceDB connection interface.""" + """An active LanceDB connection interface. + + Use [close][lancedb.db.DBConnection.close] to release the connection's + underlying resources, or use the connection as a context manager to close it + automatically when leaving the block, including when an exception is raised. + + Examples + -------- + >>> import lancedb + >>> with lancedb.connect("memory://") as db: + ... assert db.is_open() + >>> db.is_open() + False + """ + + def __enter__(self) -> DBConnection: + return self + + def __exit__(self, *_) -> None: + self.close() + + @abstractmethod + def is_open(self) -> bool: + """Return True if the connection is open.""" + pass + + @abstractmethod + def close(self) -> None: + """Close the connection, releasing any underlying resources. + + It is safe to call this method multiple times. Database operations on a + closed connection raise ``RuntimeError: Connection is closed``. + + Examples + -------- + >>> import lancedb + >>> db = lancedb.connect("memory://") + >>> db.close() + >>> db.is_open() + False + """ + pass def list_namespaces( self, @@ -1275,6 +1316,14 @@ class LanceDBConnection(DBConnection): def __repr__(self) -> str: return f"{self.__class__.__name__}(uri={self._conn.uri!r})" + @override + def is_open(self) -> bool: + return self._conn.is_open() + + @override + def close(self) -> None: + self._conn.close() + @override def serialize(self) -> str: import json diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index cf1910c68..16be63e2d 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -489,6 +489,14 @@ class LanceNamespaceDBConnection(DBConnection): ) self._uri = self._inner.uri + @override + def is_open(self) -> bool: + return self._inner.is_open() + + @override + def close(self) -> None: + self._inner.close() + @override def serialize(self) -> str: import json diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 3773ac547..52323daef 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -1054,6 +1054,10 @@ class RemoteDBConnection(DBConnection): """ return LOOP.run(self._conn.namespace_client()) - async def close(self): - """Close the connection to the database.""" + @override + def is_open(self) -> bool: + return self._conn.is_open() + + @override + def close(self) -> None: self._conn.close() diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index aeb9feeb8..29f02c09e 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import contextlib import inspect import re import sys @@ -405,6 +406,49 @@ async def test_connect(tmp_path): assert str(db) == f"ListingDatabase(uri={tmp_path}, read_consistency_interval=5s)" +@pytest.fixture(params=["memory", "local", "namespace", "from_inner"]) +def sync_db(request, tmp_path): + if request.param == "memory": + return lancedb.connect("memory://") + if request.param == "namespace": + return lancedb.connect_namespace("dir", {"root": str(tmp_path)}) + db = lancedb.connect(tmp_path) + if request.param == "from_inner": + return lancedb.db.LanceDBConnection.from_inner(db._inner, None) + return db + + +def test_sync_close(sync_db): + assert sync_db.is_open() + assert sync_db.close() is None + assert not sync_db.is_open() + assert sync_db.close() is None + + with pytest.warns(DeprecationWarning, match="table_names"): + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.table_names() + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.list_tables() + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.open_table("test") + + +@pytest.mark.parametrize("raise_error", [False, True]) +def test_sync_context_manager(sync_db, raise_error): + with contextlib.ExitStack() as stack: + if raise_error: + stack.enter_context(pytest.raises(ValueError, match="test error")) + with sync_db as db: + assert db is sync_db + assert db.is_open() + if raise_error: + raise ValueError("test error") + + assert not sync_db.is_open() + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.list_tables() + + @pytest.mark.asyncio async def test_close(mem_db_async: lancedb.AsyncConnection): assert mem_db_async.is_open() diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 8f534f678..7810fd281 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2145,15 +2145,46 @@ async def test_header_provider_overrides_static_headers(): def test_close(): - """Test that close() works without AttributeError.""" - import asyncio - def handler(req): req.send_response(200) req.end_headers() with mock_lancedb_connection(handler) as db: - asyncio.run(db.close()) + assert db.close() is None + assert not db.is_open() + assert db.close() is None + + with pytest.warns(DeprecationWarning, match="table_names"): + with pytest.raises(RuntimeError, match="Connection is closed"): + db.table_names() + with pytest.raises(RuntimeError, match="Connection is closed"): + db.list_tables() + with pytest.raises(RuntimeError, match="Connection is closed"): + db.open_table("test") + + +@pytest.mark.parametrize("raise_error", [False, True]) +def test_sync_context_manager(raise_error): + def handler(req): + req.send_response(200) + req.send_header("Content-Type", "application/json") + req.end_headers() + req.wfile.write(b'{"tables": []}') + + with mock_lancedb_connection(handler) as db: + with contextlib.ExitStack() as stack: + if raise_error: + stack.enter_context(pytest.raises(ValueError, match="test error")) + with db as conn: + assert conn is db + assert conn.is_open() + assert conn.list_tables().tables == [] + if raise_error: + raise ValueError("test error") + + assert not db.is_open() + with pytest.raises(RuntimeError, match="Connection is closed"): + db.list_tables() @pytest.mark.parametrize("exception", [KeyboardInterrupt, SystemExit, GeneratorExit])