From ede04b302ab653e79a939ff7d75f81fcbf0348f0 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 6 Aug 2026 10:06:09 -0400 Subject: [PATCH 1/3] feat: add search pagination and filter flags --- README.md | 6 ++++++ src/drs/cli.py | 11 ++++++++++- src/drs/client.py | 12 +++++++++++- src/drs/introspect.py | 18 ++++++++++++++++++ tests/test_cli.py | 33 +++++++++++++++++++++++++++++++++ tests/test_client.py | 41 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 119 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index adf2d14..681625c 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,12 @@ dremio query run "SELECT * FROM myspace.orders LIMIT 5" --output pretty # Search the catalog for anything matching "revenue" dremio search "revenue" +# Search only jobs and limit the first page size +dremio search "revenue" --filter 'category in ["JOB"]' --max-results 20 + +# Fetch the next page using the nextPageToken from a prior response +dremio search "revenue" --next-page-token 'eyJwYWdlVG9rZW4iOiJ...' + # Create a space, then a folder inside it dremio folder create "Analytics" dremio folder create Analytics.reports diff --git a/src/drs/cli.py b/src/drs/cli.py index ae2c972..f1845c6 100644 --- a/src/drs/cli.py +++ b/src/drs/cli.py @@ -170,6 +170,9 @@ def get_client() -> DremioClient: @app.command("search") def search_command( term: str = typer.Argument(help="Search term (matches table names, view names, source names)"), + filter_: str | None = typer.Option(None, "--filter", help="CEL filter expression to refine search results"), + max_results: int | None = typer.Option(None, "--max-results", min=1, help="Maximum results to return per page"), + next_page_token: str | None = typer.Option(None, "--next-page-token", help="Pagination token from a prior search response"), fmt: str = typer.Option("json", "--output", "-o", help="Output format: json, csv, pretty"), ) -> None: """Full-text search across all catalog entities (tables, views, sources).""" @@ -181,7 +184,13 @@ def search_command( async def _execute(): try: try: - return await client.search(term) + search_kwargs: dict[str, str | int | None] = { + "filter_": filter_, + "max_results": max_results, + } + if next_page_token is not None: + search_kwargs["next_page_token"] = next_page_token + return await client.search(term, **search_kwargs) except httpx.HTTPStatusError as exc: raise handle_api_error(exc) from exc finally: diff --git a/src/drs/client.py b/src/drs/client.py index 7f3af28..5762169 100644 --- a/src/drs/client.py +++ b/src/drs/client.py @@ -200,10 +200,20 @@ async def get_catalog_by_path(self, path_parts: list[str]) -> dict: joined = "/".join(path_parts) return await self._get(self._v3(f"/catalog/by-path/{joined}")) - async def search(self, query: str, filter_: str | None = None) -> dict: + async def search( + self, + query: str, + filter_: str | None = None, + max_results: int | None = None, + next_page_token: str | None = None, + ) -> dict: body: dict[str, Any] = {"query": query} if filter_: body["filter"] = filter_ + if max_results is not None: + body["maxResults"] = max_results + if next_page_token: + body["pageToken"] = next_page_token return await self._post(self._v0("/search"), json=body) async def create_catalog_entity(self, body: dict) -> dict: diff --git a/src/drs/introspect.py b/src/drs/introspect.py index bd0dbc5..c88d174 100644 --- a/src/drs/introspect.py +++ b/src/drs/introspect.py @@ -212,6 +212,24 @@ "endpoints": ["POST /v0/projects/{pid}/search"], "parameters": [ {"name": "term", "type": "string", "required": True, "positional": True, "description": "Search term"}, + { + "name": "filter", + "type": "string", + "required": False, + "description": "Optional CEL filter expression to refine search results", + }, + { + "name": "max_results", + "type": "integer", + "required": False, + "description": "Maximum number of results to return per page", + }, + { + "name": "next_page_token", + "type": "string", + "required": False, + "description": "Pagination token returned by a previous search response", + }, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, diff --git a/tests/test_cli.py b/tests/test_cli.py index 45fdfae..088d5e6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -19,6 +19,7 @@ import json from unittest.mock import AsyncMock +from unittest.mock import MagicMock from typer.testing import CliRunner @@ -217,3 +218,35 @@ async def fake_get_all_messages(client, conversation_id, page_size=200): assert '"conversationCount": 2' in report assert '"id": "conv-1"' in report assert '"id": "conv-2"' in report + + +def test_search_command_passes_filter_and_max_results(monkeypatch) -> None: + search_mock = AsyncMock(return_value={"results": []}) + close_mock = AsyncMock() + client = MagicMock() + client.search = search_mock + client.close = close_mock + + monkeypatch.setattr("drs.cli.get_client", lambda: client) + + result = runner.invoke(app, ["search", "revenue", "--filter", 'category in ["JOB"]', "--max-results", "20"]) + + assert result.exit_code == 0 + search_mock.assert_awaited_once_with("revenue", filter_='category in ["JOB"]', max_results=20) + close_mock.assert_awaited_once() + + +def test_search_command_passes_next_page_token(monkeypatch) -> None: + search_mock = AsyncMock(return_value={"results": []}) + close_mock = AsyncMock() + client = MagicMock() + client.search = search_mock + client.close = close_mock + + monkeypatch.setattr("drs.cli.get_client", lambda: client) + + result = runner.invoke(app, ["search", "revenue", "--next-page-token", "token-123"]) + + assert result.exit_code == 0 + search_mock.assert_awaited_once_with("revenue", filter_=None, max_results=None, next_page_token="token-123") + close_mock.assert_awaited_once() diff --git a/tests/test_client.py b/tests/test_client.py index b7d2982..4a38c6f 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -103,6 +103,47 @@ def test_content_type(self, client: DremioClient) -> None: assert client._client.headers["content-type"] == "application/json" +class TestSearch: + @pytest.mark.asyncio + async def test_search_includes_optional_filter_and_max_results(self, client: DremioClient) -> None: + captured: dict = {} + + async def _capture(request: httpx.Request) -> httpx.Response: + import json + + captured["body"] = json.loads(request.content) + return httpx.Response(200, json={"results": []}) + + client._client = httpx.AsyncClient(transport=httpx.MockTransport(_capture)) + + await client.search("revenue", filter_='category in ["JOB"]', max_results=20) + + assert captured["body"] == { + "query": "revenue", + "filter": 'category in ["JOB"]', + "maxResults": 20, + } + + @pytest.mark.asyncio + async def test_search_includes_page_token(self, client: DremioClient) -> None: + captured: dict = {} + + async def _capture(request: httpx.Request) -> httpx.Response: + import json + + captured["body"] = json.loads(request.content) + return httpx.Response(200, json={"results": []}) + + client._client = httpx.AsyncClient(transport=httpx.MockTransport(_capture)) + + await client.search("revenue", next_page_token="token-123") + + assert captured["body"] == { + "query": "revenue", + "pageToken": "token-123", + } + + class TestSQLBreadcrumb: @pytest.mark.asyncio async def test_submit_sql_prepends_breadcrumb(self, client: DremioClient) -> None: From ea9e19bb4bbaf87424a8a294b6e49f1adbbe900a Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 6 Aug 2026 10:54:28 -0400 Subject: [PATCH 2/3] fix: sort CLI test imports --- tests/test_cli.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 088d5e6..85ca6d4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -18,8 +18,7 @@ from __future__ import annotations import json -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock from typer.testing import CliRunner From ca41280f16ccd89704a1476319657259f5400831 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 6 Aug 2026 10:55:44 -0400 Subject: [PATCH 3/3] formatting --- src/drs/cli.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/drs/cli.py b/src/drs/cli.py index f1845c6..5ab21da 100644 --- a/src/drs/cli.py +++ b/src/drs/cli.py @@ -172,7 +172,9 @@ def search_command( term: str = typer.Argument(help="Search term (matches table names, view names, source names)"), filter_: str | None = typer.Option(None, "--filter", help="CEL filter expression to refine search results"), max_results: int | None = typer.Option(None, "--max-results", min=1, help="Maximum results to return per page"), - next_page_token: str | None = typer.Option(None, "--next-page-token", help="Pagination token from a prior search response"), + next_page_token: str | None = typer.Option( + None, "--next-page-token", help="Pagination token from a prior search response" + ), fmt: str = typer.Option("json", "--output", "-o", help="Output format: json, csv, pretty"), ) -> None: """Full-text search across all catalog entities (tables, views, sources)."""