feat(connectors): BI-5804 Add limit to Bitrix query params (#630)

* feat(connectors): BI-5804 Add limit to Bitrix query params * feat(connectors): BI-5804 Add limit to Bitrix query params * feat(connectors): BI-5804 Add limit to Bitrix query params * Fix mypy
datalens-tech · Oct 4, 2024 · 7b86f02 · 7b86f02
1 parent 1a26079
commit 7b86f02
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 8 deletions.
diff --git a/lib/dl_api_lib_testing/dl_api_lib_testing/data_api_base.py b/lib/dl_api_lib_testing/dl_api_lib_testing/data_api_base.py
@@ -223,7 +223,10 @@ def get_result_ordered(
         assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
         return data_resp
 
-    def get_preview(self, ds: Dataset, data_api: SyncHttpDataApiV2) -> HttpDataApiResponse:
-        data_resp = data_api.get_preview(dataset=ds)
+    def get_preview(self, ds: Dataset, data_api: SyncHttpDataApiV2, limit: Optional[int] = None) -> HttpDataApiResponse:
+        if limit:
+            data_resp = data_api.get_preview(dataset=ds, limit=limit)
+        else:
+            data_resp = data_api.get_preview(dataset=ds)
         assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
         return data_resp
diff --git a/lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/adapter.py b/lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/adapter.py
@@ -152,14 +152,20 @@ async def _run_query(self, dba_query: DBAdapterQuery) -> Any:
             json.dumps({k: (v if k != "key" else "...") for k, v in payload.json_body.items()}),
         )
 
+        request_params: dict[str, Any] = dict(
+            table=payload.table,
+            consumer="datalens",
+            request_id=request_id,
+        )
+
+        limit = self._extract_limit(dba_query.query)
+        if limit is not None:
+            request_params["limit"] = limit
+
         with self.handle_execution_error(query_text):
             resp = await self._session.post(
                 url=api_url,
-                params={
-                    "table": payload.table,
-                    "consumer": "datalens",
-                    "request_id": request_id,
-                },
+                params=request_params,
                 json=payload.json_body,
             )
 
@@ -219,6 +225,10 @@ def _extract_table_name(self, query: ClauseElement | str) -> str:
         assert isinstance(froms, sa.sql.TableClause)
         return froms.name
 
+    def _extract_limit(self, query: ClauseElement | str) -> int | None:
+        assert isinstance(query, sa.sql.Select)
+        return query._limit  # type: ignore  # TODO: "Select" has no attribute "_limit"  [attr-defined]
+
     def _parse_response_body(self, body: Any, dba_query: DBAdapterQuery) -> dict:
         assert isinstance(dba_query.query, sa.sql.Select)
         selected_columns_values = dba_query.query.selected_columns.values()

diff --git a/lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/test_data.py b/lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/test_data.py
@@ -46,7 +46,17 @@ class TestBitrixDataDistinct(BitrixDataApiTestBase, DefaultConnectorDataDistinct
 
 
 class TestBitrixDataPreview(BitrixDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
-    pass
+    def test_basic_preview(
+        self,
+        saved_dataset: Dataset,
+        data_api: SyncHttpDataApiV2,
+    ) -> None:
+        ds = saved_dataset
+
+        for limit in range(1, 10):
+            preview_resp = self.get_preview(ds, data_api, limit)
+            preview_rows = get_data_rows(preview_resp)
+            assert len(preview_rows) == limit
 
 
 class TestBitrixSmartTablesData(BitrixSmartTablesDataApiTestBase):