79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
import traceback
|
|
from temporalio import workflow, activity
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from psycopg2.pool import ThreadedConnectionPool
|
|
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
|
|
|
|
self.pool = ThreadedConnectionPool(
|
|
minconn=min_connections,
|
|
maxconn=max_connections,
|
|
host=self.host,
|
|
port=self.port,
|
|
user=self.user,
|
|
password=self.password,
|
|
dbname=self.dbname)
|
|
|
|
super().__init__(logger, notification_handler)
|
|
|
|
def close(self):
|
|
self.pool.closeall()
|
|
|
|
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.
|
|
"""
|
|
|
|
schema = input_data["schema"]
|
|
table_name = input_data["table_name"]
|
|
data = DataFrame(input_data["data"])
|
|
|
|
conn = self.pool.getconn()
|
|
|
|
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
|
|
)
|
|
|
|
self.logger.error(trace)
|
|
|
|
finally:
|
|
self.pool.putconn(conn)
|