eveAI/common/utils/database.py

"""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:
            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('<tenant>.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()