fix: BI-5902 typed query raw fixes (#809)

typed query raw fixes
datalens-tech · Jan 24, 2025 · bb6a5d5 · bb6a5d5
1 parent 989ac47
commit bb6a5d5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/lib/dl_api_lib/dl_api_lib/app/data_api/resources/typed_query_raw.py b/lib/dl_api_lib/dl_api_lib/app/data_api/resources/typed_query_raw.py
@@ -84,10 +84,7 @@ async def execute_query(self, connection: ConnectionBase, typed_query: TypedQuer
 
     def make_response_data(self, typed_query_result: TypedQueryRawResult) -> dict:
         """Serialize output"""
-        return {
-            "query_type": typed_query_result.query_type.name,
-            "data": TypedQueryRawResultSchema().dump(typed_query_result),
-        }
+        return TypedQueryRawResultSchema().dump(typed_query_result)
 
     @generic_profiler_async("dashsql-typed-query-raw")
     @requires(RequiredResourceDSAPI.JSON_REQUEST)

diff --git a/lib/dl_api_lib/dl_api_lib/schemas/typed_query_raw.py b/lib/dl_api_lib/dl_api_lib/schemas/typed_query_raw.py
@@ -39,7 +39,11 @@ class TypedQueryRawSchema(DefaultSchema[RawTypedQueryRaw]):
     parameters = ma_fields.Nested(TypedQueryRawParametersSchema, required=True)
 
 
-class TypedQueryRawResultSchema(BaseSchema):
-    json = ma_fields.Raw(load_default=None)
-    headers = ma_fields.Dict()  # TODO: required? or load_default = None?
+class TypedQueryRawResultDataSchema(BaseSchema):
+    body = ma_fields.Dict(load_default=None)
+    headers = ma_fields.Dict()
     status = ma_fields.Integer(required=True)
+
+
+class TypedQueryRawResultSchema(BaseSchema):
+    data = ma_fields.Nested(TypedQueryRawResultDataSchema)
diff --git a/lib/dl_core/dl_core/data_processing/typed_query_raw.py b/lib/dl_core/dl_core/data_processing/typed_query_raw.py
@@ -15,4 +15,4 @@ class CEBasedTypedQueryRawProcessor(TypedQueryRawProcessorBase):
     async_conn_executor: AsyncConnExecutorBase = attr.ib(kw_only=True)
 
     async def process_typed_query_raw(self, typed_query_raw: TypedQueryRaw) -> TypedQueryRawResult:
-        return await self.async_conn_executor.execute_typed_query_raw(typed_query=typed_query_raw)
+        return await self.async_conn_executor.execute_typed_query_raw(typed_query_raw=typed_query_raw)