diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py index b60d258b5..23c0c8135 100644 --- a/ibis-server/app/model/connector.py +++ b/ibis-server/app/model/connector.py @@ -18,11 +18,13 @@ def __init__( self.connection = self.data_source.get_connection(connection_info) self.manifest_str = manifest_str - def query(self, sql) -> pd.DataFrame: + def query(self, sql: str, limit: int) -> pd.DataFrame: rewritten_sql = rewrite(self.manifest_str, sql) - return self.connection.sql(rewritten_sql, dialect="trino").to_pandas() + return ( + self.connection.sql(rewritten_sql, dialect="trino").limit(limit).to_pandas() + ) - def dry_run(self, sql) -> None: + def dry_run(self, sql: str) -> None: try: rewritten_sql = rewrite(self.manifest_str, sql) self.connection.sql(rewritten_sql, dialect="trino") diff --git a/ibis-server/app/routers/v2/ibis.py b/ibis-server/app/routers/v2/ibis.py index 049be04eb..50e0b5882 100644 --- a/ibis-server/app/routers/v2/ibis.py +++ b/ibis-server/app/routers/v2/ibis.py @@ -27,12 +27,15 @@ def query( data_source: DataSource, dto: QueryDTO, dry_run: Annotated[bool, Query(alias="dryRun")] = False, + limit: int | None = None, ) -> Response: connector = Connector(data_source, dto.connection_info, dto.manifest_str) if dry_run: connector.dry_run(dto.sql) return Response(status_code=204) - return JSONResponse(to_json(connector.query(dto.sql), dto.column_dtypes)) + return JSONResponse( + to_json(connector.query(dto.sql, limit=limit), dto.column_dtypes) + ) @router.post("/{data_source}/validate/{rule_name}") diff --git a/ibis-server/tests/routers/ibis/test_postgres.py b/ibis-server/tests/routers/ibis/test_postgres.py index 7fb981e2f..aaca73116 100644 --- a/ibis-server/tests/routers/ibis/test_postgres.py +++ b/ibis-server/tests/routers/ibis/test_postgres.py @@ -202,6 +202,34 @@ def test_query_with_column_dtypes(self, postgres: PostgresContainer): "timestamptz": "object", } + def test_query_with_limit(self, postgres: PostgresContainer): + connection_info = self.to_connection_info(postgres) + response = client.post( + url="/v2/ibis/postgres/query", + params={"limit": 1}, + json={ + "connectionInfo": connection_info, + "manifestStr": self.manifest_str, + "sql": 'SELECT * FROM "Orders"', + }, + ) + assert response.status_code == 200 + result = response.json() + assert len(result["data"]) == 1 + + response = client.post( + url="/v2/ibis/postgres/query", + params={"limit": 1}, + json={ + "connectionInfo": connection_info, + "manifestStr": self.manifest_str, + "sql": 'SELECT * FROM "Orders" LIMIT 10', + }, + ) + assert response.status_code == 200 + result = response.json() + assert len(result["data"]) == 1 + def test_query_without_manifest(self, postgres: PostgresContainer): connection_info = self.to_connection_info(postgres) response = client.post(