SIENTIAPDE-1094
Refactor activity imports and remove unused base and logger files; update requirements for library versioning
This commit is contained in:
@@ -1,13 +1,13 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from laborious.activities.postgres import Postgres
|
||||
from sientia_do.temporal.activities.postgres import Postgres
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from laborious.activities.gates import Gates
|
||||
from laborious.activities.opc import OPC
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
|
||||
|
||||
class Activities(Postgres, MLFlow, Gates, OPC):
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from temporalio import activity
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
|
||||
|
||||
class BaseActivity:
|
||||
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
|
||||
self.logger = logger
|
||||
self.notification_handler = notification_handler
|
||||
|
||||
@activity.defn(name="prepare_activity")
|
||||
async def prepare_activity(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Prepare the activity for the notification handler.
|
||||
|
||||
Args:
|
||||
workflow_name (str): The name of the workflow.
|
||||
schedule_name (str): The name of the schedule.
|
||||
model_name (str): The name of the model.
|
||||
model_id (str): The id of the model.
|
||||
"""
|
||||
self.notification_handler.base_notification.pipeline_name = input_data['workflow_name']
|
||||
self.notification_handler.base_notification.schedule_name = input_data['schedule_name']
|
||||
self.notification_handler.base_notification.model_name = input_data['model_name']
|
||||
self.notification_handler.base_notification.model_id = input_data['model_id']
|
||||
@@ -6,7 +6,7 @@ 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 sientia_do.temporal.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 (
|
||||
|
||||
@@ -4,11 +4,11 @@ from temporalio import activity, workflow
|
||||
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from laborious.activities.base import BaseActivity
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
|
||||
|
||||
class MLFlow(BaseActivity):
|
||||
|
||||
@@ -5,7 +5,7 @@ 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 sientia_do.temporal.activities.base import BaseActivity
|
||||
from laborious.utils.repository.opc_repository import OpcRepository
|
||||
from typing import Any
|
||||
import traceback
|
||||
|
||||
@@ -1,181 +0,0 @@
|
||||
import traceback
|
||||
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
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from typing import Any
|
||||
|
||||
|
||||
class Postgres(BaseActivity):
|
||||
def __init__(self, host: str, port: int,
|
||||
user: str, password: str, dbname: str,
|
||||
min_connections: int, max_connections: int,
|
||||
logger: Logger, notification_handler: NotificationHandler):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.dbname = 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)
|
||||
|
||||
BaseActivity.__init__(self, logger, notification_handler)
|
||||
|
||||
def close(self):
|
||||
self.engine.dispose()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
@activity.defn(name="load_custom_query")
|
||||
async def load_custom_query(self, query: str) -> dict[str, Any]:
|
||||
"""
|
||||
Loads data from a custom query.
|
||||
|
||||
Args:
|
||||
query (str): The query to load data from.
|
||||
|
||||
Returns:
|
||||
dict[str, dict]: The data from the query.
|
||||
"""
|
||||
self.logger.info(f"Fetching data from query: {query}")
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
self.logger.error(trace)
|
||||
|
||||
return {}
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
if data is None:
|
||||
return {}
|
||||
|
||||
# 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: \n{data.to_string()}")
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
@activity.defn(name="repeat_last_prediction")
|
||||
async def repeat_last_prediction(self, query_items: dict[str, str]):
|
||||
"""
|
||||
Repeats the last prediction for a given model.
|
||||
|
||||
Args:
|
||||
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 (int): The model to repeat the prediction for.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
schema = query_items["schema"]
|
||||
table_name = query_items["table_name"]
|
||||
model = query_items["model"]
|
||||
|
||||
repeat_query = f"""
|
||||
INSERT INTO \"{schema}\".{table_name} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
|
||||
SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
|
||||
FROM \"{schema}\".{table_name}
|
||||
WHERE model_id = {model}
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT 1;
|
||||
"""
|
||||
self.logger.info(f"Repeating last prediction for model {model}")
|
||||
self.logger.debug(f"Query: {repeat_query}")
|
||||
|
||||
with self.session_factory() as session:
|
||||
try:
|
||||
session.execute(repeat_query)
|
||||
session.commit()
|
||||
|
||||
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)
|
||||
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
@activity.defn(name="export_data_to_postgres")
|
||||
async def export_data_to_postgres(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Exports data to a postgres table.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The data to export. Contains:
|
||||
schema (str): The schema of the table.
|
||||
table_name (str): The name of the table.
|
||||
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"])
|
||||
|
||||
with self.session_factory() as session:
|
||||
try:
|
||||
data.to_sql(table_name, self.engine, schema=schema,
|
||||
if_exists="append", index=False)
|
||||
session.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
|
||||
)
|
||||
|
||||
self.logger.error(trace)
|
||||
|
||||
else:
|
||||
self.logger.debug("Data exported to postgres")
|
||||
finally:
|
||||
session.close()
|
||||
Reference in New Issue
Block a user