import json
import logging
import os
from contextlib import contextmanager
from sqlite3 import OperationalError as SqlLiteOperationalError
from time import perf_counter, sleep
from sqlalchemy import create_engine
from sqlalchemy.exc import OperationalError
from sqlalchemy.sql import and_, select
from sqlalchemy.sql.selectable import Select
from firex_keeper.db_model import firex_run_metadata, firex_tasks, metadata
from firexapp.common import wait_until
from firexapp.events.model import COMPLETE_RUNSTATES, FireXRunMetadata, FireXTask
from firexapp.submit.uid import Uid
logger = logging.getLogger(__name__)
[docs]
class FireXWaitQueryExceeded(Exception):
pass
def task_by_uuid_exp(task_uuid):
return firex_tasks.c.uuid == task_uuid
def task_uuid_complete_exp(task_uuid):
return and_(
firex_tasks.c.uuid == task_uuid, firex_tasks.c.state.in_(COMPLETE_RUNSTATES)
)
def cur_task_by_uuid_exp():
from celery import current_task
if not current_task:
return False
return task_by_uuid_exp(current_task.request.id)
def _custom_json_loads(*args, **kwargs):
# JSON columns can still store ints in sqlite, so pass these values along even though they can't be decoded as
# JSON.
if isinstance(args[0], int):
return args[0]
return json.loads(*args, **kwargs)
def _get_pragmas(use_wal):
pragmas = ["page_size = 4096"]
if use_wal:
pragmas += [
"journal_mode=WAL",
"synchronous=NORMAL",
]
return pragmas
def execute_pragmas(engine, use_wal=False):
dbapi_connection = engine.raw_connection()
try:
cursor = dbapi_connection.cursor()
for pragma in _get_pragmas(use_wal):
cmd = "PRAGMA " + pragma
logger.debug(f"Executing: {cmd}")
cursor.execute(cmd)
cursor.close()
dbapi_connection.commit()
finally:
dbapi_connection.close()
def _db_connection_str(db_file, read_only, is_run_complete=False):
db_conn_str = f"sqlite:///file:{db_file}"
params = {"uri": "true"}
if is_run_complete:
params["immutable"] = "1"
read_only = True
if read_only:
params["mode"] = "ro"
db_conn_str += "?" + "&".join([f"{k}={v}" for k, v in params.items()])
return db_conn_str
def connect_db(
db_file, read_only=False, metadata_to_create=metadata, is_run_complete=False
):
engine = create_engine(
_db_connection_str(db_file, read_only, is_run_complete=is_run_complete),
json_deserializer=_custom_json_loads,
)
if not os.path.exists(db_file):
# WAL should not be used while concurrent read+write NFS access is still possible. Once all reads go through
# keeper process for in-progress runs, WAL is likely preferable for in-progress runs, then after the DB
# should be read-only and therefore safe for direct NFS access.
execute_pragmas(engine, use_wal=False)
logger.info(f"Creating schema for {db_file}")
metadata_to_create.create_all(engine)
logger.info(f"Schema creation complete for {db_file}")
return engine.connect()
def get_db_file_dir_path(logs_dir: str) -> str:
return os.path.join(logs_dir, Uid.debug_dirname, "keeper")
def get_keeper_complete_file_path(logs_dir):
# A way for checking if the keeper DB is complete without inspecting the DB
# file. This can be used to inform connection decisions, like if the DB
# is not expected to change.
return os.path.join(get_db_file_dir_path(logs_dir), ".keeper_complete")
def is_keeper_db_complete(logs_dir):
return os.path.isfile(get_keeper_complete_file_path(logs_dir))
def get_keeper_query_ready_file_path(logs_dir: str) -> str:
return os.path.join(get_db_file_dir_path(logs_dir), ".keeper_query_ready")
def is_keeper_db_query_ready(logs_dir: str) -> bool:
return os.path.isfile(get_keeper_query_ready_file_path(logs_dir))
def get_db_file(logs_dir: str, new=False) -> str:
db_file = os.path.join(get_db_file_dir_path(logs_dir), "firex_run.db")
if new:
assert not os.path.exists(db_file), (
f"Cannot create new DB file, it already exists: {db_file}"
)
db_file_parent = os.path.dirname(db_file)
os.makedirs(db_file_parent, exist_ok=True)
else:
assert os.path.isfile(db_file), f"DB file does not exist: {db_file}"
return db_file
[docs]
@contextmanager
def get_db_manager(logs_dir):
"Get a query-only DB manager for an existing keeper DB file."
conn = connect_db(
get_db_file(logs_dir),
read_only=True,
is_run_complete=is_keeper_db_complete(logs_dir),
)
db_manager = FireXRunDbManager(conn)
try:
yield db_manager
finally:
db_manager.close()
def _row_to_run_metadata(row):
# The first 4 columns from the table make up a FireXRunMetadata.
# Need to set firex_requester explicitly since it's not in the db_mode.firex_run_metadata
# Can't add the Column now as it won't be backward compatible
return FireXRunMetadata(*row[:4], firex_requester=None)
RETRYING_DB_EXCEPTIONS = (OperationalError, SqlLiteOperationalError)
DEFAULT_MAX_RETRY_ATTEMPTS = 20
def retry(
exceptions, max_attempts: int = DEFAULT_MAX_RETRY_ATTEMPTS, retry_delay: int = 1
):
def retry_decorator(func):
def retrying_wrapper(*args, **kwargs):
attempt = 0
while attempt < max_attempts:
attempt += 1
try:
return func(*args, **kwargs)
except exceptions:
if attempt >= max_attempts:
raise
sleep(retry_delay)
return retrying_wrapper
return retry_decorator
[docs]
class FireXRunDbManager:
"""
Read-only operations on a keeper DB connection.
"""
def __init__(self, db_conn):
self.db_conn = db_conn
def does_task_whereclause_exist(self, whereclause):
query = select([firex_tasks.c.uuid]).where(whereclause)
return self.db_conn.execute(query).scalar() is not None
def wait_before_query(self, whereclause, max_wait, error_on_wait_exceeded):
if not self._is_keeper_complete():
start_wait_time = perf_counter()
exists = wait_until(
self.does_task_whereclause_exist, max_wait, 0.5, whereclause
)
if not exists:
msg = f"Wait exceeded {max_wait} seconds for {whereclause} to exist, but it still does not."
if error_on_wait_exceeded:
raise FireXWaitQueryExceeded(msg)
else:
logger.warning(msg)
else:
wait_duration = perf_counter() - start_wait_time
logger.debug(
f"Keeper query waited {wait_duration:.2f} secs for wait query to exist."
)
@retry(RETRYING_DB_EXCEPTIONS)
def query_tasks(
self, exp, wait_for_exp_exist=None, max_wait=15, error_on_wait_exceeded=False
) -> list[FireXTask]:
if wait_for_exp_exist is not None:
self.wait_before_query(wait_for_exp_exist, max_wait, error_on_wait_exceeded)
if isinstance(exp, Select):
select_stmt = exp
else:
select_stmt = select([firex_tasks]).where(exp)
db_result = self.db_conn.execute(select_stmt)
result_tasks = []
for row in db_result:
try:
result_tasks.append(FireXTask(*row))
except TypeError:
logger.exception("Failed transforming %s", row[0])
raise
return result_tasks
@retry(RETRYING_DB_EXCEPTIONS)
def query_run_metadata(self, firex_id) -> FireXRunMetadata:
result = self.db_conn.execute(
select([firex_run_metadata]).where(
firex_run_metadata.c.firex_id == firex_id
)
)
if not result:
raise LookupError(f"Found no run data for {firex_id}")
return next(_row_to_run_metadata(row) for row in result)
def _query_single_run_metadata_row(self):
result = self.db_conn.execute(select([firex_run_metadata]))
rows = [r for r in result]
if len(rows) != 1:
raise RuntimeError(
f"Expected exactly one firex_run_metadata, but found {len(rows)}"
)
return rows[0]
@retry(RETRYING_DB_EXCEPTIONS)
def _is_keeper_complete(self) -> bool:
return self._query_single_run_metadata_row()["keeper_complete"]
@retry(RETRYING_DB_EXCEPTIONS)
def query_single_run_metadata(self) -> FireXRunMetadata:
return _row_to_run_metadata(self._query_single_run_metadata_row())
def close(self):
self.db_conn.close()