From f769833c718c5eb8aea6165c21ab31077ff992a1 Mon Sep 17 00:00:00 2001 From: Amber-Rigg <44523299+Amber-Rigg@users.noreply.github.com> Date: Thu, 11 Jan 2024 15:01:31 +0000 Subject: [PATCH] Weather Query Improvement with Forecast Runtime Query (#621) * Simplify Table Name Signed-off-by: Amber-Rigg * RAW query on Forecast Run Time Signed-off-by: Amber-Rigg * Update Tests Signed-off-by: Amber-Rigg * Update tests Signed-off-by: Amber-Rigg * Update Tests Signed-off-by: Amber-Rigg * Update Tests Signed-off-by: Amber-Rigg * Refactor raw_parameters build Signed-off-by: Amber-Rigg * Update Query Builder Table Connection Signed-off-by: Amber-Rigg * Review to include options and refactor test Signed-off-by: Amber-Rigg * black formatting Signed-off-by: Amber-Rigg * Remove 'source' parameter from WeatherQueryBuilder methods Signed-off-by: Amber-Rigg * Examples Added to Documentation Signed-off-by: Amber-Rigg * Remove duplication in tests Signed-off-by: Amber-Rigg * Remove unused Variables and include timestamp coulmn name default Signed-off-by: Amber-Rigg * Documentation and Examples Formatting Signed-off-by: Amber-Rigg * Documentation Update Signed-off-by: Amber-Rigg --------- Signed-off-by: Amber-Rigg --- .../query/functions/weather/latest.md | 61 ++++- .../query/functions/weather/raw.md | 73 +++++- .../weather/weather_query_builder.md | 8 +- .../queries/weather/_weather_query_builder.py | 130 +++++------ .../rtdip_sdk/queries/weather/latest.py | 10 +- .../python/rtdip_sdk/queries/weather/raw.py | 18 +- .../queries/weather/weather_query_builder.py | 220 ++++++++++++------ .../rtdip_sdk/queries/weather/test_latest.py | 112 +++++++-- .../rtdip_sdk/queries/weather/test_raw.py | 136 ++++++++--- .../weather/test_weather_query_builder.py | 20 +- 10 files changed, 562 insertions(+), 226 deletions(-) diff --git a/docs/sdk/code-reference/query/functions/weather/latest.md b/docs/sdk/code-reference/query/functions/weather/latest.md index 6e7d33f49..01532d537 100644 --- a/docs/sdk/code-reference/query/functions/weather/latest.md +++ b/docs/sdk/code-reference/query/functions/weather/latest.md @@ -1,2 +1,61 @@ # Weather Latest Function -::: src.sdk.python.rtdip_sdk.queries.weather.latest \ No newline at end of file +::: src.sdk.python.rtdip_sdk.queries.weather.latest + +## Example get_point + +```python +from rtdip_sdk.authentication.azure import DefaultAuth +from rtdip_sdk.queries.weather.latest import get_point +from rtdip_sdk.connectors import DatabricksSQLConnection + +auth = DefaultAuth().authenticate() +token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token +connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + +params = { + "forecast": "mock_forecast", + "forecast_type": "mock_weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "lat": 1.1, + "lon": 1.1, +} + +x = get_point(connection, params) + +print(x) +``` + +## Example get_grid + +```python +from rtdip_sdk.authentication.azure import DefaultAuth +from rtdip_sdk.queries.weather.latest import get_point +from rtdip_sdk.connectors import DatabricksSQLConnection + +auth = DefaultAuth().authenticate() +token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token +connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + +params = { + "forecast": "mock_forecast", + "forecast_type": "mock_weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "min_lat": 36, + "max_lat": 38, + "min_lon": -109.1, + "max_lon": -107.1, +} + +x = get_grid(connection, params) + +print(x) +``` + +These examples are using [```DefaultAuth()```](../../../authentication/azure.md) and [```DatabricksSQLConnection()```](../../connectors/db-sql-connector.md) to authenticate and connect. You can find other ways to authenticate [here](../../../authentication/azure.md). The alternative built in connection methods are either by [```PYODBCSQLConnection()```](../../connectors/pyodbc-sql-connector.md), [```TURBODBCSQLConnection()```](../../connectors/turbodbc-sql-connector.md) or [```SparkConnection()```](../../connectors/spark-connector.md). + +!!! note "Note" + ```server_hostname``` and ```http_path``` can be found on the [SQL Warehouses Page](../../../../queries/databricks/sql-warehouses.md).
\ No newline at end of file diff --git a/docs/sdk/code-reference/query/functions/weather/raw.md b/docs/sdk/code-reference/query/functions/weather/raw.md index 4e209aff0..bdbaa0f73 100644 --- a/docs/sdk/code-reference/query/functions/weather/raw.md +++ b/docs/sdk/code-reference/query/functions/weather/raw.md @@ -1,2 +1,73 @@ # Weather Raw Function -::: src.sdk.python.rtdip_sdk.queries.weather.raw \ No newline at end of file +::: src.sdk.python.rtdip_sdk.queries.weather.raw + +## Example get_point + +```python +from rtdip_sdk.authentication.azure import DefaultAuth +from rtdip_sdk.queries.weather.raw import get_point +from rtdip_sdk.connectors import DatabricksSQLConnection + +auth = DefaultAuth().authenticate() +token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token +connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + +params = { + "forecast": "mock_forecast", + "forecast_type": "mock_weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "lat": 1.1, + "lon": 1.1, + "start_date": "2020-01-01", + "end_date": "2020-01-02", + "forecast_run_start_date": "2020-01-01", + "forecast_run_end_date": "2020-01-02", + "timestamp_column": "EventTime", + "forecast_run_timestamp_column": "EnqueuedTime", +} + +x = get_point(connection, params) + +print(x) +``` + +## Example get_grid + +```python +from rtdip_sdk.authentication.azure import DefaultAuth +from rtdip_sdk.queries.weather.raw import get_grid +from rtdip_sdk.connectors import DatabricksSQLConnection + +auth = DefaultAuth().authenticate() +token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token +connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + +params = { + "forecast": "mock_forecast", + "forecast_type": "mock_weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "min_lat": 36, + "max_lat": 38, + "min_lon": -109.1, + "max_lon": -107.1, + "start_date": "2020-01-01", + "end_date": "2020-01-02", + "forecast_run_start_date": "2020-01-01", + "forecast_run_end_date": "2020-01-02", + "timestamp_column": "EventTime", + "forecast_run_timestamp_column": "EnqueuedTime", +} + +x = get_grid(connection, params) + +print(x) +``` + +These examples are using [```DefaultAuth()```](../../../authentication/azure.md) and [```DatabricksSQLConnection()```](../../connectors/db-sql-connector.md) to authenticate and connect. You can find other ways to authenticate [here](../../../authentication/azure.md). The alternative built in connection methods are either by [```PYODBCSQLConnection()```](../../connectors/pyodbc-sql-connector.md), [```TURBODBCSQLConnection()```](../../connectors/turbodbc-sql-connector.md) or [```SparkConnection()```](../../connectors/spark-connector.md). + +!!! note "Note" + ```server_hostname``` and ```http_path``` can be found on the [SQL Warehouses Page](../../../../queries/databricks/sql-warehouses.md).
\ No newline at end of file diff --git a/docs/sdk/code-reference/query/functions/weather/weather_query_builder.md b/docs/sdk/code-reference/query/functions/weather/weather_query_builder.md index 369267444..917bf15b9 100644 --- a/docs/sdk/code-reference/query/functions/weather/weather_query_builder.md +++ b/docs/sdk/code-reference/query/functions/weather/weather_query_builder.md @@ -1,2 +1,8 @@ # Weather Query Builder -::: src.sdk.python.rtdip_sdk.queries.weather.weather_query_builder \ No newline at end of file +::: src.sdk.python.rtdip_sdk.queries.weather.weather_query_builder + +!!! note "Note" + These examples are using [```DefaultAuth()```](../../../authentication/azure.md) and [```DatabricksSQLConnection()```](../../connectors/db-sql-connector.md) to authenticate and connect. You can find other ways to authenticate [here](../../../authentication/azure.md). The alternative built in connection methods are either by [```PYODBCSQLConnection()```](../../connectors/pyodbc-sql-connector.md), [```TURBODBCSQLConnection()```](../../connectors/turbodbc-sql-connector.md) or [```SparkConnection()```](../../connectors/spark-connector.md).
+ +!!! note "Note" + ```server_hostname``` and ```http_path``` can be found on the [SQL Warehouses Page](../../../../queries/databricks/sql-warehouses.md).
\ No newline at end of file diff --git a/src/sdk/python/rtdip_sdk/queries/weather/_weather_query_builder.py b/src/sdk/python/rtdip_sdk/queries/weather/_weather_query_builder.py index 495baf099..d7dfea5fc 100644 --- a/src/sdk/python/rtdip_sdk/queries/weather/_weather_query_builder.py +++ b/src/sdk/python/rtdip_sdk/queries/weather/_weather_query_builder.py @@ -32,51 +32,45 @@ def _build_parameters( area_type: str, table_type: str, ) -> dict: - if area_type == "grid": - raw_parameters = { - "forecast": parameters_dict.get("forecast", None), - "region": parameters_dict.get("region"), - "data_security_level": parameters_dict.get("data_security_level"), - "data_type": parameters_dict.get("data_type"), - "start_date": parameters_dict["start_date"], - "end_date": parameters_dict["end_date"], - "max_lat": parameters_dict["max_lat"], - "max_lon": parameters_dict["max_lon"], - "min_lat": parameters_dict["min_lat"], - "min_lon": parameters_dict["min_lon"], - "source": parameters_dict.get("source", None), - "limit": parameters_dict.get("limit", None), - "latitude_column": parameters_dict.get("latitude_column", "Latitude"), - "longitude_column": parameters_dict.get("longitude_column", "Longitude"), - "tagname_column": parameters_dict.get("tagname_column", "TagName"), - } - if table_type == "raw": - raw_parameters["timestamp_column"] = parameters_dict.get( - "timestamp_column", "EventTime" - ) - raw_parameters["include_status"] = False + raw_parameters = { + "source": parameters_dict.get("source", None), + "forecast": parameters_dict.get("forecast", None), + "forecast_type": parameters_dict.get("forecast_type", "weather"), + "region": parameters_dict.get("region", None), + "data_security_level": parameters_dict.get("data_security_level", None), + "data_type": parameters_dict.get("data_type", None), + "limit": parameters_dict.get("limit", None), + "latitude_column": parameters_dict.get("latitude_column", "Latitude"), + "longitude_column": parameters_dict.get("longitude_column", "Longitude"), + "tagname_column": parameters_dict.get("tagname_column", "TagName"), + } if area_type == "point": - raw_parameters = { - "forecast": parameters_dict.get("forecast", None), - "region": parameters_dict.get("region"), - "data_security_level": parameters_dict.get("data_security_level"), - "data_type": parameters_dict.get("data_type"), - "start_date": parameters_dict["start_date"], - "end_date": parameters_dict["end_date"], - "lat": parameters_dict["lat"], - "lon": parameters_dict["lon"], - "source": parameters_dict.get("source", None), - "limit": parameters_dict.get("limit", None), - "latitude_column": parameters_dict.get("latitude_column", "Latitude"), - "longitude_column": parameters_dict.get("longitude_column", "Longitude"), - "tagname_column": parameters_dict.get("tagname_column", "TagName"), - } - if table_type == "raw": - raw_parameters["timestamp_column"] = parameters_dict.get( - "timestamp_column", "EventTime" - ) - raw_parameters["include_status"] = False + raw_parameters["lat"] = parameters_dict["lat"] + raw_parameters["lon"] = parameters_dict["lon"] + + if area_type == "grid": + raw_parameters["max_lat"] = parameters_dict["max_lat"] + raw_parameters["max_lon"] = parameters_dict["max_lon"] + raw_parameters["min_lat"] = parameters_dict["min_lat"] + raw_parameters["min_lon"] = parameters_dict["min_lon"] + + if table_type == "raw": + raw_parameters["start_date"] = parameters_dict["start_date"] + raw_parameters["end_date"] = parameters_dict["end_date"] + raw_parameters["forecast_run_start_date"] = parameters_dict[ + "forecast_run_start_date" + ] + raw_parameters["forecast_run_end_date"] = parameters_dict[ + "forecast_run_end_date" + ] + raw_parameters["timestamp_column"] = parameters_dict.get( + "timestamp_column", "EventTime" + ) + raw_parameters["forecast_run_timestamp_column"] = parameters_dict.get( + "forecast_run_timestamp_column", "EnqueuedTime" + ) + raw_parameters["include_status"] = False return raw_parameters @@ -87,16 +81,14 @@ def _raw_query_grid(parameters_dict: dict) -> str: "{% if source is defined and source is not none %}" "`{{ source|lower }}` " "{% else %}" - "`{{ forecast|lower }}`.`weather`.`{{ region|lower }}_weather_{{ data_security_level|lower }}_events_{{ data_type|lower }}` " - "{% endif %}" - 'WHERE `{{ timestamp_column }}` BETWEEN to_timestamp("{{ start_date }}") AND to_timestamp("{{ end_date }}")' - "AND `{{ latitude_column }}` > '{{ min_lat}}' " - "AND `{{ latitude_column }}` < '{{ max_lat}}' " - "AND `{{ longitude_column }}` > '{{ min_lon}}' " - "AND`{{ longitude_column }}` < '{{ max_lon}}' " - "{% if source is defined and source is not none %}" - "AND SOURCE = '{{ source }}' " + "`{{ forecast|lower }}`.`{{ forecast_type|lower }}`.`{{ region|lower }}_{{ data_security_level|lower }}_events_{{ data_type|lower }}` " "{% endif %}" + 'WHERE (`{{ timestamp_column }}` BETWEEN to_timestamp("{{ start_date }}") AND to_timestamp("{{ end_date }}")) ' + 'AND (`{{ forecast_run_timestamp_column }}` BETWEEN to_timestamp("{{ forecast_run_start_date }}") AND to_timestamp("{{ forecast_run_end_date }}")) ' + "AND `{{ latitude_column }}` > {{ min_lat}} " + "AND `{{ latitude_column }}` < {{ max_lat}} " + "AND `{{ longitude_column }}` > {{ min_lon}} " + "AND `{{ longitude_column }}` < {{ max_lon}} " "ORDER BY `{{ tagname_column }}` " "{% if limit is defined and limit is not none %}" "LIMIT {{ limit }} " @@ -115,14 +107,12 @@ def _raw_query_point(parameters_dict: dict) -> str: "{% if source is defined and source is not none %}" "`{{ source|lower }}` " "{% else %}" - "`{{ forecast|lower }}`.`weather`.`{{ region|lower }}_weather_{{ data_security_level|lower }}_events_{{ data_type|lower }}` " - "{% endif %}" - 'WHERE `{{ timestamp_column }}` BETWEEN to_timestamp("{{ start_date }}") AND to_timestamp("{{ end_date }}")' - "AND `{{ latitude_column }}` > '{{lat}}' " - "AND `{{ longitude_column }}` > '{{lon}}' " - "{% if source is defined and source is not none %}" - "AND SOURCE = '{{ source }}' " + "`{{ forecast|lower }}`.`{{ forecast_type|lower }}`.`{{ region|lower }}_{{ data_security_level|lower }}_events_{{ data_type|lower }}` " "{% endif %}" + 'WHERE (`{{ timestamp_column }}` BETWEEN to_timestamp("{{ start_date }}") AND to_timestamp("{{ end_date }}")) ' + 'AND (`{{ forecast_run_timestamp_column }}` BETWEEN to_timestamp("{{ forecast_run_start_date }}") AND to_timestamp("{{ forecast_run_end_date }}")) ' + "AND `{{ latitude_column }}` == {{ lat }} " + "AND `{{ longitude_column }}` == {{ lon }} " "ORDER BY `{{ tagname_column }}` " "{% if limit is defined and limit is not none %}" "LIMIT {{ limit }} " @@ -141,15 +131,12 @@ def _latest_query_grid(parameters_dict: dict) -> str: "{% if source is defined and source is not none %}" "`{{ source|lower }}` " "{% else %}" - "`{{ forecast|lower }}`.`weather`.`{{ region|lower }}_weather_{{ data_security_level|lower }}_events_{{ data_type|lower }}` " - "{% endif %}" - "WHERE `{{ latitude_column }}` > '{{ min_lat}}' " - "AND `{{ latitude_column }}` < '{{ max_lat}}' " - "AND `{{ longitude_column }}` > '{{ min_lon}}' " - "AND`{{ longitude_column }}` < '{{ max_lon}}' " - "{% if source is defined and source is not none %}" - "AND SOURCE = '{{ source }}' " + "`{{ forecast|lower }}`.`{{ forecast_type|lower }}`.`{{ region|lower }}_{{ data_security_level|lower }}_events_{{ data_type|lower }}_latest` " "{% endif %}" + "WHERE `{{ latitude_column }}` > {{ min_lat}} " + "AND `{{ latitude_column }}` < {{ max_lat}} " + "AND `{{ longitude_column }}` > {{ min_lon}} " + "AND `{{ longitude_column }}` < {{ max_lon}} " "ORDER BY `{{ tagname_column }}` " "{% if limit is defined and limit is not none %}" "LIMIT {{ limit }} " @@ -168,13 +155,10 @@ def _latest_query_point(parameters_dict: dict) -> str: "{% if source is defined and source is not none %}" "`{{ source|lower }}` " "{% else %}" - "`{{ forecast|lower }}`.`weather`.`{{ region|lower }}_weather_{{ data_security_level|lower }}_events_{{ data_type|lower }}` " - "{% endif %}" - "WHERE `{{ latitude_column }}` == '{{lat}}' " - "AND `{{ longitude_column }}` == '{{lon}}' " - "{% if source is defined and source is not none %}" - "AND SOURCE = '{{ source }}' " + "`{{ forecast|lower }}`.`{{ forecast_type|lower }}`.`{{ region|lower }}_{{ data_security_level|lower }}_events_{{ data_type|lower }}_latest` " "{% endif %}" + "WHERE `{{ latitude_column }}` == {{ lat }} " + "AND `{{ longitude_column }}` == {{ lon }} " "ORDER BY `{{ tagname_column }}` " "{% if limit is defined and limit is not none %}" "LIMIT {{ limit }} " diff --git a/src/sdk/python/rtdip_sdk/queries/weather/latest.py b/src/sdk/python/rtdip_sdk/queries/weather/latest.py index 611760b43..124e0f98a 100644 --- a/src/sdk/python/rtdip_sdk/queries/weather/latest.py +++ b/src/sdk/python/rtdip_sdk/queries/weather/latest.py @@ -34,7 +34,9 @@ def get_grid(connection: object, parameters_dict: dict) -> pd.DataFrame: parameters_dict: A dictionary of parameters (see Attributes table below) Attributes: - forecast (str): Business unit + source (optional str): Source of the data the full table name + forecast (str): Any specific identifier for forecast + forecast_type(str): Type of forecast ie weather, solar, power, etc region (str): Region data_security_level (str): Level of data security data_type (str): Type of the data (float, integer, double, string) @@ -42,7 +44,6 @@ def get_grid(connection: object, parameters_dict: dict) -> pd.DataFrame: max_lon (float): Maximum longitude min_lat (float): Minimum latitude min_lon (float): Minimum longitude - source (optional str): Source of the data ie ECMWF limit (optional int): The number of rows to be returned Returns: @@ -84,13 +85,14 @@ def get_point(connection: object, parameters_dict: dict) -> pd.DataFrame: parameters_dict: A dictionary of parameters (see Attributes table below) Attributes: - forecast (str): Business unit + source (optional str): Source of the data the full table name + forecast (str): Any specific identifier for forecast + forecast_type(str): Type of forecast ie weather, solar, power, etc region (str): Region data_security_level (str): Level of data security data_type (str): Type of the data (float, integer, double, string) lat (float): latitude lon (float): longitude - source (optional str): Source of the data ie ECMWF limit (optional int): The number of rows to be returned Returns: diff --git a/src/sdk/python/rtdip_sdk/queries/weather/raw.py b/src/sdk/python/rtdip_sdk/queries/weather/raw.py index 9941f2789..48f4af5a4 100644 --- a/src/sdk/python/rtdip_sdk/queries/weather/raw.py +++ b/src/sdk/python/rtdip_sdk/queries/weather/raw.py @@ -34,17 +34,22 @@ def get_grid(connection: object, parameters_dict: dict) -> pd.DataFrame: parameters_dict: A dictionary of parameters (see Attributes table below) Attributes: - forecast (str): Business unit + source (optional str): Source of the data the full table name + forecast (str): Any specific identifier for forecast + forecast_type(str): Type of forecast ie weather, solar, power, etc region (str): Region data_security_level (str): Level of data security data_type (str): Type of the data (float, integer, double, string) start_date (str): Start date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) end_date (str): End date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_start_date (str): Start date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_end_date (str): End date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + timestamp_column (str): The column which contains the the forecast output time. Default "EventTime". + forecast_run_timestamp_column (str): The column which contains whent the forecast was run. Default "EnqueuedTime". max_lat (float): Maximum latitude max_lon (float): Maximum longitude min_lat (float): Minimum latitude min_lon (float): Minimum longitude - source (optional str): Source of the data ie ECMWF limit (optional int): The number of rows to be returned } @@ -89,15 +94,20 @@ def get_point(connection: object, parameters_dict: dict) -> pd.DataFrame: parameters_dict: A dictionary of parameters (see Attributes table below) Attributes: - forecast (str): Business unit + source (optional str): Source of the data the full table name + forecast (str): Any specific identifier for forecast + forecast_type(str): Type of forecast ie weather, solar, power, etc region (str): Region data_security_level (str): Level of data security data_type (str): Type of the data (float, integer, double, string) start_date (str): Start date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) end_date (str): End date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_start_date (str): Start date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_end_date (str): End date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + timestamp_column (str): The column which contains the the forecast output time. Default "EventTime". + forecast_run_timestamp_column (str): The column which contains whent the forecast was run. Default "EnqueuedTime. lat (float): latitude lon (float): longitude - source (optional str): Source of the data ie ECMWF limit (optional int): The number of rows to be returned } diff --git a/src/sdk/python/rtdip_sdk/queries/weather/weather_query_builder.py b/src/sdk/python/rtdip_sdk/queries/weather/weather_query_builder.py index eb78fde10..f6be94175 100644 --- a/src/sdk/python/rtdip_sdk/queries/weather/weather_query_builder.py +++ b/src/sdk/python/rtdip_sdk/queries/weather/weather_query_builder.py @@ -23,11 +23,18 @@ class WeatherQueryBuilder: """ - A builder for developing RTDIP queries using any delta table + A builder for developing RTDIP forecast queries using any delta table + """ parameters: dict connection: ConnectionInterface + close_connection: bool + data_source: str + tagname_column: str + timestamp_column: str + status_column: str + value_column: str def connect(self, connection: ConnectionInterface): """ @@ -44,6 +51,7 @@ def source( source: str, tagname_column: str = "TagName", timestamp_column: str = "EventTime", + forecast_run_timestamp_column: str = "EnqueuedTime", status_column: Union[str, None] = "Status", value_column: str = "Value", ): @@ -54,63 +62,82 @@ def source( source (str): Source of the query can be a Unity Catalog table, Hive metastore table or path tagname_column (optional str): The column name in the source that contains the tagnames or series timestamp_column (optional str): The timestamp column name in the source + forecast_run_timestamp_column (optional str): The forecast run timestamp column name in the source status_column (optional str): The status column name in the source indicating `Good` or `Bad`. If this is not available, specify `None` value_column (optional str): The value column name in the source which is normally a float or string value for the time series event """ self.data_source = "`.`".join(source.split(".")) self.tagname_column = tagname_column self.timestamp_column = timestamp_column + self.forecast_run_timestamp_column = forecast_run_timestamp_column self.status_column = status_column self.value_column = value_column return self def raw_point( self, - forecast: str, - region: str, - data_security_level: str, - data_type: str, start_date: str, end_date: str, + forecast_run_start_date: str, + forecast_run_end_date: str, lat: float, lon: float, - source: str = None, - time_zone: str = None, - include_bad_data: bool = False, limit: int = None, ) -> DataFrame: """ A function to return back raw data for a point. + **Example:** + ```python + from rtdip_sdk.queries.weather.weather_query_builder import ( + WeatherQueryBuilder, + ) + from rtdip_sdk.authentication.azure import DefaultAuth + from rtdip_sdk.connectors import DatabricksSQLConnection + + auth = DefaultAuth().authenticate() + token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token + connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + + data = ( + WeatherQueryBuilder() + .connect(connection) + .source("example.forecast.table") + .raw_point( + start_date="2021-01-01", + end_date="2021-01-02", + forecast_run_start_date="2021-01-01", + forecast_run_end_date="2021-01-02", + lat=0.1, + lon=0.1, + ) + ) + + print(data) + ``` + Args: - forecast (str): Business unit - region (str): Region - data_security_level (str): Level of data security - data_type (str): Type of the data (float, integer, double, string) start_date (str): Start date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) end_date (str): End date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_start_date (str): Start date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_end_date (str): End date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) lat (float): latitude lon (float): longitude - source (optional str): Source of the data ie ECMWF - time_zone (str): Timezone of the data - include_bad_data (bool): Include "Bad" data points with True or remove "Bad" data points with False limit (optional int): The number of rows to be returned Returns: DataFrame: A dataframe of raw timeseries data. """ raw_parameters = { - "forecast": forecast, - "region": region, - "data_security_level": data_security_level, - "data_type": data_type, + "source": self.data_source, "start_date": start_date, "end_date": end_date, + "forecast_run_start_date": forecast_run_start_date, + "forecast_run_end_date": forecast_run_end_date, + "timestamp_column": self.timestamp_column, + "forecast_run_timestamp_column": self.forecast_run_timestamp_column, "lat": lat, "lon": lon, - "source": source, - "time_zone": time_zone, - "include_bad_data": include_bad_data, "limit": limit, "supress_warning": True, } @@ -119,39 +146,50 @@ def raw_point( def latest_point( self, - forecast: str, - region: str, - data_security_level: str, - data_type: str, lat: float, lon: float, - source: str = None, limit: int = None, ) -> DataFrame: """ A function to return back the latest data for a point. + **Example:** + ```python + from rtdip_sdk.queries.weather.weather_query_builder import ( + WeatherQueryBuilder, + ) + from rtdip_sdk.authentication.azure import DefaultAuth + from rtdip_sdk.connectors import DatabricksSQLConnection + + auth = DefaultAuth().authenticate() + token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token + connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + + data = ( + WeatherQueryBuilder() + .connect(connection) + .source("example.forecast.table") + .latest_point( + lat=0.1, + lon=0.1, + ) + ) + + print(data) + ``` + Args: - forecast (str): Business unit - region (str): Region - data_security_level (str): Level of data security - data_type (str): Type of the data (float, integer, double, string) lat (float): latitude lon (float): longitude - source (optional str): Source of the data ie ECMWF limit (optional int): The number of rows to be returned Returns: DataFrame: A dataframe of raw timeseries data. """ raw_parameters = { - "forecast": forecast, - "region": region, - "data_security_level": data_security_level, - "data_type": data_type, + "source": self.data_source, "lat": lat, "lon": lon, - "source": source, "limit": limit, "supress_warning": True, } @@ -160,57 +198,76 @@ def latest_point( def raw_grid( # NOSONAR self, # NOSONAR - forecast: str, - region: str, - data_security_level: str, - data_type: str, start_date: str, end_date: str, + forecast_run_start_date: str, + forecast_run_end_date: str, min_lat: float, min_lon: float, max_lat: float, max_lon: float, - source: str = None, - time_zone: str = None, - include_bad_data: bool = False, limit: int = None, # NOSONAR ) -> DataFrame: """ - A function to return back raw data for a point. + A function to return back raw data for a grid. + + **Example:** + ```python + from rtdip_sdk.queries.weather.weather_query_builder import ( + WeatherQueryBuilder, + ) + from rtdip_sdk.authentication.azure import DefaultAuth + from rtdip_sdk.connectors import DatabricksSQLConnection + + auth = DefaultAuth().authenticate() + token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token + connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + + data = ( + WeatherQueryBuilder() + .connect(connection) + .source("example.forecast.table") + .raw_grid( + start_date="2021-01-01", + end_date="2021-01-02", + forecast_run_start_date="2021-01-01", + forecast_run_end_date="2021-01-02", + min_lat=0.1, + max_lat=0.1, + min_lon=0.1, + max_lon=0.1, + ) + ) + + print(data) + ``` Args: - forecast (str): Business unit - region (str): Region - data_security_level (str): Level of data security - data_type (str): Type of the data (float, integer, double, string) start_date (str): Start date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) end_date (str): End date (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_start_date (str): Start date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) + forecast_run_end_date (str): End date of the forecast run (Either a date in the format YY-MM-DD or a datetime in the format YYY-MM-DDTHH:MM:SS or specify the timezone offset in the format YYYY-MM-DDTHH:MM:SS+zz:zz) min_lat (float): Min latitude min_lon (float): Min longitude max_lat (float): Max latitude max_lon (float): Max longitude - source (optional str): Source of the data ie ECMWF - time_zone (str): Timezone of the data - include_bad_data (bool): Include "Bad" data points with True or remove "Bad" data points with False limit (optional int): The number of rows to be returned Returns: DataFrame: A dataframe of raw timeseries data. """ raw_parameters = { - "forecast": forecast, - "region": region, - "data_security_level": data_security_level, - "data_type": data_type, + "source": self.data_source, "start_date": start_date, "end_date": end_date, + "forecast_run_start_date": forecast_run_start_date, + "forecast_run_end_date": forecast_run_end_date, + "timestamp_column": self.timestamp_column, + "forecast_run_timestamp_column": self.forecast_run_timestamp_column, "min_lat": min_lat, "min_lon": min_lon, "max_lat": max_lat, "max_lon": max_lon, - "source": source, - "time_zone": time_zone, - "include_bad_data": include_bad_data, "limit": limit, "supress_warning": True, } @@ -219,45 +276,58 @@ def raw_grid( # NOSONAR def latest_grid( self, - forecast: str, - region: str, - data_security_level: str, - data_type: str, min_lat: float, min_lon: float, max_lat: float, max_lon: float, - source: str = None, limit: int = None, ) -> DataFrame: """ - A function to return back the latest data for a point. + A function to return back the latest data for a grid. + + **Example:** + ```python + from rtdip_sdk.queries.weather.weather_query_builder import ( + WeatherQueryBuilder, + ) + from rtdip_sdk.authentication.azure import DefaultAuth + from rtdip_sdk.connectors import DatabricksSQLConnection + + auth = DefaultAuth().authenticate() + token = auth.get_token("2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default").token + connection = DatabricksSQLConnection("{server_hostname}", "{http_path}", token) + + data = ( + WeatherQueryBuilder() + .connect(connection) + .source("example.forecast.table") + .latest_grid( + min_lat=0.1, + max_lat=0.1, + min_lon=0.1, + max_lon=0.1, + ) + ) + + print(data) + ``` Args: - forecast (str): Business unit - region (str): Region - data_security_level (str): Level of data security - data_type (str): Type of the data (float, integer, double, string) min_lat (float): Min latitude min_lon (float): Min longitude max_lat (float): Max latitude max_lon (float): Max longitude - source (optional str): Source of the data ie ECMWF limit (optional int): The number of rows to be returned Returns: DataFrame: A dataframe of raw timeseries data. """ raw_parameters = { - "forecast": forecast, - "region": region, - "data_security_level": data_security_level, - "data_type": data_type, + "source": self.data_source, "min_lat": min_lat, "min_lon": min_lon, "max_lat": max_lat, "max_lon": max_lon, - "source": source, "limit": limit, "supress_warning": True, } diff --git a/tests/sdk/python/rtdip_sdk/queries/weather/test_latest.py b/tests/sdk/python/rtdip_sdk/queries/weather/test_latest.py index 8fdb641df..1e85a62d3 100644 --- a/tests/sdk/python/rtdip_sdk/queries/weather/test_latest.py +++ b/tests/sdk/python/rtdip_sdk/queries/weather/test_latest.py @@ -37,39 +37,45 @@ DATABRICKS_SQL_CONNECT = "databricks.sql.connect" DATABRICKS_SQL_CONNECT_CURSOR = "databricks.sql.connect.cursor" INTERPOLATION_METHOD = "test/test/test" -MOCKED_QUERY_GRID = "SELECT * FROM `mocked-forecast`.`weather`.`mocked-region_weather_mocked-data-security-level_events_mocked-data-type` WHERE `Latitude` > '0' AND `Latitude` < '0.1' AND `Longitude` > '0' AND`Longitude` < '0.1' ORDER BY `TagName` " -MOCKED_QUERY_POINT = "SELECT * FROM `mocked-forecast`.`weather`.`mocked-region_weather_mocked-data-security-level_events_mocked-data-type` WHERE `Latitude` == '0' AND `Longitude` == '0' ORDER BY `TagName` " +MOCKED_QUERY_GRID = "SELECT * FROM `forecast`.`weather`.`mock_region_mock_security_events_mock_data_type_latest` WHERE `Latitude` > 36 AND `Latitude` < 38 AND `Longitude` > -109.1 AND `Longitude` < -107.1 ORDER BY `TagName` " +MOCKED_QUERY_POINT = "SELECT * FROM `forecast`.`weather`.`mock_region_mock_security_events_mock_data_type_latest` WHERE `Latitude` == 37 AND `Longitude` == -108.1 ORDER BY `TagName` " MOCKED_QUERY_OFFSET_LIMIT = "LIMIT 10 OFFSET 10 " + MOCKED_PARAMETER_DICT_GRID = { - "forecast": "mocked-forecast", - "region": "mocked-region", - "data_security_level": "mocked-data-security-level", - "data_type": "mocked-data-type", - "min_lat": 0, - "max_lat": 0.1, - "min_lon": 0, - "max_lon": 0.1, - "start_date": "2020-01-01", - "end_date": "2020-01-02", + "forecast": "forecast", + "forecast_type": "weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "min_lat": 36, + "max_lat": 38, + "min_lon": -109.1, + "max_lon": -107.1, } MOCKED_PARAMETER_DICT_POINT = { - "forecast": "mocked-forecast", - "region": "mocked-region", - "data_security_level": "mocked-data-security-level", - "data_type": "mocked-data-type", - "lat": 0, - "lon": 0, - "start_date": "2020-01-01", - "end_date": "2020-01-02", + "forecast": "forecast", + "forecast_type": "weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "lat": 37, + "lon": -108.1, +} + +MOCKED_PARAMETER_DICT_GRID_SOURCE = { + "source": "forecast`.`weather`.`mock_region_mock_security_events_mock_data_type_latest", + "min_lat": 36, + "max_lat": 38, + "min_lon": -109.1, + "max_lon": -107.1, } -MOCKED_NO_TAG_QUERY = "SELECT * FROM `mocked-business-unit`.`sensors`.`mocked-asset_mocked-data-security-level_events_latest` ORDER BY `TagName` " -MOCKED_PARAMETER_NO_TAGS_DICT = { - "forecast": "mocked-forecast", - "region": "mocked-region", - "data_security_level": "mocked-data-security-level", +MOCKED_PARAMETER_DICT_POINT_SOURCE = { + "source": "forecast`.`weather`.`mock_region_mock_security_events_mock_data_type_latest", + "lat": 37, + "lon": -108.1, } @@ -101,6 +107,34 @@ def test_latest_grid(mocker: MockerFixture): assert isinstance(actual, pd.DataFrame) +def test_latest_grid_source(mocker: MockerFixture): + mocked_cursor = mocker.spy(MockedDBConnection, "cursor") + mocked_connection_close = mocker.spy(MockedDBConnection, "close") + mocked_execute = mocker.spy(MockedCursor, "execute") + mocked_fetch_all = mocker.patch.object( + MockedCursor, + "fetchall_arrow", + return_value=pa.Table.from_pandas( + pd.DataFrame(data={"a": [1], "c": [2], "g": [3], "i": [4]}) + ), + ) + mocked_close = mocker.spy(MockedCursor, "close") + mocker.patch(DATABRICKS_SQL_CONNECT, return_value=MockedDBConnection()) + + mocked_connection = DatabricksSQLConnection( + SERVER_HOSTNAME, HTTP_PATH, ACCESS_TOKEN + ) + + actual = latest_grid(mocked_connection, MOCKED_PARAMETER_DICT_GRID_SOURCE) + + mocked_cursor.assert_called_once() + mocked_connection_close.assert_called_once() + mocked_execute.assert_called_once_with(mocker.ANY, query=MOCKED_QUERY_GRID) + mocked_fetch_all.assert_called_once() + mocked_close.assert_called_once() + assert isinstance(actual, pd.DataFrame) + + def test_latest_grid_fails(mocker: MockerFixture): mocker.spy(MockedDBConnection, "cursor") mocker.spy(MockedDBConnection, "close") @@ -145,6 +179,34 @@ def test_latest_point(mocker: MockerFixture): assert isinstance(actual, pd.DataFrame) +def test_latest_point_source(mocker: MockerFixture): + mocked_cursor = mocker.spy(MockedDBConnection, "cursor") + mocked_connection_close = mocker.spy(MockedDBConnection, "close") + mocked_execute = mocker.spy(MockedCursor, "execute") + mocked_fetch_all = mocker.patch.object( + MockedCursor, + "fetchall_arrow", + return_value=pa.Table.from_pandas( + pd.DataFrame(data={"h": [1], "o": [2], "u": [3], "w": [4]}) + ), + ) + mocked_close = mocker.spy(MockedCursor, "close") + mocker.patch(DATABRICKS_SQL_CONNECT, return_value=MockedDBConnection()) + + mocked_connection = DatabricksSQLConnection( + SERVER_HOSTNAME, HTTP_PATH, ACCESS_TOKEN + ) + + actual = latest_point(mocked_connection, MOCKED_PARAMETER_DICT_POINT_SOURCE) + + mocked_cursor.assert_called_once() + mocked_connection_close.assert_called_once() + mocked_execute.assert_called_once_with(mocker.ANY, query=MOCKED_QUERY_POINT) + mocked_fetch_all.assert_called_once() + mocked_close.assert_called_once() + assert isinstance(actual, pd.DataFrame) + + def test_latest_point_fails(mocker: MockerFixture): mocker.spy(MockedDBConnection, "cursor") mocker.spy(MockedDBConnection, "close") diff --git a/tests/sdk/python/rtdip_sdk/queries/weather/test_raw.py b/tests/sdk/python/rtdip_sdk/queries/weather/test_raw.py index 41aa0af29..aeea9b887 100644 --- a/tests/sdk/python/rtdip_sdk/queries/weather/test_raw.py +++ b/tests/sdk/python/rtdip_sdk/queries/weather/test_raw.py @@ -37,45 +37,73 @@ DATABRICKS_SQL_CONNECT = "databricks.sql.connect" DATABRICKS_SQL_CONNECT_CURSOR = "databricks.sql.connect.cursor" INTERPOLATION_METHOD = "test/test/test" -MOCKED_QUERY_GRID = "SELECT * FROM `mocked-forecast`.`weather`.`mocked-region_weather_mocked-data-security-level_events_mocked-data-type` WHERE `EventTime` BETWEEN to_timestamp(\"2020-01-01\") AND to_timestamp(\"2020-01-02\")AND `Latitude` > '1.1' AND `Latitude` < '1.1' AND `Longitude` > '1.1' AND`Longitude` < '1.1' ORDER BY `TagName` " -MOCKED_QUERY_POINT = "SELECT * FROM `mocked-forecast`.`weather`.`mocked-region_weather_mocked-data-security-level_events_mocked-data-type` WHERE `EventTime` BETWEEN to_timestamp(\"2020-01-01\") AND to_timestamp(\"2020-01-02\")AND `Latitude` > '1.1' AND `Longitude` > '1.1' ORDER BY `TagName` " +MOCKED_QUERY_GRID = 'SELECT * FROM `forecast`.`weather`.`mock_region_mock_security_events_mock_data_type` WHERE (`EventTime` BETWEEN to_timestamp("2024-01-01") AND to_timestamp("2024-01-03")) AND (`EnqueuedTime` BETWEEN to_timestamp("2023-12-28") AND to_timestamp("2023-12-31")) AND `Latitude` > 36 AND `Latitude` < 38 AND `Longitude` > -109.1 AND `Longitude` < -107.1 ORDER BY `TagName` ' +MOCKED_QUERY_POINT = 'SELECT * FROM `forecast`.`weather`.`mock_region_mock_security_events_mock_data_type` WHERE (`EventTime` BETWEEN to_timestamp("2024-01-01") AND to_timestamp("2024-01-03")) AND (`EnqueuedTime` BETWEEN to_timestamp("2023-12-28") AND to_timestamp("2023-12-31")) AND `Latitude` == 37 AND `Longitude` == -108.1 ORDER BY `TagName` ' MOCKED_QUERY_OFFSET_LIMIT = "LIMIT 10 OFFSET 10 " + MOCKED_PARAMETER_DICT_GRID = { - "forecast": "mocked-forecast", - "region": "mocked-region", - "data_security_level": "mocked-data-security-level", - "data_type": "mocked-data-type", - "min_lat": 1.1, - "max_lat": 1.1, - "min_lon": 1.1, - "max_lon": 1.1, - "start_date": "2020-01-01", - "end_date": "2020-01-02", + "forecast": "forecast", + "forecast_type": "weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "min_lat": 36, + "max_lat": 38, + "min_lon": -109.1, + "max_lon": -107.1, + "start_date": "2024-01-01", + "end_date": "2024-01-03", + "forecast_run_start_date": "2023-12-28", + "forecast_run_end_date": "2023-12-31", "timestamp_column": "EventTime", + "forecast_run_timestamp_column": "EnqueuedTime", } MOCKED_PARAMETER_DICT_POINT = { - "forecast": "mocked-forecast", - "region": "mocked-region", - "data_security_level": "mocked-data-security-level", - "data_type": "mocked-data-type", - "lat": 1.1, - "lon": 1.1, - "start_date": "2020-01-01", - "end_date": "2020-01-02", + "forecast": "forecast", + "forecast_type": "weather", + "region": "mock_region", + "data_security_level": "mock_security", + "data_type": "mock_data_type", + "lat": 37, + "lon": -108.1, + "start_date": "2024-01-01", + "end_date": "2024-01-03", + "forecast_run_start_date": "2023-12-28", + "forecast_run_end_date": "2023-12-31", "timestamp_column": "EventTime", + "forecast_run_timestamp_column": "EnqueuedTime", } -MOCKED_NO_TAG_QUERY = "SELECT * FROM `mocked-business-unit`.`sensors`.`mocked-asset_mocked-data-security-level_events_raw` ORDER BY `TagName` " -MOCKED_PARAMETER_NO_TAGS_DICT = { - "forecast": "mocked-forecast", - "region": "mocked-region", - "data_security_level": "mocked-data-security-level", +MOCKED_PARAMETER_DICT_GRID_SOURCE = { + "source": "forecast`.`weather`.`mock_region_mock_security_events_mock_data_type", + "min_lat": 36, + "max_lat": 38, + "min_lon": -109.1, + "max_lon": -107.1, + "start_date": "2024-01-01", + "end_date": "2024-01-03", + "forecast_run_start_date": "2023-12-28", + "forecast_run_end_date": "2023-12-31", + "timestamp_column": "EventTime", + "forecast_run_timestamp_column": "EnqueuedTime", } +MOCKED_PARAMETER_DICT_POINT_SOURCE = { + "source": "forecast`.`weather`.`mock_region_mock_security_events_mock_data_type", + "lat": 37, + "lon": -108.1, + "start_date": "2024-01-01", + "end_date": "2024-01-03", + "forecast_run_start_date": "2023-12-28", + "forecast_run_end_date": "2023-12-31", + "timestamp_column": "EventTime", + "forecast_run_timestamp_column": "EnqueuedTime", +} -def test_raw_grid(mocker: MockerFixture): + +def test_raw_grid_(mocker: MockerFixture): mocked_cursor = mocker.spy(MockedDBConnection, "cursor") mocked_connection_close = mocker.spy(MockedDBConnection, "close") mocked_execute = mocker.spy(MockedCursor, "execute") @@ -103,6 +131,34 @@ def test_raw_grid(mocker: MockerFixture): assert isinstance(actual, pd.DataFrame) +def test_raw_grid_source(mocker: MockerFixture): + mocked_cursor = mocker.spy(MockedDBConnection, "cursor") + mocked_connection_close = mocker.spy(MockedDBConnection, "close") + mocked_execute = mocker.spy(MockedCursor, "execute") + mocked_fetch_all = mocker.patch.object( + MockedCursor, + "fetchall_arrow", + return_value=pa.Table.from_pandas( + pd.DataFrame(data={"a": [1], "d": [2], "e": [3], "r": [4]}) + ), + ) + mocked_close = mocker.spy(MockedCursor, "close") + mocker.patch(DATABRICKS_SQL_CONNECT, return_value=MockedDBConnection()) + + mocked_connection = DatabricksSQLConnection( + SERVER_HOSTNAME, HTTP_PATH, ACCESS_TOKEN + ) + + actual = raw_grid(mocked_connection, MOCKED_PARAMETER_DICT_GRID_SOURCE) + + mocked_cursor.assert_called_once() + mocked_connection_close.assert_called_once() + mocked_execute.assert_called_once_with(mocker.ANY, query=MOCKED_QUERY_GRID) + mocked_fetch_all.assert_called_once() + mocked_close.assert_called_once() + assert isinstance(actual, pd.DataFrame) + + def test_raw_grid_fails(mocker: MockerFixture): mocker.spy(MockedDBConnection, "cursor") mocker.spy(MockedDBConnection, "close") @@ -147,6 +203,34 @@ def test_raw_point(mocker: MockerFixture): assert isinstance(actual, pd.DataFrame) +def test_raw_point_source(mocker: MockerFixture): + mocked_cursor = mocker.spy(MockedDBConnection, "cursor") + mocked_connection_close = mocker.spy(MockedDBConnection, "close") + mocked_execute = mocker.spy(MockedCursor, "execute") + mocked_fetch_all = mocker.patch.object( + MockedCursor, + "fetchall_arrow", + return_value=pa.Table.from_pandas( + pd.DataFrame(data={"u": [1], "v": [2], "w": [3], "x": [4]}) + ), + ) + mocked_close = mocker.spy(MockedCursor, "close") + mocker.patch(DATABRICKS_SQL_CONNECT, return_value=MockedDBConnection()) + + mocked_connection = DatabricksSQLConnection( + SERVER_HOSTNAME, HTTP_PATH, ACCESS_TOKEN + ) + + actual = raw_point(mocked_connection, MOCKED_PARAMETER_DICT_POINT_SOURCE) + + mocked_cursor.assert_called_once() + mocked_connection_close.assert_called_once() + mocked_execute.assert_called_once_with(mocker.ANY, query=MOCKED_QUERY_POINT) + mocked_fetch_all.assert_called_once() + mocked_close.assert_called_once() + assert isinstance(actual, pd.DataFrame) + + def test_raw_point_fails(mocker: MockerFixture): mocker.spy(MockedDBConnection, "cursor") mocker.spy(MockedDBConnection, "close") diff --git a/tests/sdk/python/rtdip_sdk/queries/weather/test_weather_query_builder.py b/tests/sdk/python/rtdip_sdk/queries/weather/test_weather_query_builder.py index 31ae7b67a..6b25f6520 100644 --- a/tests/sdk/python/rtdip_sdk/queries/weather/test_weather_query_builder.py +++ b/tests/sdk/python/rtdip_sdk/queries/weather/test_weather_query_builder.py @@ -33,12 +33,10 @@ def test_weather_query_builder_raw_point(mocker: MockerFixture): .connect(MOCK_CONNECTION) .source(MOCK_TABLE, status_column=None) .raw_point( - forecast="mock_forecast", - region="mock_region", - data_security_level="mock_data_security_level", - data_type="mock_data_type", start_date="2021-01-01", end_date="2021-01-02", + forecast_run_start_date="2021-01-01", + forecast_run_end_date="2021-01-02", lat=0.1, lon=0.1, ) @@ -57,10 +55,6 @@ def test_query_builder_latest_point(mocker: MockerFixture): .connect(MOCK_CONNECTION) .source(MOCK_TABLE, status_column=None) .latest_point( - forecast="mock_forecast", - region="mock_region", - data_security_level="mock_data_security_level", - data_type="mock_data_type", lat=0.1, lon=0.1, ) @@ -79,12 +73,10 @@ def test_weather_query_builder_raw_grid(mocker: MockerFixture): .connect(MOCK_CONNECTION) .source(MOCK_TABLE, status_column=None) .raw_grid( - forecast="mock_forecast", - region="mock_region", - data_security_level="mock_data_security_level", - data_type="mock_data_type", start_date="2021-01-01", end_date="2021-01-02", + forecast_run_start_date="2021-01-01", + forecast_run_end_date="2021-01-02", min_lat=0.1, max_lat=0.1, min_lon=0.1, @@ -105,10 +97,6 @@ def test_query_builder_latest_grid(mocker: MockerFixture): .connect(MOCK_CONNECTION) .source(MOCK_TABLE, status_column=None) .latest_grid( - forecast="mock_forecast", - region="mock_region", - data_security_level="mock_data_security_level", - data_type="mock_data_type", min_lat=0.1, max_lat=0.1, min_lon=0.1,