diff --git a/python/python/tests/test_closed_table_query.py b/python/python/tests/test_closed_table_query.py new file mode 100644 index 000000000..2ff3a8e98 --- /dev/null +++ b/python/python/tests/test_closed_table_query.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest +from lancedb.db import AsyncConnection + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["count_rows", "query", "head"]) +async def test_closed_table_queries_raise_runtime_error( + mem_db_async: AsyncConnection, operation +): + table = await mem_db_async.create_table("closed_query", data=[{"id": 0}]) + assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}] + table.close() + table.close() + assert not table.is_open() + + with pytest.raises(RuntimeError, match="Table closed_query is closed"): + if operation == "query": + table.query() + else: + await getattr(table, operation)() + + +@pytest.mark.asyncio +async def test_open_table_queries(mem_db_async: AsyncConnection): + table = await mem_db_async.create_table("open_query", data=[{"id": 0}]) + assert await table.count_rows() == 1 + assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}] + assert (await table.head()).to_pylist() == [{"id": 0}] + table.close() diff --git a/python/src/table.rs b/python/src/table.rs index c1d23d151..27ed8be8b 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1428,8 +1428,8 @@ impl Table { }) } - pub fn query(&self) -> Query { - Query::new(self.inner_ref().unwrap().query()) + pub fn query(&self) -> PyResult { + Ok(Query::new(self.inner_ref()?.query())) } #[getter]