exobrain-database 0.1.0__py3-none-any.whl

exobrain/database/migrations/README.md

@@ -0,0 +1 @@
+Generic single-database configuration.

exobrain/database/migrations/__init__.py

@@ -0,0 +1,9 @@
+"""Database migration utilities for multi-tenant organization schemas."""
+
+from exobrain.database.migrations.migration_manager import (
+    configure,
+    downgrade,
+    upgrade,
+)
+
+__all__ = ["configure", "upgrade", "downgrade"]

exobrain/database/migrations/alembic.ini

@@ -0,0 +1,104 @@
+# A generic, single database configuration.
+
+[alembic]
+script_location = %(here)s
+version_locations = %(here)s/versions
+file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+timezone = UTC
+
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python-dateutil library that can be
+# installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the
+# "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to alembic/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "version_path_separator" below.
+
+# version path separator; As mentioned above, this is the character used to split
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
+# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
+# Valid values for version_path_separator are:
+#
+# version_path_separator = :
+# version_path_separator = ;
+# version_path_separator = space
+version_path_separator = :
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

exobrain/database/migrations/env.py

@@ -0,0 +1,198 @@
+import logging
+import os
+import typing as t
+from logging.config import fileConfig
+
+from alembic import context
+from alembic.runtime.environment import NameFilterType
+from sqlalchemy import URL, Connection, Engine, Result, Table, column, create_engine, select, table
+from sqlalchemy.schema import CreateSchema
+from sqlalchemy.sql import text
+from sqlalchemy.sql.schema import SchemaItem
+
+from exobrain.database.model.base import BASE
+
+# Apply logging configuration from alembic.ini
+if context.config.config_file_name:
+    fileConfig(context.config.config_file_name)
+
+logger = logging.getLogger("alembic.runtime.migration")
+
+base_metadata = [BASE.metadata]
+"""Metadata for schema migrations which contains all declared models."""
+
+
+def list_schemas(connection: Connection) -> list[str]:
+    results: Result[t.Any] = connection.execute(
+        select(column("schema_name"))
+        .select_from(table(schema="information_schema", name="schemata"))
+        .where(column("schema_name").like("org_%"))
+    )
+    return [row.schema_name for row in results]
+
+
+class MigrationException(SystemExit):
+    """Custom exception for migration errors."""
+
+
+def _get_alembic_commands() -> tuple[t.Any, ...]:
+    """
+    Get the command-line options passed to the ``alembic`` script.
+    """
+    cmd_opts = context.config.cmd_opts
+    if not cmd_opts or "cmd" not in cmd_opts or not cmd_opts.cmd:
+        logger.warning("No command were passed to the alembic script: %s", repr(cmd_opts))
+        return ()
+    return tuple(cmd_opts.cmd)
+
+
+# Use to generate the schema even if no org has been defined
+def has_autogenerate() -> bool:
+    """
+    Check if the current alembic command is an upgrade or autogenerate.
+
+    Possible commands include:
+
+    - `alembic upgrade`
+    - `alembic revision --autogenerate`
+    - `alembic revision --autogenerate -m "message"`
+    """
+    commands = _get_alembic_commands()
+    is_upgrade = len(commands) > 0 and callable(commands[0]) and commands[0].__name__ == "upgrade"
+    is_autogenerate = len(commands) > 2 and "autogenerate" in commands[2]  # noqa PLR2004
+    return is_upgrade or is_autogenerate
+
+
+def is_public(
+    schema_item: SchemaItem,
+    name: str | None,  # noqa
+    type_: NameFilterType,
+    reflected: bool,  # noqa
+    compare_to: SchemaItem | None,  # noqa
+) -> bool:
+    if type_ == "table" and schema_item.info.get("skip_autogenerate"):
+        return False
+
+    table: Table
+    if type_ == "table":
+        table = schema_item  # type: ignore
+    elif type_ == "foreign_key_constraint":
+        table = schema_item.parent  # type: ignore
+    else:
+        table = schema_item.table  # type: ignore
+
+    return table.schema == "public" or table.info.get("schema_type") == "base"
+
+
+def is_org(
+    schema_item: SchemaItem,
+    name: str | None,
+    type_: NameFilterType,
+    reflected: bool,  # noqa
+    compare_to: SchemaItem | None,
+) -> bool:
+    if type_ == "table" and schema_item.info.get("skip_autogenerate"):
+        return False
+
+    return not is_public(schema_item, name, type_, reflected, compare_to)
+
+
+# for public schemas
+def run_base_migrations_online(connection: Connection) -> None:
+    logger.info("======= Processing public schema =======")
+    args = context.get_x_argument(as_dictionary=True)
+    schema = args.get("schema")
+    if schema and schema != "public":
+        logger.info("Skipping base migrations for schema: %s", schema)
+        return
+
+    context.configure(
+        connection=connection,
+        target_metadata=base_metadata,
+        include_object=is_public,
+        upgrade_token="base_up",  # noqa: S106
+        downgrade_token="base_down",  # noqa: S106
+        compare_type=True,
+    )
+
+    with context.begin_transaction():
+        connection.execute(text("SET search_path TO public"))
+        connection.dialect.default_schema_name = "public"
+        context.run_migrations(schema_type="base")
+        connection.commit()
+
+
+# for organizations schemas
+def run_orgs_migrations_online(connection: Connection) -> None:
+    args = context.get_x_argument(as_dictionary=True)
+    schema = args.get("schema")
+    if schema and schema == "public":
+        logger.info("Skipping organization migrations for schema: %s", schema)
+        return
+    schemas = [schema] if schema else list_schemas(connection)
+
+    # When the database is empty (without organizations),
+    # we create a default schema "org_1" for autogeneration.
+    if has_autogenerate() and not schemas:
+        schemas = ["org_1"]
+
+    for schema in filter(None, schemas):
+        logger.info("======= Processing organization schema: %s =======", schema)
+
+        context.configure(
+            connection=connection,
+            target_metadata=base_metadata,
+            include_object=is_org,
+            upgrade_token="org_up",  # noqa: S106
+            downgrade_token="org_down",  # noqa: S106
+            compare_type=True,
+        )
+
+        with context.begin_transaction():
+            if not connection.dialect.has_schema(connection, schema):
+                connection.execute(CreateSchema(schema, if_not_exists=True))
+                connection.commit()
+            connection.execute(text(f"SET search_path TO {schema}"))
+            connection.dialect.default_schema_name = schema
+            context.run_migrations(schema_type="org")
+            connection.commit()
+
+
+def run_migrations_online() -> None:
+    """
+    Execute database migrations for both base and organization schemas.
+    """
+    sql_engine = os.getenv("SQL_ENGINE", "postgresql+psycopg")
+    sql_host = os.getenv("SQL_HOST", "localhost")
+    sql_database = os.getenv("SQL_DATABASE", "exobrain")
+    sql_user = os.getenv("SQL_USER", "postgres")
+    sql_port = int(os.getenv("SQL_PORT", "5432"))
+    sql_password = os.getenv("SQL_PASSWORD")
+    if not sql_password:
+        raise MigrationException("SQL_PASSWORD environment variable is required.")
+
+    db_url = URL.create(
+        drivername=sql_engine,
+        username=sql_user,
+        password=sql_password,
+        host=sql_host,
+        port=sql_port,
+        database=sql_database,
+    )
+    db_engine: Engine = create_engine(db_url)
+
+    with db_engine.connect() as connection:
+        # If the first argument is a function, and it is a downgrade, we reverse the order of migrations
+        # to ensure that base migrations are applied before organization migrations.
+        commands = _get_alembic_commands()
+        is_downgrade = len(commands) > 0 and callable(commands[0]) and commands[0].__name__ == "downgrade"
+        if is_downgrade:
+            run_orgs_migrations_online(connection)
+            run_base_migrations_online(connection)
+        else:
+            run_base_migrations_online(connection)
+            run_orgs_migrations_online(connection)
+        connection.commit()
+
+
+run_migrations_online()

exobrain/database/migrations/migration_manager.py

@@ -0,0 +1,109 @@
+"""
+Migration manager for Alembic database migrations.
+
+This module provides utilities to run database schema migrations for multi-tenant
+organization schemas using Alembic. Each organization has its own schema
+and migrations can be applied independently to each organization's schema.
+
+Functions:
+    - configure(org_id): Configure Alembic for a specific organization schema.
+    - upgrade(org_id): Upgrade the database schema to the latest version.
+    - downgrade(org_id): Downgrade the database schema to the base version.
+"""
+
+import argparse
+import functools
+from collections import abc
+from pathlib import Path
+from threading import Lock
+from typing import ParamSpec, TypeVar
+from uuid import UUID
+
+from alembic import command
+from alembic.config import Config
+
+HERE = Path(__file__).parent
+"""Path to the current directory of this module, which contains the Alembic configuration."""
+
+_migration_lock = Lock()
+"""Shared lock to ensure thread-safe migration operations."""
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _synchronized(func: abc.Callable[P, T]) -> abc.Callable[P, T]:
+    """
+    Decorator to synchronize access to migration functions.
+
+    This ensures that only one migration operation can run at a time,
+    preventing concurrent migrations from interfering with each other.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        with _migration_lock:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def configure(org_id: UUID | int | str) -> Config:
+    """
+    Create an Alembic configuration for a specific organization schema.
+
+    Args:
+        org_id: The organization identifier.
+
+    Returns:
+        An Alembic Config object configured for the organization's schema.
+
+    Raises:
+        TypeError: If org_id is not a valid type (UUID, int, or numeric string).
+    """
+    if isinstance(org_id, UUID):
+        schema = f"org_{org_id.int}"
+    elif isinstance(org_id, int) or (isinstance(org_id, str) and org_id.isdigit()):
+        schema = f"org_{org_id}"
+    else:
+        error = f"Invalid org_id type: '{type(org_id).__name__}'. Must be UUID, int, or numeric str."
+        raise TypeError(error)
+
+    cmd_opts = argparse.Namespace(x=[f"schema={schema}"])
+
+    path = HERE.joinpath("alembic.ini")
+    if not path.is_file():  # pragma: no cover
+        error = f"Alembic configuration file not found at {path}"
+        raise NotImplementedError(error)
+
+    return Config(path, cmd_opts=cmd_opts)
+
+
+@_synchronized
+def upgrade(org_id: UUID | int | str) -> None:
+    """
+    Upgrade the database schema to the latest version for a specific organization.
+
+    Args:
+        org_id: The organization identifier.
+
+    Raises:
+        TypeError: If org_id is not a valid type (UUID, int, or numeric string).
+    """
+    config = configure(org_id)
+    command.upgrade(config, "head")
+
+
+@_synchronized
+def downgrade(org_id: UUID | int | str) -> None:
+    """
+    Downgrade the database schema to the base version for a specific organization.
+
+    Args:
+        org_id: The organization identifier.
+
+    Raises:
+        TypeError: If org_id is not a valid type (UUID, int, or numeric string).
+    """
+    config = configure(org_id)
+    command.downgrade(config, "head")

exobrain/database/migrations/script.py.mako

@@ -0,0 +1,47 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+"""
+
+from collections.abc import Sequence
+
+from alembic import op
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: str | None = ${repr(down_revision)}
+branch_labels: str | Sequence[str] | None = ${repr(branch_labels)}
+depends_on: str | Sequence[str] | None = ${repr(depends_on)}
+
+
+def upgrade(schema_type: str) -> None:
+    if schema_type == "base":
+        upgrade_base()
+    else:
+        upgrade_organization()
+
+
+def downgrade(schema_type: str) -> None:
+    if schema_type == "base":
+        downgrade_base()
+    else:
+        downgrade_organization()
+
+
+def upgrade_base() -> None:
+    ${context.get("base_up", "pass")}
+
+
+def downgrade_base() -> None:
+    ${context.get("base_down", "pass")}
+
+
+def upgrade_organization() -> None:
+    ${context.get("org_up", "pass")}
+
+
+def downgrade_organization() -> None:
+    ${context.get("org_down", "pass")}

exobrain/database/migrations/versions/2023_08_17_0937-3ed3cf31c9d4_.py

@@ -0,0 +1,161 @@
+"""Add new tables for organizations and roles
+
+Revision ID: 3ed3cf31c9d4
+Revises:
+Create Date: 2023-08-17 09:37:26.315461+00:00
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "3ed3cf31c9d4"
+down_revision: str | None = None
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade(schema_type):
+    if schema_type == "base":
+        return upgrade_base()
+    else:
+        return upgrade_organization()
+
+
+def downgrade(schema_type):
+    if schema_type == "base":
+        return downgrade_base()
+    else:
+        return downgrade_organization()
+
+
+# BASE SCHEMA MIGRATIONS
+
+
+def upgrade_base():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "organizations",
+        sa.Column("name", sa.Text(), nullable=False),
+        sa.Column("description", sa.Text(), nullable=True),
+        sa.Column("emailAdmin", sa.Text(), nullable=False),
+        sa.Column("currency", sa.Enum("DOLLAR", "EURO", name="currency"), nullable=True),
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("created_by", sa.Text(), nullable=False),
+        sa.Column("updated_by", sa.Text(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+        schema="public",
+    )
+    op.create_table(
+        "roles",
+        sa.Column("name", sa.Text(), nullable=False),
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("created_by", sa.Text(), nullable=False),
+        sa.Column("updated_by", sa.Text(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+        schema="public",
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade_base():
+    # First, drop the foreign key constraints in organization schemas
+    connection = op.get_bind()
+    result = connection.execute(
+        sa.text(
+            "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE 'org_%'"
+        )
+    )
+    org_schemas = [row[0] for row in result]
+
+    # Drop constraints in each organization schema that reference the public.roles table
+    for schema in org_schemas:
+        # Constraints referencing roles
+        op.drop_constraint(
+            "roles_users_role_id_fkey",
+            "roles_users",
+            schema=schema,
+            type_="foreignkey",
+        )
+        op.drop_constraint(
+            "roles_action_conf_role_id_fkey",
+            "roles_action_conf",
+            schema=schema,
+            type_="foreignkey",
+        )
+        op.drop_constraint(
+            "roles_exec_conf_role_id_fkey",
+            "roles_exec_conf",
+            schema=schema,
+            type_="foreignkey",
+        )
+        op.drop_constraint(
+            "roles_running_executions_role_id_fkey",
+            "roles_running_executions",
+            schema=schema,
+            type_="foreignkey",
+        )
+        op.drop_constraint(
+            "roles_running_actions_role_id_fkey",
+            "roles_running_actions",
+            schema=schema,
+            type_="foreignkey",
+        )
+
+    # Now proceed with the original downgrade logic
+    op.drop_table("roles", schema="public")
+    op.drop_table("organizations", schema="public")
+
+
+# END BASE SCHEMA MIGRATIONS
+
+
+# ORGANIZATION SCHEMA MIGRATIONS
+
+
+def upgrade_organization():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "users",
+        sa.Column("email", sa.Text(), nullable=False),
+        sa.Column("firstname", sa.Text(), nullable=True),
+        sa.Column("lastname", sa.Text(), nullable=True),
+        sa.Column("displayname", sa.Text(), nullable=True),
+        sa.Column("preferredLanguage", sa.Text(), nullable=True),
+        sa.Column("officeLocation", sa.Text(), nullable=True),
+        sa.Column("jobTitle", sa.Text(), nullable=True),
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("created_by", sa.Text(), nullable=False),
+        sa.Column("updated_by", sa.Text(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_table(
+        "roles_users",
+        sa.Column("role_id", sa.UUID(), nullable=False),
+        sa.Column("users_id", sa.UUID(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["role_id"], ["public.roles.id"], onupdate="CASCADE", ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(["users_id"], ["users.id"], onupdate="CASCADE", ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("role_id", "users_id"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade_organization():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("roles_users")
+    op.drop_table("users")
+    # ### end Alembic commands ###
+
+
+# END ORGANIZATION SCHEMA MIGRATIONS