Snippets for SQLAlchemy scripts.

extensions.py

"""Implement custom logics over SQLAlchemy SQL engines.

SQLite Extensions: https://github.com/nalgeon/sqlean
Reference: https://docs.sqlalchemy.org/en/14/core/engines.html
"""

# IMPORTS

import sys
import sqlite3
from pathlib import Path

import sqlalchemy as sql

# CONFIGS

SQLITE_EXTENSIONS = ["math", "stats"]

_PARENT = Path(__file__).resolve().parent
EXTENSIONS_FOLDER = _PARENT / "extensions"

# HELPERS


def get_extension(name: str) -> str:
    """Get the extension path from its name."""
    # NOTE: function only tested on windows so far
    if sys.platform == "linux" or sys.platform == "linux2":
        return str(EXTENSIONS_FOLDER / f"{name}.so")
    if sys.platform == "darwin":
        return str(EXTENSIONS_FOLDER / f"{name}.dylib")
    if sys.platform == "win32":
        return str(EXTENSIONS_FOLDER / f"{name}.dll")
    raise ValueError(f"Unknown extension platform: {sys.platform}")


# FUNCTIONS


def create_engine(*args, **kwargs) -> sql.engine.Engine:
    """Create a new SQL engine with custom logics."""
    engine = sql.create_engine(*args, **kwargs)

    @sql.event.listens_for(engine, "connect")
    def on_engine_connect(dbapi_connection, _):
        """Hook on connect to apply custom logics."""
        # apply logic only for SQLite database connection
        if isinstance(dbapi_connection, sqlite3.Connection):
            dbapi_connection.enable_load_extension(True)
            for extension in SQLITE_EXTENSIONS:
                extension_path = get_extension(extension)
                dbapi_connection.load_extension(extension_path)

    return engine