rtdip · GBBBAS · Jul 25, 2024 · Jul 25, 2024 · Jul 25, 2024 · Jul 25, 2024
diff --git a/src/api/v1/common.py b/src/api/v1/common.py
@@ -188,7 +188,7 @@ def json_response(
 ) -> Response:
     schema_df = pd.DataFrame()
     if data["data"] is not None and data["data"] != "":
-        json_str = data["data"][0 : data["data"].find("}") + 1]
+        json_str = data["sample_row"]
         json_dict = json.loads(json_str, object_hook=datetime_parser)
         schema_df = pd.json_normalize(json_dict)
 

diff --git a/src/api/v1/sql.py b/src/api/v1/sql.py
@@ -55,7 +55,11 @@ def sql_get(
             else int(parameters["offset"])
         )
         data = SQLQueryBuilder().get(
-            connection, parameters["sql_statement"], limit, offset
+            connection,
+            parameters["sql_statement"],
+            parameters["to_json"],
+            limit,
+            offset,
         )
 
         return json_response(data, limit_offset_parameters)

diff --git a/src/sdk/python/rtdip_sdk/connectors/odbc/db_sql_connector.py b/src/sdk/python/rtdip_sdk/connectors/odbc/db_sql_connector.py
@@ -121,6 +121,7 @@ def fetch_all(self, fetch_size=5_000_000) -> Union[list, dict]:
             get_next_result = True
             results = None if self.return_type == ConnectionReturnType.String else []
             count = 0
+            sample_row = None
             while get_next_result:
                 result = self.cursor.fetchmany_arrow(fetch_size)
                 count += result.num_rows
@@ -133,8 +134,9 @@ def fetch_all(self, fetch_size=5_000_000) -> Union[list, dict]:
                     column_list = []
                     for column in result.columns:
                         column_list.append(column.to_pylist())
-
-                    strings = ",".join([str(item[0]) for item in zip(*column_list)])
+                    rows = [str(item[0]) for item in zip(*column_list)]
+                    sample_row = rows[0]
+                    strings = ",".join(rows)
                     if results is None:
                         results = strings
                     else:
@@ -155,6 +157,7 @@ def fetch_all(self, fetch_size=5_000_000) -> Union[list, dict]:
             elif self.return_type == ConnectionReturnType.String:
                 return {
                     "data": results,
+                    "sample_row": sample_row,
                     "count": count,
                 }
         except Exception as e:

diff --git a/src/sdk/python/rtdip_sdk/queries/sql/sql_query.py b/src/sdk/python/rtdip_sdk/queries/sql/sql_query.py
@@ -27,7 +27,7 @@ class SQLQueryBuilder:
     connection: ConnectionInterface
 
     def get(
-        self, connection=object, sql_query=str, limit=None, offset=None
+        self, connection=object, sql_query=str, to_json=False, limit=None, offset=None
     ) -> pd.DataFrame:
         """
         A function to return back raw data by querying databricks SQL Warehouse using a connection specified by the user.
@@ -49,6 +49,7 @@ def get(
         """
         try:
             parameters_dict = {"sql_statement": sql_query}
+            parameters_dict["to_json"] = to_json
             parameters_dict["supress_warning"] = True
             if limit:
                 parameters_dict["limit"] = limit

diff --git a/src/sdk/python/rtdip_sdk/queries/time_series/_time_series_query_builder.py b/src/sdk/python/rtdip_sdk/queries/time_series/_time_series_query_builder.py
@@ -101,19 +101,26 @@ def _raw_query(parameters_dict: dict) -> str:
 
 def _sql_query(parameters_dict: dict) -> str:
     sql_query = (
-        "{{ sql_statement }}"
+        "{% if to_json is defined and to_json == true %}"
+        'SELECT to_json(struct(*), map("timestampFormat", "yyyy-MM-dd\'T\'HH:mm:ss.SSSSSSSSSXXX")) as Value FROM ('
+        "{% endif %}"
+        "{{ sql_statement }} "
         "{% if limit is defined and limit is not none %}"
         "LIMIT {{ limit }} "
         "{% endif %}"
         "{% if offset is defined and offset is not none %}"
         "OFFSET {{ offset }} "
         "{% endif %}"
+        "{% if to_json is defined and to_json == true %}"
+        ")"
+        "{% endif %}"
     )
 
     sql_parameters = {
         "sql_statement": parameters_dict.get("sql_statement"),
         "limit": parameters_dict.get("limit", None),
         "offset": parameters_dict.get("offset", None),
+        "to_json": parameters_dict.get("to_json", False),
     }
 
     sql_template = Template(sql_query)
@@ -1013,10 +1020,13 @@ def _query_builder(parameters_dict: dict, query_type: str) -> str:
             + " "
             + parameters_dict["time_interval_unit"][0]
         )
+        to_json = parameters_dict.get("to_json", False)
+        parameters_dict["to_json"] = False
         sample_prepared_query, sample_query, sample_parameters = _sample_query(
             parameters_dict
         )
         sample_parameters["is_resample"] = False
+        sample_parameters["to_json"] = to_json
         return _interpolation_query(parameters_dict, sample_query, sample_parameters)
 
     if query_type == "time_weighted_average":

diff --git a/tests/api/v1/test_api_raw.py b/tests/api/v1/test_api_raw.py
@@ -14,10 +14,6 @@
 
 import pytest
 from pytest_mock import MockerFixture
-import pandas as pd
-import numpy as np
-from datetime import datetime, timezone
-from src.sdk.python.rtdip_sdk.authentication.azure import DefaultAuth
 from tests.api.v1.api_test_objects import (
     RAW_MOCKED_PARAMETER_DICT,
     RAW_MOCKED_PARAMETER_ERROR_DICT,
@@ -27,9 +23,6 @@
     TEST_HEADERS,
     BASE_URL,
 )
-from src.api.v1.models import (
-    RawResponse,
-)
 from pandas.io.json import build_table_schema
 from httpx import AsyncClient
 from src.api.v1 import app

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -99,7 +99,11 @@ def api_test_data():
     }
     mock_raw_data = test_raw_data.copy()
     mock_raw_data["EventTime"] = mock_raw_data["EventTime"].strftime(datetime_format)
-    mock_raw_df = {"data": json.dumps(mock_raw_data, separators=(",", ":")), "count": 1}
+    mock_raw_df = {
+        "data": json.dumps(mock_raw_data, separators=(",", ":")),
+        "count": 1,
+        "sample_row": json.dumps(mock_raw_data, separators=(",", ":")),
+    }
     expected_raw = expected_result(test_raw_data)
 
     # Mock Aggregated Data
@@ -110,7 +114,11 @@ def api_test_data():
     }
     mock_agg_data = test_agg_data.copy()
     mock_agg_data["EventTime"] = mock_agg_data["EventTime"].strftime(datetime_format)
-    mock_agg_df = {"data": json.dumps(mock_agg_data, separators=(",", ":")), "count": 1}
+    mock_agg_df = {
+        "data": json.dumps(mock_agg_data, separators=(",", ":")),
+        "count": 1,
+        "sample_row": json.dumps(mock_agg_data, separators=(",", ":")),
+    }
     expected_agg = expected_result(test_agg_data)
 
     # Summary Data
@@ -130,6 +138,7 @@ def api_test_data():
     mock_plot_df = {
         "data": json.dumps(mock_plot_data, separators=(",", ":")),
         "count": 1,
+        "sample_row": json.dumps(mock_plot_data, separators=(",", ":")),
     }
     expected_plot = expected_result(test_plot_data)
 
@@ -147,6 +156,7 @@ def api_test_data():
     mock_summary_df = {
         "data": json.dumps(test_summary_data, separators=(",", ":")),
         "count": 1,
+        "sample_row": json.dumps(test_summary_data, separators=(",", ":")),
     }
     expected_summary = expected_result(test_summary_data)
 
@@ -159,6 +169,7 @@ def api_test_data():
     mock_metadata_df = {
         "data": json.dumps(test_metadata, separators=(",", ":")),
         "count": 1,
+        "sample_row": json.dumps(test_metadata, separators=(",", ":")),
     }
     expected_metadata = expected_result(test_metadata)
 
@@ -183,6 +194,7 @@ def api_test_data():
     mock_latest_df = {
         "data": json.dumps(mock_latest_data, separators=(",", ":")),
         "count": 1,
+        "sample_row": json.dumps(mock_latest_data, separators=(",", ":")),
     }
     expected_latest = expected_result(test_latest_data)
 

diff --git a/tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py b/tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py
@@ -44,7 +44,7 @@ def test_sql_query(mocker: MockerFixture):
 
     mocked_cursor.assert_called_once()
     mocked_connection_close.assert_called_once()
-    mocked_execute.assert_called_once_with(mocker.ANY, query=MOCKED_SQL_QUERY)
+    mocked_execute.assert_called_once_with(mocker.ANY, query=MOCKED_SQL_QUERY + " ")
     mocked_fetch_all.assert_called_once()
     mocked_close.assert_called_once()
     assert isinstance(actual, pd.DataFrame)
@@ -78,8 +78,7 @@ def test_sql_query_fail(mocker: MockerFixture):
         # Add more test cases as needed
     ],
 )
-def test_raw_query(spark_connection, parameters, expected):
+def test_sql_query(spark_connection, parameters, expected):
     df = SQLQueryBuilder().get(spark_connection, parameters["sql_statement"])
     assert df.columns == ["EventTime", "TagName", "Status", "Value"]
-    df.show()
     assert df.count() == expected["count"]