Source code for palaestrai.store.session

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Dict, Tuple

import os

import jsonpickle
import jsonpickle.ext.numpy as jsonpickle_numpy

import sqlalchemy
import sqlalchemy.exc
import sqlalchemy.orm
import sqlalchemy.engine
from sqlite3 import Connection as SQLite3Connection

from palaestrai.core.runtime_config import RuntimeConfig

if TYPE_CHECKING:
    import sqlalchemy.engine
    import sqlalchemy.orm.session


LOG = logging.getLogger(__name__)
_session_makers: Dict[
    int, Tuple[sqlalchemy.engine.Engine, sqlalchemy.orm.session.sessionmaker]
] = {}


def _get_session_maker() -> sqlalchemy.orm.session.sessionmaker:
    """Internal getter ensure one session maker per process"""
    pid = os.getpid()
    if (
        pid in _session_makers
        and str(_session_makers[pid][0].url) == RuntimeConfig().store_uri
    ):
        LOG.debug(
            "Returning existing DB session maker for PID %d: %s",
            pid,
            _session_makers[pid][1],
        )
        return _session_makers[pid][1]

    jsonpickle_numpy.register_handlers()
    jsonpickle.set_preferred_backend("simplejson")
    jsonpickle.set_encoder_options("simplejson", ignore_nan=True)

    db_engine = sqlalchemy.create_engine(
        RuntimeConfig().store_uri,
        json_serializer=jsonpickle.dumps,
        json_deserializer=jsonpickle.loads,
    )
    db_session_maker = sqlalchemy.orm.sessionmaker()
    db_session_maker.configure(bind=db_engine)

    _session_makers[pid] = (db_engine, db_session_maker)
    LOG.debug(
        "Created new DB session maker for PID %d; state dict is now: %s",
        pid,
        _session_makers,
    )
    return db_session_maker


@sqlalchemy.event.listens_for(sqlalchemy.engine.Engine, "connect")
def _set_sqlite_pragma(dbapi_connection, connection_record):
    if not isinstance(dbapi_connection, SQLite3Connection):
        return
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON;")
    # WAL lets a reader (e.g. the REST API) and a writer (e.g. the log
    # handler) work on the same database file concurrently; busy_timeout
    # makes a connection wait instead of immediately raising
    # "database is locked"; synchronous=NORMAL is the recommended,
    # crash-safe pairing with WAL. journal_mode=WAL is a harmless no-op on
    # a pure ``:memory:`` connection, so this is safe for tests.
    cursor.execute("PRAGMA journal_mode=WAL;")
    cursor.execute("PRAGMA busy_timeout=5000;")
    cursor.execute("PRAGMA synchronous=NORMAL;")
    cursor.close()




def Session() -> sqlalchemy.orm.Session:
    """Creates a new, connected database session to run queries on.

    This is a convenience function that creates and returns a new, opened
    database session. It uses the access data provided by
    ::`RuntimeConfig.store_uri`. It exists in order to facilitate working
    with the store, such as this::

        from palaestrai.store import Session, database_model as dbm

        session = Session()
        q = session.query(dbm.Experiment)

    Returns:
    --------

    sqlalchemy.orm.Session
        The initialized, opened database session.
    """
    db_session_maker = _get_session_maker()
    return db_session_maker()