From 46c7d8692ee7673db125980c5239e83cb4bde637 Mon Sep 17 00:00:00 2001 From: Rahim Bhojani Date: Mon, 13 Apr 2026 10:11:32 -0500 Subject: [PATCH 1/2] feat: support listing all reflections without a dataset argument (#10) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make `dremio reflection list` work without a dataset path — queries sys.project.reflections with no WHERE filter. Add --limit/-l flag to cap result size for large projects. Closes #10 Co-Authored-By: Claude Opus 4.6 (1M context) --- src/drs/commands/reflection.py | 34 +++++++++++++------- src/drs/introspect.py | 7 ++-- tests/test_commands/test_reflection.py | 44 ++++++++++++++++++++++++-- 3 files changed, 68 insertions(+), 17 deletions(-) diff --git a/src/drs/commands/reflection.py b/src/drs/commands/reflection.py index 2171ea2..fb5cb42 100644 --- a/src/drs/commands/reflection.py +++ b/src/drs/commands/reflection.py @@ -60,15 +60,24 @@ async def create(client: DremioClient, path: str, rtype: str, display_fields: li raise handle_api_error(exc) from exc -async def list_reflections(client: DremioClient, path: str) -> dict: - """List reflections on a dataset via sys.project.reflections.""" - parts = parse_path(path) - try: - entity = await client.get_catalog_by_path(parts) - except httpx.HTTPStatusError as exc: - raise handle_api_error(exc) from exc - dataset_id = entity["id"] - sql = f"SELECT * FROM sys.project.reflections WHERE dataset_id = '{dataset_id}'" +async def list_reflections(client: DremioClient, path: str | None = None, limit: int | None = None) -> dict: + """List reflections via sys.project.reflections. + + When *path* is given, only reflections for that dataset are returned. + When omitted, all reflections in the project are returned. + An optional *limit* caps the number of rows returned. + """ + sql = "SELECT * FROM sys.project.reflections" + if path is not None: + parts = parse_path(path) + try: + entity = await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + dataset_id = entity["id"] + sql += f" WHERE dataset_id = '{dataset_id}'" + if limit is not None: + sql += f" LIMIT {limit}" return await run_query(client, sql) @@ -142,12 +151,13 @@ def cli_create( @app.command("list") def cli_list( - path: str = typer.Argument(help="Dot-separated dataset path"), + path: str = typer.Argument(None, help="Dot-separated dataset path (omit to list all reflections)"), + limit: int = typer.Option(None, "--limit", "-l", help="Maximum number of reflections to return"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """List all reflections defined on a dataset.""" + """List reflections. Shows all project reflections, or those for a specific dataset.""" client = _get_client() - _run_command(list_reflections(client, path), client, fmt) + _run_command(list_reflections(client, path, limit=limit), client, fmt) @app.command("get") diff --git a/src/drs/introspect.py b/src/drs/introspect.py index d68bc3c..b70c91c 100644 --- a/src/drs/introspect.py +++ b/src/drs/introspect.py @@ -302,11 +302,12 @@ "reflection.list": { "group": "reflection", "command": "list", - "description": "List all reflections defined on a dataset.", + "description": "List reflections. Shows all project reflections, or those for a specific dataset.", "mechanism": "SQL", - "sql_template": "SELECT * FROM sys.project.reflections WHERE dataset_id = '{dataset_id}'", + "sql_template": "SELECT * FROM sys.project.reflections [WHERE dataset_id = '{dataset_id}'] [LIMIT {limit}]", "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "path", "type": "string", "required": False, "positional": True}, + {"name": "limit", "type": "integer", "required": False, "flag": "--limit/-l"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, diff --git a/tests/test_commands/test_reflection.py b/tests/test_commands/test_reflection.py index a6d657a..4d6c41a 100644 --- a/tests/test_commands/test_reflection.py +++ b/tests/test_commands/test_reflection.py @@ -17,11 +17,51 @@ from __future__ import annotations -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import pytest -from drs.commands.reflection import delete, get_reflection, refresh +from drs.commands.reflection import delete, get_reflection, list_reflections, refresh + +QUERY_RESULT = {"job_id": "j1", "state": "COMPLETED", "rowCount": 2, "rows": [{"id": "r1"}, {"id": "r2"}]} + + +@pytest.mark.asyncio +async def test_list_reflections_all(mock_client) -> None: + """Omitting path queries all reflections without a WHERE clause.""" + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + result = await list_reflections(mock_client) + mock_rq.assert_called_once_with(mock_client, "SELECT * FROM sys.project.reflections") + assert result["rowCount"] == 2 + + +@pytest.mark.asyncio +async def test_list_reflections_for_dataset(mock_client) -> None: + """Providing a path filters by dataset_id.""" + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "ds-123"}) + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + result = await list_reflections(mock_client, path="space.my_table") + mock_rq.assert_called_once_with(mock_client, "SELECT * FROM sys.project.reflections WHERE dataset_id = 'ds-123'") + assert result["rowCount"] == 2 + + +@pytest.mark.asyncio +async def test_list_reflections_with_limit(mock_client) -> None: + """--limit appends a SQL LIMIT clause.""" + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + await list_reflections(mock_client, limit=50) + mock_rq.assert_called_once_with(mock_client, "SELECT * FROM sys.project.reflections LIMIT 50") + + +@pytest.mark.asyncio +async def test_list_reflections_dataset_with_limit(mock_client) -> None: + """Both path and limit combine WHERE and LIMIT.""" + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "ds-456"}) + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + await list_reflections(mock_client, path="space.ds", limit=10) + mock_rq.assert_called_once_with( + mock_client, "SELECT * FROM sys.project.reflections WHERE dataset_id = 'ds-456' LIMIT 10" + ) @pytest.mark.asyncio From c2f08ff331fe4ad8f1e997adb0c990ba07a6a565 Mon Sep 17 00:00:00 2001 From: Rahim Bhojani Date: Mon, 13 Apr 2026 11:37:35 -0500 Subject: [PATCH 2/2] feat: add --type, --status, --dataset-name filters to reflection list MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback — support filtering reflections by type (raw/aggregation), status (CAN_ACCELERATE, FAILED, etc.), and dataset name (substring match via ILIKE). All filters compose with each other and with the existing path argument and --limit. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/drs/commands/reflection.py | 31 ++++++++++++++++++-- src/drs/introspect.py | 5 +++- tests/test_commands/test_reflection.py | 39 ++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 4 deletions(-) diff --git a/src/drs/commands/reflection.py b/src/drs/commands/reflection.py index fb5cb42..79b0ead 100644 --- a/src/drs/commands/reflection.py +++ b/src/drs/commands/reflection.py @@ -60,14 +60,24 @@ async def create(client: DremioClient, path: str, rtype: str, display_fields: li raise handle_api_error(exc) from exc -async def list_reflections(client: DremioClient, path: str | None = None, limit: int | None = None) -> dict: +async def list_reflections( + client: DremioClient, + path: str | None = None, + *, + rtype: str | None = None, + status: str | None = None, + dataset_name: str | None = None, + limit: int | None = None, +) -> dict: """List reflections via sys.project.reflections. When *path* is given, only reflections for that dataset are returned. When omitted, all reflections in the project are returned. + Optional filters narrow results by *rtype*, *status*, or *dataset_name*. An optional *limit* caps the number of rows returned. """ sql = "SELECT * FROM sys.project.reflections" + conditions: list[str] = [] if path is not None: parts = parse_path(path) try: @@ -75,7 +85,15 @@ async def list_reflections(client: DremioClient, path: str | None = None, limit: except httpx.HTTPStatusError as exc: raise handle_api_error(exc) from exc dataset_id = entity["id"] - sql += f" WHERE dataset_id = '{dataset_id}'" + conditions.append(f"dataset_id = '{dataset_id}'") + if rtype is not None: + conditions.append(f"type = '{rtype.upper()}'") + if status is not None: + conditions.append(f"status = '{status.upper()}'") + if dataset_name is not None: + conditions.append(f"dataset_name ILIKE '%{dataset_name}%'") + if conditions: + sql += " WHERE " + " AND ".join(conditions) if limit is not None: sql += f" LIMIT {limit}" return await run_query(client, sql) @@ -152,12 +170,19 @@ def cli_create( @app.command("list") def cli_list( path: str = typer.Argument(None, help="Dot-separated dataset path (omit to list all reflections)"), + rtype: str = typer.Option(None, "--type", "-t", help="Filter by reflection type: raw or aggregation"), + status: str = typer.Option(None, "--status", "-s", help="Filter by status (e.g. CAN_ACCELERATE, FAILED, EXPIRED)"), + dataset_name: str = typer.Option(None, "--dataset-name", "-d", help="Filter by dataset name (substring match)"), limit: int = typer.Option(None, "--limit", "-l", help="Maximum number of reflections to return"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: """List reflections. Shows all project reflections, or those for a specific dataset.""" client = _get_client() - _run_command(list_reflections(client, path, limit=limit), client, fmt) + _run_command( + list_reflections(client, path, rtype=rtype, status=status, dataset_name=dataset_name, limit=limit), + client, + fmt, + ) @app.command("get") diff --git a/src/drs/introspect.py b/src/drs/introspect.py index b70c91c..b317e0a 100644 --- a/src/drs/introspect.py +++ b/src/drs/introspect.py @@ -304,9 +304,12 @@ "command": "list", "description": "List reflections. Shows all project reflections, or those for a specific dataset.", "mechanism": "SQL", - "sql_template": "SELECT * FROM sys.project.reflections [WHERE dataset_id = '{dataset_id}'] [LIMIT {limit}]", + "sql_template": "SELECT * FROM sys.project.reflections [WHERE ...] [LIMIT {limit}]", "parameters": [ {"name": "path", "type": "string", "required": False, "positional": True}, + {"name": "type", "type": "string", "required": False, "flag": "--type/-t"}, + {"name": "status", "type": "string", "required": False, "flag": "--status/-s"}, + {"name": "dataset_name", "type": "string", "required": False, "flag": "--dataset-name/-d"}, {"name": "limit", "type": "integer", "required": False, "flag": "--limit/-l"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], diff --git a/tests/test_commands/test_reflection.py b/tests/test_commands/test_reflection.py index 4d6c41a..b2ca0de 100644 --- a/tests/test_commands/test_reflection.py +++ b/tests/test_commands/test_reflection.py @@ -64,6 +64,45 @@ async def test_list_reflections_dataset_with_limit(mock_client) -> None: ) +@pytest.mark.asyncio +async def test_list_reflections_filter_by_type(mock_client) -> None: + """--type filters by reflection type.""" + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + await list_reflections(mock_client, rtype="raw") + mock_rq.assert_called_once_with(mock_client, "SELECT * FROM sys.project.reflections WHERE type = 'RAW'") + + +@pytest.mark.asyncio +async def test_list_reflections_filter_by_status(mock_client) -> None: + """--status filters by reflection status.""" + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + await list_reflections(mock_client, status="failed") + mock_rq.assert_called_once_with(mock_client, "SELECT * FROM sys.project.reflections WHERE status = 'FAILED'") + + +@pytest.mark.asyncio +async def test_list_reflections_filter_by_dataset_name(mock_client) -> None: + """--dataset-name filters with ILIKE substring match.""" + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + await list_reflections(mock_client, dataset_name="orders") + mock_rq.assert_called_once_with( + mock_client, "SELECT * FROM sys.project.reflections WHERE dataset_name ILIKE '%orders%'" + ) + + +@pytest.mark.asyncio +async def test_list_reflections_combined_filters(mock_client) -> None: + """Multiple filters combine with AND.""" + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "ds-789"}) + with patch("drs.commands.reflection.run_query", new_callable=AsyncMock, return_value=QUERY_RESULT) as mock_rq: + await list_reflections(mock_client, path="space.ds", rtype="raw", status="can_accelerate", limit=5) + mock_rq.assert_called_once_with( + mock_client, + "SELECT * FROM sys.project.reflections" + " WHERE dataset_id = 'ds-789' AND type = 'RAW' AND status = 'CAN_ACCELERATE' LIMIT 5", + ) + + @pytest.mark.asyncio async def test_get_reflection(mock_client) -> None: mock_client.get_reflection = AsyncMock(return_value={"id": "r1", "status": "CAN_ACCELERATE"})