b203b7d22c
Implement workflows for fake data generation, scouter processing, and core scouter operations - Added `FakeData` workflow to generate random data and send it to a Kafka topic. - Implemented `Scouter` workflow to load data from Kafka and trigger the core scouter workflow. - Created `CoreScouter` workflow to process data through quality gates, aggregation, and export to PostgreSQL. - Developed comprehensive unit tests for activities and workflows, ensuring proper functionality and error handling. - Enhanced Redis and Postgres activities with robust testing for data handling and error notifications. - Introduced quality filters for data validation and implemented tests to verify their functionality.
86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
import traceback
|
|
from temporalio import workflow, activity
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.pool import QueuePool
|
|
from pandas import DataFrame
|
|
from logging import Logger
|
|
from sientia_do.notifications.handlers import NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from scouter.activities.base import BaseActivity
|
|
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="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()
|