Source code for palaestrai.store.database_util

from __future__ import annotations
from typing import Optional

import time
import logging

from sqlalchemy import create_engine
from sqlalchemy.sql.expression import text
from sqlalchemy.exc import ProgrammingError, OperationalError
from sqlalchemy_utils import database_exists, create_database

from .database_model import Model

LOG = logging.getLogger(__name__)

# Default chunk_time_interval. Might become configurable at some point:
TIMESCALEDB_DEFAULT_CHUNK_SIZE_INTERVAL = 512


def _create_timescaledb_extension(engine):
    """Create the timescaledb extension.

    :param engine: The database engine.
    """
    with engine.begin() as conn:
        conn.execute("CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE;")
    timescale_tables = {
        "world_states",
        "muscle_actions",
    }
    with engine.begin() as conn:
        for tbl in timescale_tables:
            cmd = (
                f"SELECT * FROM create_hypertable("
                f"'{tbl}', "  # Table name
                f"'id', "  # Primary partitioning column
                f"chunk_time_interval => "
                f"{TIMESCALEDB_DEFAULT_CHUNK_SIZE_INTERVAL}, "
                f"if_not_exists => TRUE)"
            )
            res = conn.execute(text(cmd))
            LOG.debug(
                'Result of executing "%s" during setup: %s',
                cmd,
                res.fetchall(),
            )
            res.close()
    LOG.info(
        "Created TimescaleDB hypertables: %s, set 'chunk_time_interval' "
        "parameter to %d. HINT: The chunk_time_interval should be chosen such "
        "that all active chunks of all your hypertables fit in 25%% of your "
        "RAM. You can change the value with TimescaleDB's "
        "set_chunk_time_interval() function.",
        ", ".join(timescale_tables),
        TIMESCALEDB_DEFAULT_CHUNK_SIZE_INTERVAL,
    )




def setup_database(uri: Optional[str] = None):
    """Creates the database from the current model in one go.

    Parameters
    ----------
    uri : str
        The complete database connection URI.
    """
    if not uri:
        from palaestrai.core import RuntimeConfig

        uri = RuntimeConfig().store_uri
    engine = create_engine(uri)
    while not database_exists(uri):
        i = 1
        if i > 3:  # Hardcoded max tries. No real reason to configure this.
            LOG.critical(
                "Could not create the database. See errors above for more "
                "details. Giving up now."
            )
            raise RuntimeError("Could not create database")
        try:
            create_database(uri)
        except OperationalError as e:
            try:
                import psycopg2.errors

                if isinstance(e.orig, psycopg2.errors.ObjectInUse):
                    LOG.warning(
                        "Could not create database because the template was "
                        "in use. Retrying in %d seconds.",
                        i,
                    )
                    time.sleep(i)
                else:
                    break
            except ImportError:
                pass
        except ProgrammingError as e:
            LOG.error(
                "There was an error creating the database. I will continue "
                "and hope for the best. The error was: %s",
                e,
            )
        i += 1

    with engine.begin() as conn:
        try:
            Model.metadata.create_all(engine)
        except ProgrammingError as e:
            LOG.error("Could not create database: %s" % e)
            raise e
        try:
            from midas.tools.palaestrai.database_view import (  # type: ignore
                create_midas_views,
            )

            if engine.url.drivername == "psycopg2":
                # type: ignore
                create_midas_views(conn)
        except ModuleNotFoundError:
            pass  # Ok, don't create specific views if the tools are not pres.
    try:
        _create_timescaledb_extension(engine)
    except OperationalError as e:
        LOG.warning(
            "Could not create extension timescaledb and create hypertables: "
            "%s. "
            "Your database setup might lead to noticeable slowdowns with "
            "larger experiment runs. Please upgrade to PostgreSQL with "
            "TimescaleDB for the best performance." % e
        )