diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index e7821c763..509abc19c 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -482,6 +482,16 @@ class LanceNamespaceDBConnection(DBConnection): def serialize(self) -> str: import json + if ( + self._namespace_client_impl is None + or self._namespace_client_properties is None + ): + raise ValueError( + "Cannot serialize a namespace connection constructed from an " + "opaque namespace client. Pass namespace_client_impl and " + "namespace_client_properties when constructing the connection." + ) + return json.dumps( { "connection_type": "namespace", diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 14d086fdf..e533414b5 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2113,6 +2113,7 @@ class _LanceTableReopenState: connection_state: Optional[str] can_reopen_after_fork: bool + fork_reopen_error: Optional[str] name: str namespace_path: List[str] storage_options: Optional[Dict[str, str]] @@ -2196,18 +2197,34 @@ class LanceTable(Table): # A native table owns object-store clients and connection pools. Those # handles must not be used after fork, so retain a process-independent # connection description while it is still safe to inspect the parent - # connection. Some tables constructed from a bare Rust handle cannot - # be serialized; those retain the historical best-effort behavior. + # connection. Connections without reconstructible metadata are not + # safe to reuse in a forked child, so retain a clear diagnostic rather + # than advertising them as reopenable based on JSON encoding alone. + try: + connection_uri: Optional[str] = self._conn.uri + except Exception: + connection_uri = None + + fork_reopen_error: Optional[str] = None try: connection_state: Optional[str] = self._conn.serialize() - can_reopen_after_fork = not self._conn.uri.startswith("memory://") - except Exception: + can_reopen_after_fork = connection_uri is not None and not ( + connection_uri.startswith("memory://") + ) + except Exception as error: connection_state = None can_reopen_after_fork = False + if connection_uri is not None and not connection_uri.startswith( + "memory://" + ): + fork_reopen_error = ( + f"Cannot reopen table {name!r} in a forked process: {error}" + ) self._reopen_state = _LanceTableReopenState( connection_state=connection_state, can_reopen_after_fork=can_reopen_after_fork, + fork_reopen_error=fork_reopen_error, name=name, namespace_path=list(self._namespace_path), storage_options=( @@ -2330,6 +2347,9 @@ class LanceTable(Table): return state = getattr(self, "_reopen_state", None) + fork_reopen_error = getattr(state, "fork_reopen_error", None) + if fork_reopen_error is not None: + raise RuntimeError(fork_reopen_error) if ( state is None or not state.can_reopen_after_fork diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index a5ace932f..0ef146d5e 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -6,10 +6,13 @@ import tempfile import shutil import importlib +import multiprocessing as mp +import sys from datetime import timedelta import pytest import pyarrow as pa import lancedb +from lance_namespace import connect as namespace_connect from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError from lancedb.namespace import _MAX_QUERY_K from lancedb.table import AsyncTable, LanceTable @@ -73,6 +76,16 @@ def _namespace_lance_table(namespace_client: _NamespaceClient) -> LanceTable: return table +def _direct_namespace_fork_child(table, result_queue): + from lancedb.permutation import Permutation + + try: + permutation = Permutation.identity(table) + result_queue.put(("ok", permutation.num_rows)) + except Exception as error: + result_queue.put((type(error).__name__, str(error))) + + class TestNamespaceConnection: """Test namespace-based LanceDB connection using DirectoryNamespace.""" @@ -446,6 +459,70 @@ class TestNamespaceConnection: restored = lancedb.deserialize_conn(db.serialize()) assert restored.read_consistency_interval == timedelta(0) + @pytest.mark.skipif( + sys.platform != "linux", + reason="fork() is only supported safely for this test on Linux", + ) + def test_direct_namespace_with_descriptor_reopens_after_fork(self): + properties = {"root": self.temp_dir} + namespace = namespace_connect("dir", properties) + db = lancedb.LanceNamespaceDBConnection( + namespace, + namespace_client_impl="dir", + namespace_client_properties=properties, + ) + table = db.create_table("items", pa.table({"id": [1]})) + + ctx = mp.get_context("fork") + result_queue = ctx.Queue() + process = ctx.Process( + target=_direct_namespace_fork_child, + args=(table, result_queue), + ) + process.start() + process.join(10) + + if process.is_alive(): + process.terminate() + process.join(5) + pytest.fail("Direct namespace table hung while reopening after fork") + + assert process.exitcode == 0 + assert result_queue.get(timeout=2) == ("ok", 1) + + @pytest.mark.skipif( + sys.platform != "linux", + reason="fork() is only supported safely for this test on Linux", + ) + def test_opaque_direct_namespace_reports_unsupported_fork(self): + namespace = namespace_connect("dir", {"root": self.temp_dir}) + db = lancedb.LanceNamespaceDBConnection(namespace) + table = db.create_table("items", pa.table({"id": [1]})) + + with pytest.raises(ValueError, match="opaque namespace client"): + db.serialize() + assert not table._can_reopen_after_fork + + ctx = mp.get_context("fork") + result_queue = ctx.Queue() + process = ctx.Process( + target=_direct_namespace_fork_child, + args=(table, result_queue), + ) + process.start() + process.join(10) + + if process.is_alive(): + process.terminate() + process.join(5) + pytest.fail("Opaque namespace table hung after fork") + + assert process.exitcode == 0 + error_type, message = result_queue.get(timeout=2) + assert error_type == "RuntimeError" + assert "Cannot reopen table 'items' in a forked process" in message + assert "namespace_client_impl and namespace_client_properties" in message + def test_namespace_operations(self): """Test namespace management operations.""" db = lancedb.connect_namespace("dir", {"root": self.temp_dir})