"""Database related functions""" from os import popen from sqlalchemy import text, event from sqlalchemy.schema import CreateSchema from sqlalchemy.exc import InternalError from sqlalchemy.orm import sessionmaker, scoped_session, Session as SASession from sqlalchemy.exc import SQLAlchemyError from flask import current_app from common.extensions import db, migrate class Database: """used for managing tenant databases related operations""" def __init__(self, tenant: str) -> None: self.schema = str(tenant) # --- Session / Transaction events to ensure correct search_path per transaction --- @event.listens_for(SASession, "after_begin") def _set_search_path_per_tx(session, transaction, connection): """Ensure each transaction sees the right tenant schema, regardless of which pooled connection is used. Uses SET LOCAL so it is scoped to the tx. """ schema = session.info.get("tenant_schema") if schema: current_app.logger.debug(f"DBCTX tx_begin schema={schema}") try: connection.exec_driver_sql(f'SET LOCAL search_path TO "{schema}", public') # Optional visibility/logging for debugging sp = connection.exec_driver_sql("SHOW search_path").scalar() try: current_app.logger.info(f"DBCTX tx_begin conn_id={id(connection.connection)} search_path={sp}") except Exception: pass except Exception as e: try: current_app.logger.error(f"Failed to SET LOCAL search_path for schema {schema}: {e!r}") except Exception: pass def _log_db_context(self, origin: str = "") -> None: """Log key DB context info to diagnose schema/search_path issues. Collects and logs in a single structured line: - current_database() - inet_server_addr(), inet_server_port() - SHOW search_path - current_schema() - to_regclass('interaction') - to_regclass('.interaction') """ try: db_name = db.session.execute(text("SELECT current_database()"))\ .scalar() host = db.session.execute(text("SELECT inet_server_addr()"))\ .scalar() port = db.session.execute(text("SELECT inet_server_port()"))\ .scalar() search_path = db.session.execute(text("SHOW search_path"))\ .scalar() current_schema = db.session.execute(text("SELECT current_schema()"))\ .scalar() reg_unqualified = db.session.execute(text("SELECT to_regclass('interaction')"))\ .scalar() qualified = f"{self.schema}.interaction" reg_qualified = db.session.execute( text("SELECT to_regclass(:qn)"), {"qn": qualified} ).scalar() current_app.logger.info( "DBCTX origin=%s db=%s host=%s port=%s search_path=%s current_schema=%s to_regclass(interaction)=%s to_regclass(%s)=%s", origin, db_name, host, port, search_path, current_schema, reg_unqualified, qualified, reg_qualified ) except SQLAlchemyError as e: current_app.logger.error( f"DBCTX logging failed at {origin} for schema {self.schema}: {e!r}" ) def get_engine(self): """create new schema engine""" return db.engine.execution_options( schema_translate_map={None: self.schema} ) def get_session(self): """To get session of tenant/public database schema for quick use returns: session: session of tenant/public database schema """ return scoped_session( sessionmaker(bind=self.get_engine(), expire_on_commit=True) ) def create_schema(self): """create new database schema, mostly used on tenant creation""" try: db.session.execute(CreateSchema(self.schema)) # db.session.commit() db.session.execute(text(f"SET search_path TO {self.schema}, public")) db.session.commit() except SQLAlchemyError as e: db.session.rollback() db.session.close() current_app.logger.error(f"Error creating schema {self.schema}: {e.args}") def create_tables(self): """create tables in for schema""" try: db.metadata.create_all(self.get_engine()) except SQLAlchemyError as e: current_app.logger.error(f"💔 Error creating tables for schema {self.schema}: {e.args}") def switch_schema(self): """switch between tenant/public database schema with diagnostics logging""" # Record the desired tenant schema on the active Session so events can use it try: db.session.info["tenant_schema"] = self.schema except Exception: pass # Log the context before switching self._log_db_context("before_switch") try: db.session.execute(text(f'set search_path to "{self.schema}", public')) db.session.commit() except SQLAlchemyError as e: # Rollback on error to avoid InFailedSqlTransaction and log details try: db.session.rollback() except Exception: pass current_app.logger.error( f"Error switching search_path to {self.schema}: {e!r}" ) # Also log context after failure self._log_db_context("after_switch_failed") # Re-raise to let caller decide handling if needed raise # Log the context after successful switch self._log_db_context("after_switch") def migrate_tenant_schema(self): """migrate tenant database schema for new tenant""" # Get the current revision for a database. # NOTE: using popen may have a minor performance impact on the application # you can store it in a different table in public schema and use it from there # may be a faster approach # last_revision = heads(directory="migrations/tenant", verbose=True, resolve_dependencies=False) last_revision = popen("scripts/db_heads.sh -d migrations/tenant").read() print("LAST REVISION") print(last_revision) last_revision = last_revision.splitlines()[-1].split(" ")[0] # creating revision table in tenant schema session = self.get_session() session.execute( text( f'CREATE TABLE "{self.schema}".alembic_version (version_num ' "VARCHAR(32) NOT NULL)" ) ) session.commit() # Insert last revision to alembic_version table session.execute( text( f'INSERT INTO "{self.schema}".alembic_version (version_num) ' "VALUES (:version)" ), {"version": last_revision}, ) session.commit() session.close() def create_tenant_schema(self): """create tenant used for creating new schema and its tables""" self.create_schema() self.create_tables() self.migrate_tenant_schema()