From 6828b1a9eab28be12b0e20dcf9028b31d935d41e Mon Sep 17 00:00:00 2001 From: GBBBAS Date: Thu, 25 Jul 2024 10:27:09 +0100 Subject: [PATCH 1/3] API Updates Signed-off-by: GBBBAS --- src/api/v1/common.py | 2 +- src/api/v1/sql.py | 6 +++++- .../connectors/odbc/db_sql_connector.py | 7 +++++-- .../python/rtdip_sdk/queries/sql/sql_query.py | 3 ++- .../time_series/_time_series_query_builder.py | 12 +++++++++++- tests/api/v1/test_api_raw.py | 7 ------- tests/conftest.py | 16 ++++++++++++++-- 7 files changed, 38 insertions(+), 15 deletions(-) diff --git a/src/api/v1/common.py b/src/api/v1/common.py index 714c66e73..4aac0eaa8 100644 --- 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 index 3b56e37cd..7b9c36ceb 100644 --- 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 index aef9063f4..fde262e7e 100644 --- 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 index 1ad92d41e..ad5a1f912 100644 --- 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 index 6d687e152..5c19dd98e 100644 --- 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 index bd4e64a5c..4430e9bc7 100644 --- 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 index 7b0f7c624..4dcabf888 100644 --- 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) From adb618e360b23490404ff22f2153e31371578d02 Mon Sep 17 00:00:00 2001 From: GBBBAS Date: Thu, 25 Jul 2024 10:34:40 +0100 Subject: [PATCH 2/3] Update test for sql query Signed-off-by: GBBBAS --- tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 72b49635a..919922447 100644 --- a/tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py +++ b/tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py @@ -17,7 +17,7 @@ ACCESS_TOKEN = "mock_databricks_token" DATABRICKS_SQL_CONNECT = "databricks.sql.connect" DATABRICKS_SQL_CONNECT_CURSOR = "databricks.sql.connect.cursor" -MOCKED_SQL_QUERY = "SELECT * FROM MOCKEDTABLE" +MOCKED_SQL_QUERY = "SELECT * FROM MOCKEDTABLE " def test_sql_query(mocker: MockerFixture): From e6ee83cbb368b9f612287c411595b11429bf5648 Mon Sep 17 00:00:00 2001 From: GBBBAS Date: Thu, 25 Jul 2024 10:44:53 +0100 Subject: [PATCH 3/3] Update SQL Query Test Signed-off-by: GBBBAS --- tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 index 919922447..667950f95 100644 --- a/tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py +++ b/tests/sdk/python/rtdip_sdk/queries/sql/test_sql_query.py @@ -17,7 +17,7 @@ ACCESS_TOKEN = "mock_databricks_token" DATABRICKS_SQL_CONNECT = "databricks.sql.connect" DATABRICKS_SQL_CONNECT_CURSOR = "databricks.sql.connect.cursor" -MOCKED_SQL_QUERY = "SELECT * FROM MOCKEDTABLE " +MOCKED_SQL_QUERY = "SELECT * FROM MOCKEDTABLE" def test_sql_query(mocker: MockerFixture): @@ -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"]