SIENTIAPDE-994
Refactor and enhance the laborious workflow and utilities - Removed outdated test file `test_predictions_batch.py` from workflows. - Added `input_sample.json` for standardized input configuration. - Introduced `connectors_config.py` to manage database and service configurations. - Implemented a logging utility in `logger.py` for consistent logging across the application. - Created `policies.py` to define retry policies for workflows. - Developed comprehensive tests for `MLFlowRepository` in `test_model_repository.py`. - Added extensive tests for `OpcRepository` in `test_opc_repository.py`. - Updated `test_predictions_batch.py` to reflect new workflow structure and testing methodology.
This commit is contained in:
@@ -40,15 +40,10 @@ class Activities(Postgres, MLFlow, Gates, OPC):
|
||||
notification_handler=notification_handler)
|
||||
|
||||
OPC.__init__(self,
|
||||
name=opc_config['name'],
|
||||
url=opc_config['url'],
|
||||
server_uri=opc_config['server_uri'],
|
||||
cert_path=opc_config['cert_path'],
|
||||
private_key_path=opc_config['private_key_path'],
|
||||
server_cert_path=opc_config['server_cert_path'],
|
||||
opc_servers=opc_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
@activity.defn(name="prepare_activity")
|
||||
def prepare_activity(self, input_data: dict[str, Any]):
|
||||
super().prepare_activity(input_data)
|
||||
async def prepare_activity(self, input_data: dict[str, Any]):
|
||||
await super().prepare_activity(input_data)
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from temporalio import activity
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
@@ -9,7 +10,7 @@ class BaseActivity:
|
||||
self.notification_handler = notification_handler
|
||||
|
||||
@activity.defn(name="prepare_activity")
|
||||
def prepare_activity(self, input_data: dict[str, Any]):
|
||||
async def prepare_activity(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Prepare the activity for the notification handler.
|
||||
|
||||
|
||||
@@ -5,67 +5,85 @@ with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from laborious.activities.base import BaseActivity
|
||||
from typing import Any
|
||||
from laborious.utils.filters.conditional_filters import filter_empty_data, filter_specific_variables_null_values
|
||||
from pandas import DataFrame
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from laborious.activities.base import BaseActivity
|
||||
from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
|
||||
from typing import Any
|
||||
from laborious.utils.filters.conditional_filters import (
|
||||
filter_empty_data,
|
||||
filter_specific_variables_null_values
|
||||
)
|
||||
from pandas import DataFrame
|
||||
from datetime import datetime
|
||||
|
||||
input_filter_functions = {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
|
||||
'EMPTY_DATA': filter_empty_data,
|
||||
'path_confidence': {
|
||||
'stop': -1,
|
||||
'continue': 2,
|
||||
'repeat': -1
|
||||
'STOP': -1,
|
||||
'CONTINUE': 2,
|
||||
'REPEAT': -1
|
||||
}
|
||||
}
|
||||
|
||||
mlflow_response_filter_functions = {
|
||||
'API_ERROR': api_error_filter,
|
||||
'path_confidence': {
|
||||
'stop': -1,
|
||||
'continue': 10,
|
||||
'repeat': -1
|
||||
'STOP': -1,
|
||||
'CONTINUE': 10,
|
||||
'REPEAT': -1
|
||||
},
|
||||
}
|
||||
|
||||
mlflow_content_filter_functions = {
|
||||
'NAN_VALUES': nan_values_filter,
|
||||
'path_confidence': {
|
||||
'stop': -1,
|
||||
'continue': 18,
|
||||
'repeat': -1
|
||||
'STOP': -1,
|
||||
'CONTINUE': 18,
|
||||
'REPEAT': -1
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class Gates(BaseActivity):
|
||||
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
|
||||
super().__init__(logger, notification_handler)
|
||||
BaseActivity.__init__(self, logger, notification_handler)
|
||||
|
||||
@activity.defn(name="input_gate")
|
||||
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
|
||||
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
"""
|
||||
Filters the data based on the filters. The return value is a tuple with the first element
|
||||
being the policy and the second element being the confidence status.
|
||||
Args:
|
||||
input_data (dict): The input data. Contains:
|
||||
filters (dict): The filters to apply.
|
||||
The key is the filter name and the value is the filter configuration.
|
||||
data (dict[str, Any]): The data to filter.
|
||||
path_priority (list[str]): The path priority.
|
||||
Returns:
|
||||
tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions.
|
||||
tuple[str | None, int, str]: (policy, confidence) based in priority
|
||||
list and filter configuration and functions.
|
||||
"""
|
||||
|
||||
self.logger.debug("Performing input gate...")
|
||||
|
||||
filters = input_data['filters']
|
||||
data = DataFrame(input_data['data'])
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
|
||||
self.logger.debug(f"Input data:\n {data.to_string()}")
|
||||
self.logger.debug(f"Filters: {filters}")
|
||||
|
||||
for fil, config in filters.items():
|
||||
if fil not in input_filter_functions:
|
||||
self.logger.error(f"Filter {fil} not found")
|
||||
continue
|
||||
try:
|
||||
if input_filter_functions[fil](data, config):
|
||||
self.logger.debug(
|
||||
f"Data not passed the input filter {fil}:{config}")
|
||||
filter_output.append(config['POLICY'])
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
@@ -79,14 +97,18 @@ class Gates(BaseActivity):
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
return path_flag, input_filter_functions['path_confidence'][path_flag]
|
||||
self.logger.debug(f"Input gate result: {path_flag}")
|
||||
return path_flag, input_filter_functions['path_confidence'][path_flag], \
|
||||
"Input data with bad quality"
|
||||
|
||||
return None, 0
|
||||
self.logger.debug("Nothing was filtered by the input gate")
|
||||
return None, 0, ""
|
||||
|
||||
@activity.defn(name="mlflow_response_gate")
|
||||
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
|
||||
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
"""
|
||||
Filters the data based on the mlflow response filters. The return value is a tuple with the first element
|
||||
Filters the data based on the mlflow response filters.
|
||||
The return value is a tuple with the first element
|
||||
being the policy and the second element being the confidence status.
|
||||
Args:
|
||||
input_data (dict): The input data. Contains:
|
||||
@@ -95,17 +117,29 @@ class Gates(BaseActivity):
|
||||
path_priority (list[str]): The path priority list.
|
||||
type (str): The type of the gate.
|
||||
Returns:
|
||||
tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions.
|
||||
tuple[str | None, int, str]: (policy, confidence) based in priority list
|
||||
and filter configuration and functions.
|
||||
"""
|
||||
|
||||
self.logger.debug("Performing mlflow response gate...")
|
||||
|
||||
filters = input_data['filters']
|
||||
data = input_data['data']
|
||||
gate_type = input_data['type']
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
|
||||
self.logger.debug(f"Input data:\n {data}")
|
||||
self.logger.debug(f"Filters: {filters}")
|
||||
|
||||
comments = []
|
||||
for fil, config in filters.items():
|
||||
if fil not in mlflow_response_filter_functions:
|
||||
continue
|
||||
if mlflow_response_filter_functions[fil](data, config):
|
||||
filter_output.append(config['POLICY'])
|
||||
comments.append(data['content']['message'])
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
|
||||
message=data['content']['message'],
|
||||
@@ -116,14 +150,18 @@ class Gates(BaseActivity):
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag]
|
||||
self.logger.debug(f"Mlflow response gate result: {path_flag}")
|
||||
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
|
||||
", ".join(comments)
|
||||
|
||||
return None, 0
|
||||
self.logger.debug("Nothing was filtered by the mlflow response gate")
|
||||
return None, 0, ""
|
||||
|
||||
@activity.defn(name="mlflow_content_gate")
|
||||
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
|
||||
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
"""
|
||||
Filters the data based on the mlflow content filters. The return value is a tuple with the first element
|
||||
Filters the data based on the mlflow content filters.
|
||||
The return value is a tuple with the first element
|
||||
being the policy and the second element being the confidence status.
|
||||
Args:
|
||||
input_data (dict): The input data. Contains:
|
||||
@@ -132,9 +170,12 @@ class Gates(BaseActivity):
|
||||
path_priority (list[str]): The path priority list.
|
||||
type (str): The type of the gate.
|
||||
Returns:
|
||||
tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions.
|
||||
tuple[str | None, int, str]: (policy, confidence) based in priority
|
||||
list and filter configuration and functions.
|
||||
"""
|
||||
|
||||
self.logger.debug("Performing mlflow content gate...")
|
||||
|
||||
filters = input_data['filters']
|
||||
data = DataFrame(input_data['data'])
|
||||
gate_type = input_data['type']
|
||||
@@ -142,7 +183,12 @@ class Gates(BaseActivity):
|
||||
|
||||
filter_output = []
|
||||
|
||||
self.logger.debug(f"Input data:\n {data}")
|
||||
self.logger.debug(f"Filters: {filters}")
|
||||
|
||||
for fil, config in filters.items():
|
||||
if fil not in mlflow_content_filter_functions:
|
||||
continue
|
||||
if mlflow_content_filter_functions[fil](data, config):
|
||||
filter_output.append(config['POLICY'])
|
||||
self.notification_handler.build_and_send_notification(
|
||||
@@ -155,9 +201,12 @@ class Gates(BaseActivity):
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag]
|
||||
self.logger.debug(f"Mlflow content gate result: {path_flag}")
|
||||
return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
|
||||
"Transformed data not passed the content filter"
|
||||
|
||||
return None, 0
|
||||
self.logger.debug("Nothing was filtered by the mlflow content gate")
|
||||
return None, 0, ""
|
||||
|
||||
@activity.defn(name="format_prediction")
|
||||
async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -172,12 +221,14 @@ class Gates(BaseActivity):
|
||||
Returns:
|
||||
dict: The formatted data.
|
||||
"""
|
||||
self.logger.debug("Formatting prediction...")
|
||||
|
||||
data = DataFrame(input_data['data'])
|
||||
data['timestamp'] = input_data['timestamp']
|
||||
data['model_id'] = input_data['model_id']
|
||||
data['prediction_confidence'] = input_data['prediction_confidence']
|
||||
data['prediction_status'] = 'Good'
|
||||
data['comment'] = ""
|
||||
data['comments'] = ""
|
||||
data.sort_values(by='timestamp', inplace=True)
|
||||
|
||||
return data.to_dict()
|
||||
@@ -198,6 +249,8 @@ class Gates(BaseActivity):
|
||||
dict: The formatted data.
|
||||
"""
|
||||
|
||||
self.logger.debug("Formatting default prediction...")
|
||||
|
||||
return DataFrame({
|
||||
'prediction': [0],
|
||||
'response_time': [0],
|
||||
@@ -205,7 +258,7 @@ class Gates(BaseActivity):
|
||||
'model_id': [input_data['model_id']],
|
||||
'prediction_confidence': [input_data['prediction_confidence']],
|
||||
'prediction_status': ['Bad'],
|
||||
'comment': [input_data['comment']]
|
||||
'comments': [input_data['comment']]
|
||||
}).to_dict()
|
||||
|
||||
@activity.defn(name="get_last_timestamp")
|
||||
@@ -219,4 +272,6 @@ class Gates(BaseActivity):
|
||||
str: The last timestamp of the data.
|
||||
"""
|
||||
data = DataFrame(input_data['data'])
|
||||
if data.empty:
|
||||
return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||||
return max(data['timestamp'].values.tolist())
|
||||
|
||||
@@ -5,17 +5,16 @@ from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from laborious.activities.base import BaseActivity
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
|
||||
class MLFlow(BaseActivity):
|
||||
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
|
||||
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
|
||||
super().__init__(logger, notification_handler)
|
||||
BaseActivity.__init__(self, logger, notification_handler)
|
||||
self.mlflow_host = mlflow_host
|
||||
self.mlflow_port = mlflow_port
|
||||
self.mlflow_username = mlflow_username
|
||||
@@ -33,7 +32,7 @@ class MLFlow(BaseActivity):
|
||||
input_data (dict): The input data. Contains:
|
||||
data (dict[str, Any]): The data to transform.
|
||||
model_name (str): The name of the model.
|
||||
model_retention (int): The retention of the model.
|
||||
model_retention (int): The retention of the model in minutes.
|
||||
Returns:
|
||||
dict[str, Any]: The transformed data.
|
||||
"""
|
||||
@@ -54,6 +53,8 @@ class MLFlow(BaseActivity):
|
||||
response_data = self.model_monitoring_repository.transform(
|
||||
model_name, data, model_retention)
|
||||
|
||||
self.logger.debug(response_data)
|
||||
|
||||
return response_data
|
||||
|
||||
@activity.defn(name="request_predict")
|
||||
@@ -80,4 +81,6 @@ class MLFlow(BaseActivity):
|
||||
response_data = self.model_monitoring_repository.predict(
|
||||
model_name, data, model_retention)
|
||||
|
||||
self.logger.debug(response_data)
|
||||
|
||||
return response_data
|
||||
|
||||
@@ -4,40 +4,55 @@ from temporalio import activity, workflow
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from laborious.activities.base import BaseActivity
|
||||
from laborious.utils.repository.opc_repository import OpcRepository
|
||||
from typing import Any
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
import traceback
|
||||
from pandas import DataFrame
|
||||
|
||||
|
||||
class OPC(BaseActivity):
|
||||
def __init__(self,
|
||||
name: str, url: str, server_uri: str,
|
||||
cert_path: str, private_key_path: str, server_cert_path: str,
|
||||
def __init__(self, opc_servers: dict[str, dict[str, Any]],
|
||||
logger: Logger, notification_handler: NotificationHandler):
|
||||
|
||||
self.logger = logger
|
||||
self.notification_handler = notification_handler
|
||||
self.name = name
|
||||
self.url = url
|
||||
self.server_uri = server_uri
|
||||
self.cert_path = cert_path
|
||||
self.private_key_path = private_key_path
|
||||
self.server_cert_path = server_cert_path
|
||||
self.opc_servers = opc_servers
|
||||
|
||||
self.opc_repository = OpcRepository(
|
||||
name=self.name,
|
||||
url=self.url,
|
||||
logger=self.logger,
|
||||
server_uri=self.server_uri,
|
||||
cert_path=self.cert_path,
|
||||
private_key_path=self.private_key_path,
|
||||
server_cert_path=self.server_cert_path
|
||||
)
|
||||
self.opc_repository = {}
|
||||
for name, server in opc_servers.items():
|
||||
self.opc_repository[name] = OpcRepository(
|
||||
name=name,
|
||||
url=server['url'],
|
||||
logger=self.logger,
|
||||
server_uri=server['server_uri'],
|
||||
cert_path=server['cert_path'],
|
||||
private_key_path=server['private_key_path'],
|
||||
server_cert_path=server['server_cert_path'],
|
||||
notification_handler=self.notification_handler,
|
||||
reconnection_interval=server['reconnection_interval'],
|
||||
)
|
||||
self.opc_repository[name].connect()
|
||||
|
||||
self.opc_repository.connect()
|
||||
BaseActivity.__init__(self, logger, notification_handler)
|
||||
|
||||
def write_data(self, server: str, tag: str, data: Any,
|
||||
data_type: str, tag_type: str):
|
||||
try:
|
||||
self.opc_repository[server].write_data(
|
||||
tag, data, data_type)
|
||||
self.logger.debug(f"Wrote {tag_type} to {tag}")
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
|
||||
message=f"Error writing data to OPC server: {e}",
|
||||
block="write_opc_data",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
self.logger.error(trace)
|
||||
|
||||
@activity.defn(name='write_opc_data')
|
||||
async def write_opc_data(self, input_data: dict[str, Any]):
|
||||
@@ -47,46 +62,40 @@ class OPC(BaseActivity):
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data. Contains the following keys:
|
||||
data (dict[str, Any]): The dataframe that contains the data to write to the OPC servers.
|
||||
opc_servers (list[str]): The OPC servers to write to.
|
||||
opc_output_config (dict[str, Any]): The OPC writing configuration. Contains:
|
||||
- data (dict[str, Any]): The dataframe that contains the data to write
|
||||
to the OPC servers.
|
||||
- opc_output_config (dict[str, Any]): The OPC writing configuration.
|
||||
The keys are the OPC server names and the values contain:
|
||||
prediction_tags (dict[str, Any]): The tags to write to the OPC servers.
|
||||
confidence_tags (dict[str, Any]): The tags to write to the OPC servers.
|
||||
|
||||
Returns:
|
||||
"""
|
||||
self.logger.debug("Writing data to OPC servers...")
|
||||
data = DataFrame(input_data['data'])
|
||||
_opc_servers = input_data['opc_servers']
|
||||
opc_output_config = input_data['opc_output_config']
|
||||
self.logger.debug(data)
|
||||
|
||||
if 'prediction_tags' in opc_output_config:
|
||||
for tag, config in opc_output_config['prediction_tags'].items():
|
||||
try:
|
||||
self.opc_repository.write_data(
|
||||
tag, data.head(1)['prediction'].values[0], config['data_type'])
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="WRITE_OPC_PREDICTION_ERROR",
|
||||
message=f"Error writing data to OPC server: {e}",
|
||||
block="write_opc_data",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
self.logger.error(trace)
|
||||
for server, config in opc_output_config.items():
|
||||
if self.opc_repository.get(server) is None:
|
||||
self.logger.error(f"OPC server {server} not found")
|
||||
continue
|
||||
|
||||
if 'confidence_tags' in opc_output_config:
|
||||
for tag, config in opc_output_config['confidence_tags'].items():
|
||||
try:
|
||||
self.opc_repository.write_data(
|
||||
tag, data.head(1)['prediction_confidence'].values[0], config['data_type'])
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="WRITE_OPC_CONFIDENCE_ERROR",
|
||||
message=f"Error writing data to OPC server: {e}",
|
||||
block="write_opc_data",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
if 'prediction_tags' in config:
|
||||
for tag, tag_config in config['prediction_tags'].items():
|
||||
self.write_data(
|
||||
server=server,
|
||||
tag=tag,
|
||||
data=data.head(1)['prediction'].values[0],
|
||||
data_type=tag_config['data_type'],
|
||||
tag_type='prediction'
|
||||
)
|
||||
if 'confidence_tags' in config:
|
||||
for tag, tag_config in config['confidence_tags'].items():
|
||||
self.write_data(
|
||||
server=server,
|
||||
tag=tag,
|
||||
data=data.head(1)['prediction_confidence'].values[0],
|
||||
data_type=tag_config['data_type'],
|
||||
tag_type='confidence'
|
||||
)
|
||||
self.logger.error(trace)
|
||||
|
||||
@@ -3,6 +3,9 @@ from temporalio import workflow, activity
|
||||
|
||||
from laborious.activities.base import BaseActivity
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy.pool import QueuePool
|
||||
from psycopg2.pool import ThreadedConnectionPool
|
||||
from pandas import read_sql_query, DataFrame
|
||||
from logging import Logger
|
||||
@@ -22,19 +25,20 @@ class Postgres(BaseActivity):
|
||||
self.password = password
|
||||
self.dbname = dbname
|
||||
|
||||
self.pool = ThreadedConnectionPool(
|
||||
minconn=min_connections,
|
||||
maxconn=max_connections,
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
user=self.user,
|
||||
password=self.password,
|
||||
dbname=self.dbname)
|
||||
# Create SQLAlchemy engine with connection pooling
|
||||
self.engine = create_engine(
|
||||
f'postgresql://{user}:{password}@{host}:{port}/{dbname}',
|
||||
poolclass=QueuePool,
|
||||
pool_size=min_connections,
|
||||
max_overflow=max_connections - min_connections,
|
||||
pool_pre_ping=True
|
||||
)
|
||||
self.session_factory = sessionmaker(bind=self.engine)
|
||||
|
||||
super().__init__(logger, notification_handler)
|
||||
BaseActivity.__init__(self, logger, notification_handler)
|
||||
|
||||
def close(self):
|
||||
self.pool.closeall()
|
||||
self.engine.dispose()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
@@ -52,28 +56,36 @@ class Postgres(BaseActivity):
|
||||
"""
|
||||
self.logger.info(f"Fetching data from query: {query}")
|
||||
|
||||
conn = self.pool.getconn()
|
||||
try:
|
||||
data = read_sql_query(query, conn)
|
||||
data = None
|
||||
with self.session_factory() as session:
|
||||
try:
|
||||
data = read_sql_query(query, self.engine)
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="ERROR_LOADING_CUSTOM_QUERY",
|
||||
message=f"Error fetching data from query: {e}",
|
||||
block="load_custom_query",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="ERROR_LOADING_CUSTOM_QUERY",
|
||||
message=f"Error fetching data from query: {e}",
|
||||
block="load_custom_query",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
|
||||
self.logger.error(trace)
|
||||
self.logger.error(trace)
|
||||
|
||||
return {}
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
if data is None:
|
||||
return {}
|
||||
finally:
|
||||
self.pool.putconn(conn)
|
||||
|
||||
# Converts any datetime datatype columns to string
|
||||
for col in data.select_dtypes(include=['datetime64']).columns:
|
||||
data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S')
|
||||
|
||||
self.logger.info(f"Fetched {len(data)} rows")
|
||||
self.logger.debug(f"Data: {data.to_string()}")
|
||||
self.logger.debug(f"Data: \n{data.to_string()}")
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
@@ -86,7 +98,7 @@ class Postgres(BaseActivity):
|
||||
query_items (dict[str, str]): The query items. Contains:
|
||||
schema (str): The schema of the table.
|
||||
table_name (str): The name of the table.
|
||||
model (str): The model to repeat the prediction for.
|
||||
model (int): The model to repeat the prediction for.
|
||||
|
||||
Returns:
|
||||
None
|
||||
@@ -106,28 +118,25 @@ class Postgres(BaseActivity):
|
||||
self.logger.info(f"Repeating last prediction for model {model}")
|
||||
self.logger.debug(f"Query: {repeat_query}")
|
||||
|
||||
conn = self.pool.getconn()
|
||||
with self.session_factory() as session:
|
||||
try:
|
||||
session.execute(repeat_query)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(repeat_query)
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="ERROR_REPEATING_LAST_PREDICTION",
|
||||
message=f"Error repeating last prediction: {e}",
|
||||
block="repeat_last_prediction",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="ERROR_REPEATING_LAST_PREDICTION",
|
||||
message=f"Error repeating last prediction: {e}",
|
||||
block="repeat_last_prediction",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
self.logger.error(trace)
|
||||
|
||||
self.logger.error(trace)
|
||||
|
||||
finally:
|
||||
self.pool.putconn(conn)
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
@activity.defn(name="export_data_to_postgres")
|
||||
async def export_data_to_postgres(self, input_data: dict[str, Any]):
|
||||
@@ -141,28 +150,32 @@ class Postgres(BaseActivity):
|
||||
data (DataFrame): The data to export.
|
||||
"""
|
||||
|
||||
self.logger.debug(
|
||||
f"Exporting data to postgres: {input_data['data']}")
|
||||
|
||||
schema = input_data["schema"]
|
||||
table_name = input_data["table_name"]
|
||||
data = DataFrame(input_data["data"])
|
||||
|
||||
conn = self.pool.getconn()
|
||||
with self.session_factory() as session:
|
||||
try:
|
||||
data.to_sql(table_name, self.engine, schema=schema,
|
||||
if_exists="append", index=False)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
data.to_sql(table_name, conn, schema=schema,
|
||||
if_exists="append", index=False)
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
|
||||
message=f"Error exporting data to postgres: {e}",
|
||||
block="export_data_to_postgres",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
|
||||
message=f"Error exporting data to postgres: {e}",
|
||||
block="export_data_to_postgres",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
self.logger.error(trace)
|
||||
|
||||
self.logger.error(trace)
|
||||
|
||||
finally:
|
||||
self.pool.putconn(conn)
|
||||
else:
|
||||
self.logger.debug("Data exported to postgres")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
Reference in New Issue
Block a user