feat: add query_arrow/pandas/polars/pylist helpers on Client

Adds one-shot output ergonomics so callers don't have to reach for
.read_all().to_pandas() on the underlying Flight or ADBC reader. Each
helper accepts an optional `params` list that routes through the existing
ADBC parameterized path, otherwise it goes through Flight as before.

polars is gated behind an optional extra (`spicepy[polars]`) and raises
a clear ImportError if missing.

Author: claude

pyproject.toml

@@ -56,11 +56,15 @@ test = [
     "pyarrow>=23.0.1",
     "adbc-driver-flightsql>=1.11.0",
     "adbc-driver-manager>=1.11.0",
+    "polars>=1.0.0",
 ]
 params = [
     "adbc-driver-flightsql>=1.11.0",
     "adbc-driver-manager>=1.11.0",
 ]
+polars = [
+    "polars>=1.0.0",
+]
 
 # ============== Tool Configuration ==============
 
@@ -93,6 +97,8 @@ module = [
     "adbc_driver_manager.*",
     "certifi",
     "pandas",
+    "polars",
+    "polars.*",
     "_pytest.*",
     "pytest.*",
 ]

spicepy/_client.py

@@ -3,11 +3,15 @@
 from pathlib import Path
 import platform
 import threading
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
 import certifi
 import pyarrow as pa
 
+if TYPE_CHECKING:
+    import pandas as pd
+    import polars as pl
+
 # pylint: disable=E0611
 from pyarrow._flight import (
     FlightCallOptions,
@@ -409,6 +413,91 @@ def query_with_params(
         adbc = self._ensure_adbc_client()
         return adbc.query_with_params(sql, params)
 
+    def _read_table(
+        self,
+        sql: str,
+        params: list[Any] | None,
+        timeout: int | None,
+    ) -> pa.Table:
+        if params is not None:
+            return self.query_with_params(sql, params).read_all()
+        kwargs: dict[str, Any] = {}
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+        return self.query(sql, **kwargs).read_all()
+
+    def query_arrow(
+        self,
+        sql: str,
+        *,
+        params: list[Any] | None = None,
+        timeout: int | None = None,
+    ) -> pa.Table:
+        """Execute a SQL query and return results as a PyArrow Table.
+
+        Args:
+            sql: SQL query string. Use $1, $2, ... placeholders if passing params.
+            params: Optional list of parameter values. When provided, the query is
+                executed via ADBC FlightSQL with prepared statements. See
+                :meth:`query_with_params` for parameter format.
+            timeout: Optional query timeout in seconds (ignored when params is set).
+
+        Returns:
+            Arrow Table with all query results materialized in memory.
+        """
+        return self._read_table(sql, params, timeout)
+
+    def query_pandas(
+        self,
+        sql: str,
+        *,
+        params: list[Any] | None = None,
+        timeout: int | None = None,
+    ) -> "pd.DataFrame":
+        """Execute a SQL query and return results as a pandas DataFrame.
+
+        See :meth:`query_arrow` for argument semantics.
+        """
+        return cast("pd.DataFrame", self._read_table(sql, params, timeout).to_pandas())
+
+    def query_polars(
+        self,
+        sql: str,
+        *,
+        params: list[Any] | None = None,
+        timeout: int | None = None,
+    ) -> "pl.DataFrame":
+        """Execute a SQL query and return results as a polars DataFrame.
+
+        Requires the optional ``polars`` dependency:
+        ``pip install spicepy[polars]``.
+
+        See :meth:`query_arrow` for argument semantics.
+        """
+        try:
+            import polars as pl
+        except ImportError as exc:
+            raise ImportError(
+                "polars is not installed. Install it with: pip install spicepy[polars]"
+            ) from exc
+        return cast("pl.DataFrame", pl.from_arrow(self._read_table(sql, params, timeout)))
+
+    def query_pylist(
+        self,
+        sql: str,
+        *,
+        params: list[Any] | None = None,
+        timeout: int | None = None,
+    ) -> list[dict[str, Any]]:
+        """Execute a SQL query and return results as a list of row dicts.
+
+        See :meth:`query_arrow` for argument semantics.
+        """
+        return cast(
+            "list[dict[str, Any]]",
+            self._read_table(sql, params, timeout).to_pylist(),
+        )
+
     def refresh_dataset(
         self, dataset: str, refresh_opts: RefreshOpts | None = None
     ) -> Any:

tests/test_client.py

@@ -1017,6 +1017,202 @@ def test_thread_join_raises_exception(self) -> None:
                 thread.join()
 
 
+class TestClientQueryHelpers:
+    """Test Client query_arrow / query_pandas / query_polars / query_pylist."""
+
+    @staticmethod
+    def _sample_table() -> pa.Table:
+        return pa.table({"id": [1, 2, 3], "name": ["a", "b", "c"]})
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    def test_query_arrow_returns_table(
+        self,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        table = self._sample_table()
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = table
+        mock_flight = MagicMock()
+        mock_flight.query.return_value = mock_reader
+        mock_flight_class.return_value = mock_flight
+
+        client = Client()
+        result = client.query_arrow("SELECT * FROM t")
+
+        assert result is table
+        mock_flight.query.assert_called_once_with("SELECT * FROM t")
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    def test_query_arrow_passes_timeout(
+        self,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = self._sample_table()
+        mock_flight = MagicMock()
+        mock_flight.query.return_value = mock_reader
+        mock_flight_class.return_value = mock_flight
+
+        client = Client()
+        client.query_arrow("SELECT 1", timeout=30)
+
+        mock_flight.query.assert_called_once_with("SELECT 1", timeout=30)
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    @patch("spicepy._client._ADBCClient")
+    def test_query_arrow_routes_params_to_adbc(
+        self,
+        mock_adbc_class: MagicMock,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        table = self._sample_table()
+        adbc_reader = MagicMock()
+        adbc_reader.read_all.return_value = table
+        mock_adbc = MagicMock()
+        mock_adbc.query_with_params.return_value = adbc_reader
+        mock_adbc_class.return_value = mock_adbc
+
+        flight_instance = MagicMock()
+        mock_flight_class.return_value = flight_instance
+
+        client = Client()
+        result = client.query_arrow("SELECT * FROM t WHERE id = $1", params=[1])
+
+        assert result is table
+        mock_adbc.query_with_params.assert_called_once_with(
+            "SELECT * FROM t WHERE id = $1", [1]
+        )
+        flight_instance.query.assert_not_called()
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    def test_query_pandas_returns_dataframe(
+        self,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        import pandas as pd
+
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = self._sample_table()
+        mock_flight = MagicMock()
+        mock_flight.query.return_value = mock_reader
+        mock_flight_class.return_value = mock_flight
+
+        client = Client()
+        df = client.query_pandas("SELECT * FROM t")
+
+        assert isinstance(df, pd.DataFrame)
+        assert list(df.columns) == ["id", "name"]
+        assert len(df) == 3
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    def test_query_pylist_returns_row_dicts(
+        self,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = self._sample_table()
+        mock_flight = MagicMock()
+        mock_flight.query.return_value = mock_reader
+        mock_flight_class.return_value = mock_flight
+
+        client = Client()
+        rows = client.query_pylist("SELECT * FROM t")
+
+        assert rows == [
+            {"id": 1, "name": "a"},
+            {"id": 2, "name": "b"},
+            {"id": 3, "name": "c"},
+        ]
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    def test_query_polars_returns_dataframe(
+        self,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        pl = pytest.importorskip("polars")
+
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = self._sample_table()
+        mock_flight = MagicMock()
+        mock_flight.query.return_value = mock_reader
+        mock_flight_class.return_value = mock_flight
+
+        client = Client()
+        df = client.query_polars("SELECT * FROM t")
+
+        assert isinstance(df, pl.DataFrame)
+        assert df.columns == ["id", "name"]
+        assert df.height == 3
+
+    @patch("spicepy._client._SpiceFlight")
+    @patch("spicepy._client._Cert")
+    def test_query_polars_missing_raises_importerror(
+        self,
+        mock_cert_class: MagicMock,
+        mock_flight_class: MagicMock,
+    ) -> None:
+        mock_cert = MagicMock()
+        mock_cert.tls_root_certs = b"cert"
+        mock_cert_class.return_value = mock_cert
+
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = self._sample_table()
+        mock_flight = MagicMock()
+        mock_flight.query.return_value = mock_reader
+        mock_flight_class.return_value = mock_flight
+
+        client = Client()
+
+        import builtins
+
+        real_import = builtins.__import__
+
+        def fake_import(name, *args, **kwargs):
+            if name == "polars":
+                raise ImportError("No module named 'polars'")
+            return real_import(name, *args, **kwargs)
+
+        with patch("builtins.__import__", side_effect=fake_import):
+            with pytest.raises(ImportError, match="polars is not installed"):
+                client.query_polars("SELECT 1")
+
+
 class TestIsMacosArm64:
     """Test is_macos_arm64 function."""