Add query parameter limit

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")

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 = 100,
 ) -> 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}")

ibis-server/tests/routers/ibis/test_postgres.py

@@ -202,6 +202,21 @@ 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
+
     def test_query_without_manifest(self, postgres: PostgresContainer):
         connection_info = self.to_connection_info(postgres)
         response = client.post(