"""Connect to the database and provide a session scope for queries."""
from os import path, makedirs
from contextlib import contextmanager
import platformdirs
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine, inspect as sa_inspect
from sqlalchemy.pool import NullPool
from autowisp.database.data_model.base import DataModelBase
from autowisp.database.frozen_row import FrozenRow
from autowisp.database.migrate import (
check_project_schema,
create_project_schema,
migrate_project,
)
from autowisp.database.initialize_data_reduction_structure import (
get_default_data_reduction_structure,
)
from autowisp.database.initialize_light_curve_structure import (
get_default_light_curve_structure,
)
from autowisp.exceptions import DatabaseError
_db_engine = None
# pylint false positive - Session is actually a class name.
# pylint: disable=invalid-name
_Session = None # sessionmaker(db_engine, expire_on_commit=False)
# pylint: enable=invalid-name
_project_home = None
DB_URL_FNAME = "autowisp_db.url" # pylint: disable=invalid-name
"""
Filename (relative to project home) where a non-SQLite connection URL is stored.
When a project is initialised with a centralised database (MySQL, MariaDB,
etc.) the connection URL is written to this file so that subsequent calls to
:func:`set_project_home` with only the directory path can reconnect without
requiring the caller to supply the URL again.
"""
[docs]
def get_db_engine():
"""Return the database engine."""
print(f"Returning engine {_db_engine!r}")
return _db_engine
[docs]
@contextmanager
def start_db_session():
"""Context manager to start a database session."""
with _Session.begin() as db_session: # pylint: disable=no-member
yield db_session
[docs]
def get_project_home():
"""Return the project home directory currently being used."""
return _project_home
[docs]
def snapshot_row(orm_obj, *, exclude=()):
"""Freeze all mapped columns of a live ORM instance into a FrozenRow.
Must be called while ``orm_obj`` is still attached/loaded (i.e.
inside the ``start_db_session()`` block that produced it), otherwise
touching an expired column would raise ``DetachedInstanceError``.
Args:
orm_obj: A SQLAlchemy ORM instance.
exclude(Iterable[str]): Column keys to omit (e.g. large or
sensitive columns).
Returns:
FrozenRow: Snapshot of the instance's column values, detached
from the session and safe to pickle.
"""
mapper = sa_inspect(orm_obj).mapper
exclude = set(exclude)
return FrozenRow(
table=mapper.local_table.name,
columns={
attr.key: getattr(orm_obj, attr.key)
for attr in mapper.column_attrs
if attr.key not in exclude
},
)
[docs]
def initialize_cmdline_database():
"""Initialize the current database HDF5 structure tables."""
create_project_schema(_db_engine)
with start_db_session() as db_session:
db_session.add(get_default_data_reduction_structure())
db_session.add(get_default_light_curve_structure(db_session))
[docs]
def set_project_home(
project_home,
db_url=None,
new_project=False,
*,
migrate=False,
assume_backed_up=False,
):
"""
Set the database engine and session for the given project home.
On first use with a non-SQLite ``db_url`` the URL is persisted to
``<project_home>/autowisp_db.url`` so that subsequent calls with only
``project_home`` reconnect to the same database automatically.
Args:
project_home:
Directory used as the project home. For SQLite (the default), the
database file ``autowisp.db`` is created here. For centralised
databases the directory is still used for other project files (HDF5
products, etc.). Pass ``None`` to use the platform-appropriate user
data directory.
db_url:
SQLAlchemy connection URL. When omitted (or ``None``) the function
first checks for a previously saved URL in
``<project_home>/autowisp_db.url``; if none is found it falls back
to an SQLite database in ``project_home``:
``sqlite:///<project_home>/autowisp.db?timeout=100&uri=true``.
To connect to a centralised server pass the full URL, e.g.:
``"mysql+pymysql://user:password@host:3306/dbname"``
``"mariadb+pymysql://user:password@host:3306/dbname"``
Passing an explicit URL raises an error if a saved URL is found.
new_project:
Pass ``True`` when this call is creating a project rather than
opening an existing one. The target database must then contain
none of the AutoWISP tables, since project creation goes on to
drop and recreate them (see ``initialize_database``) -- which
on a centralised server would silently destroy whichever
project already lives in that database. Independent of
backend: it equally catches an SQLite project home that
already holds a database.
migrate:
Bring the schema up to date rather than merely requiring that it
already is. Off by default because every process opening a
project calls this, including every pipeline worker, and
concurrent schema changes from dozens of workers is exactly what
must not happen. Set it only where a single process is known to
be in charge: the ``wisp-migrate`` command, the browser
interface selecting a project, and the main process of a
pipeline run.
assume_backed_up:
Passed through to :func:`migrate_project`; only meaningful with
``migrate``. Required to migrate a server database, which cannot
be copied aside automatically.
Returns:
dict or None: The result of :func:`migrate_project` when
``migrate`` is set and the project already existed, otherwise
None.
Raises:
DatabaseError: If ``new_project`` is set and the target database
already contains AutoWISP tables. Nothing is written or
dropped in that case: the URL file is persisted only after
this check passes.
"""
global _db_engine, _Session, _project_home # pylint: disable=global-statement
# print(f"Setting project home to {project_home!r}")
if _db_engine is not None:
_db_engine.dispose()
if project_home is None:
project_home = platformdirs.user_data_dir("autowisp")
else:
assert path.isdir(
project_home
), f"Project home {project_home!r} is not a directory."
# Ensure directory exists
makedirs(project_home, exist_ok=True)
_project_home = path.abspath(project_home)
url_file = path.join(_project_home, DB_URL_FNAME)
persist_db_url = None
if db_url is not None:
assert not path.exists(url_file), (
f"Attempting to set a new db_url in {_project_home!r} which already"
f" contains {url_file!r}"
)
# Persisted (further down) only once the target database has been
# accepted, so a rejected one leaves no half-created project whose
# URL file points at somebody else's data.
persist_db_url = db_url
elif path.exists(url_file):
with open(url_file, encoding="utf-8") as fobj:
db_url = fobj.read().strip()
engine_kwargs = {
"echo": False,
"pool_pre_ping": True,
"pool_recycle": 3600,
}
if db_url is None:
db_path = path.join(_project_home, "autowisp.db")
db_url = f"sqlite:///{path.abspath(db_path)}?timeout=600&uri=true"
if db_url.startswith("sqlite"):
engine_kwargs["poolclass"] = NullPool
engine_kwargs["connect_args"] = {"timeout": 600}
# Built into a local first: an engine rejected below must not become
# the module-wide one, or the caller would be left connected to the
# database it was just refused.
engine = create_engine(db_url, **engine_kwargs)
existing_tables = set(sa_inspect(engine).get_table_names())
already_present = set(DataModelBase.metadata.tables) & existing_tables
if new_project and already_present:
# The URL can carry a password, so report the sanitised form.
target = engine.url.render_as_string(hide_password=True)
engine.dispose()
raise DatabaseError(
f"Refusing to create a new project in {_project_home!r}: its "
f"database ({target}) already contains "
f"{len(already_present)} AutoWISP table(s), including "
f"{', '.join(sorted(already_present)[:5])}. Creating the "
"project would drop them, destroying whatever project they "
"belong to. Point the new project at an empty database, or "
"delete the existing project first."
)
_db_engine = engine
_Session = sessionmaker(_db_engine, expire_on_commit=False)
if persist_db_url is not None:
# Lets future calls with only project_home reconnect here.
with open(url_file, "w", encoding="utf-8") as fobj:
fobj.write(persist_db_url)
if not already_present:
initialize_cmdline_database()
return None
if migrate:
return migrate_project(_db_engine, assume_backed_up=assume_backed_up)
# Read-only on purpose: every process opening a project runs this,
# including every pipeline worker, so it must never issue DDL. Migrating
# is migrate_project()'s job -- see autowisp.database.migrate.
check_project_schema(_db_engine)
return None