Refactoring part 1

Some changes for workers, but stopped due to refactoring

common/utils/app_hooks.py

@@ -0,0 +1,12 @@
+from flask import session
+
+
+def make_session_permanent():
+    session.permanent = True  # Refresh the session timeout on every request
+    # if 'user_id' in session and session.get('was_authenticated'):
+    #     if session.modified:  # Check if the session was modified
+    #         session['was_authenticated'] = True
+    # else:
+    #     session.pop('user_id', None)  # Clear session
+    #     session.pop('was_authenticated', None)
+    #     return redirect(url_for('login'))  # Redirect to login page if session expired

common/utils/database.py

@@ -0,0 +1,91 @@
+"""Database related functions"""
+from os import popen
+from sqlalchemy import text
+from sqlalchemy.schema import CreateSchema
+from sqlalchemy.exc import InternalError
+from sqlalchemy.orm import sessionmaker, scoped_session
+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)
+
+    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.execute(text(f"CREATE EXTENSION IF NOT EXISTS pgvector SCHEMA {self.schema}"))
+            db.session.commit()
+        except InternalError 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"""
+        db.metadata.create_all(self.get_engine())
+
+    def switch_schema(self):
+        """switch between tenant/public database schema"""
+        db.session.execute(text(f'set search_path to "{self.schema}"'))
+        db.session.commit()
+
+    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(".venv/bin/flask db heads -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()

common/utils/middleware.py

@@ -0,0 +1,28 @@
+"""Middleware for the API
+
+for handling tenant requests
+"""
+
+from flask_security import current_user
+from flask import session
+
+from .database import Database
+
+
+def mw_before_request():
+    """Before request
+
+    switch tenant schema
+    """
+
+    tenant_id = session['tenant']['id']
+    if not tenant_id:
+        return {"message": "You are not logged into any tenant"}, 403
+
+    # user = User.query.get(current_user.id)
+    if current_user.has_roles(['Super User']) or current_user.tenant_id == tenant_id:
+        Database(tenant_id).switch_schema()
+    else:
+        return {"message": "You are not a member of this tenant"}, 403
+
+

common/utils/security.py

@@ -0,0 +1,18 @@
+from flask import session
+from common.models import User, Tenant
+
+
+# Definition of Trigger Handlers
+def set_tenant_session_data(sender, user, **kwargs):
+    tenant = Tenant.query.filter_by(id=user.tenant_id).first()
+    session['tenant'] = tenant.to_dict()
+    session['default_language'] = tenant.default_language
+    session['default_embedding_model'] = tenant.default_embedding_model
+    session['default_llm_model'] = tenant.default_llm_model
+
+
+def clear_tenant_session_data(sender, user, **kwargs):
+    session.pop('tenant', None)
+    session.pop('default_language', None)
+    session.pop('default_embedding_model', None)
+    session.pop('default_llm_model', None)