matrix-alertbot/matrix_alertbot/storage.py

127 lines
4.2 KiB
Python
Raw Normal View History

2019-09-25 14:31:56 +02:00
import logging
2021-01-10 04:30:07 +01:00
from typing import Any, Dict
2019-09-25 14:26:29 +02:00
# The latest migration version of the database.
#
# Database migrations are applied starting from the number specified in the database's
# `migration_version` table + 1 (or from 0 if this table does not yet exist) up until
# the version specified here.
#
# When a migration is performed, the `migration_version` table should be incremented.
latest_migration_version = 0
2019-09-25 14:26:29 +02:00
2019-09-25 14:31:56 +02:00
logger = logging.getLogger(__name__)
2019-09-25 14:26:29 +02:00
2021-01-10 04:30:07 +01:00
class Storage:
def __init__(self, database_config: Dict[str, str]):
"""Setup the database.
2019-09-25 14:26:29 +02:00
Runs an initial setup or migrations depending on whether a database file has already
been created.
2019-09-25 14:26:29 +02:00
Args:
database_config: a dictionary containing the following keys:
* type: A string, one of "sqlite" or "postgres".
* connection_string: A string, featuring a connection string that
be fed to each respective db library's `connect` method.
2019-09-25 14:26:29 +02:00
"""
self.conn = self._get_database_connection(
database_config["type"], database_config["connection_string"]
)
self.cursor = self.conn.cursor()
self.db_type = database_config["type"]
2019-09-25 14:26:29 +02:00
# Try to check the current migration version
migration_level = 0
try:
self._execute("SELECT version FROM migration_version")
row = self.cursor.fetchone()
migration_level = row[0]
except Exception:
2019-09-25 14:26:29 +02:00
self._initial_setup()
finally:
if migration_level < latest_migration_version:
self._run_migrations(migration_level)
logger.info(f"Database initialization of type '{self.db_type}' complete")
2021-01-10 04:30:07 +01:00
def _get_database_connection(
self, database_type: str, connection_string: str
) -> Any:
"""Creates and returns a connection to the database"""
if database_type == "sqlite":
import sqlite3
# Initialize a connection to the database, with autocommit on
return sqlite3.connect(connection_string, isolation_level=None)
elif database_type == "postgres":
import psycopg2
conn = psycopg2.connect(connection_string)
# Autocommit on
conn.set_isolation_level(0)
return conn
2019-09-25 14:26:29 +02:00
2021-01-10 04:30:07 +01:00
def _initial_setup(self) -> None:
2019-09-25 14:26:29 +02:00
"""Initial setup of the database"""
2019-09-25 14:31:56 +02:00
logger.info("Performing initial database setup...")
2019-09-25 14:26:29 +02:00
# Set up the migration_version table
self._execute(
"""
CREATE TABLE migration_version (
version INTEGER PRIMARY KEY
)
"""
)
2019-09-25 14:26:29 +02:00
# Initially set the migration version to 0
self._execute(
"""
INSERT INTO migration_version (
version
) VALUES (?)
""",
(0,),
2020-08-10 00:02:07 +02:00
)
2019-09-25 14:26:29 +02:00
# Set up any other necessary database tables here
2019-09-25 14:31:56 +02:00
logger.info("Database setup complete")
2021-01-10 04:30:07 +01:00
def _run_migrations(self, current_migration_version: int) -> None:
"""Execute database migrations. Migrates the database to the
`latest_migration_version`.
Args:
current_migration_version: The migration version that the database is
2021-01-10 04:30:07 +01:00
currently at.
"""
logger.debug("Checking for necessary database migrations...")
# if current_migration_version < 1:
# logger.info("Migrating the database from v0 to v1...")
#
# # Add new table, delete old ones, etc.
#
# # Update the stored migration version
# self._execute("UPDATE migration_version SET version = 1")
#
# logger.info("Database migrated to v1")
2021-01-10 04:30:07 +01:00
def _execute(self, *args) -> None:
"""A wrapper around cursor.execute that transforms placeholder ?'s to %s for postgres.
This allows for the support of queries that are compatible with both postgres and sqlite.
Args:
args: Arguments passed to cursor.execute.
"""
if self.db_type == "postgres":
self.cursor.execute(args[0].replace("?", "%s"), *args[1:])
else:
self.cursor.execute(*args)