capsul-flask/capsulflask/db.py

131 lines
4.2 KiB
Python
Raw Normal View History

import psycopg2
import re
import sys
from os import listdir
from os.path import isfile, join
from psycopg2 import pool
from flask import current_app
from flask import g
from capsulflask.db_model import DBModel
from capsulflask.shared import my_exec_info_message
def init_app(app, is_running_server):
app.config['PSYCOPG2_CONNECTION_POOL'] = psycopg2.pool.SimpleConnectionPool(
2020-10-30 02:25:29 +00:00
1,
20,
app.config['POSTGRES_CONNECTION_PARAMETERS']
)
# tell the app to clean up the DB connection when shutting down.
app.teardown_appcontext(close_db)
# only run the migrations if we are running the server.
# If we are just running a cli command (e.g. to fix a broken migration 😅), skip it
if not is_running_server:
return
schemaMigrations = {}
schemaMigrationsPath = join(app.root_path, 'schema_migrations')
2020-05-16 04:19:01 +00:00
app.logger.info("loading schema migration scripts from {}".format(schemaMigrationsPath))
for filename in listdir(schemaMigrationsPath):
2020-05-13 05:28:53 +00:00
result = re.search(r"^\d+_(up|down)", filename)
if not result:
2020-05-16 04:19:01 +00:00
app.logger.error(f"schemaVersion {filename} must match ^\\d+_(up|down). exiting.")
continue
2020-05-13 05:28:53 +00:00
key = result.group()
with open(join(schemaMigrationsPath, filename), 'rb') as file:
schemaMigrations[key] = file.read().decode("utf8")
2020-10-30 02:25:29 +00:00
2020-05-10 01:36:14 +00:00
connection = app.config['PSYCOPG2_CONNECTION_POOL'].getconn()
hasSchemaVersionTable = False
actionWasTaken = False
schemaVersion = 0
desiredSchemaVersion = 18
2020-05-10 01:36:14 +00:00
cursor = connection.cursor()
cursor.execute("""
SELECT table_name, table_schema FROM information_schema.tables WHERE table_schema = '{}'
""".format(app.config['DATABASE_SCHEMA']))
rows = cursor.fetchall()
for row in rows:
if row[0] == "schemaversion":
hasSchemaVersionTable = True
if hasSchemaVersionTable == False:
2020-05-16 04:19:01 +00:00
app.logger.info("no table named schemaversion found in the {} schema. running migration 01_up".format(app.config['DATABASE_SCHEMA']))
try:
cursor.execute(schemaMigrations["01_up"])
2020-05-10 01:36:14 +00:00
connection.commit()
except:
2020-05-16 04:19:01 +00:00
app.logger.error("unable to create the schemaversion table because: {}".format(my_exec_info_message(sys.exc_info())))
exit(1)
actionWasTaken = True
cursor.execute("SELECT Version FROM schemaversion")
schemaVersion = cursor.fetchall()[0][0]
if schemaVersion > desiredSchemaVersion:
2020-05-16 04:19:01 +00:00
app.logger.critical("schemaVersion ({}) > desiredSchemaVersion ({}). schema downgrades are not supported yet. exiting.".format(
schemaVersion, desiredSchemaVersion
))
exit(1)
while schemaVersion < desiredSchemaVersion:
migrationKey = "%02d_up" % (schemaVersion+1)
2020-05-16 04:19:01 +00:00
app.logger.info("schemaVersion ({}) < desiredSchemaVersion ({}). running migration {}".format(
schemaVersion, desiredSchemaVersion, migrationKey
))
try:
cursor.execute(schemaMigrations[migrationKey])
2020-05-10 01:36:14 +00:00
connection.commit()
except KeyError:
2020-05-16 04:19:01 +00:00
app.logger.critical("missing schema migration script: {}_xyz.sql".format(migrationKey))
exit(1)
except:
2020-05-16 04:19:01 +00:00
app.logger.critical("unable to execute the schema migration {} because: {}".format(migrationKey, my_exec_info_message(sys.exc_info())))
exit(1)
actionWasTaken = True
schemaVersion += 1
cursor.execute("SELECT Version FROM schemaversion")
versionFromDatabase = cursor.fetchall()[0][0]
if schemaVersion != versionFromDatabase:
2020-05-16 04:19:01 +00:00
app.logger.critical("incorrect schema version value \"{}\" after running migration {}, expected \"{}\". exiting.".format(
versionFromDatabase,
2020-10-30 02:25:29 +00:00
migrationKey,
schemaVersion
))
exit(1)
2020-10-30 02:25:29 +00:00
cursor.close()
2020-05-10 01:36:14 +00:00
app.config['PSYCOPG2_CONNECTION_POOL'].putconn(connection)
2020-05-16 04:19:01 +00:00
app.logger.info("{} current schemaVersion: \"{}\"".format(
("schema migration completed." if actionWasTaken else "schema is already up to date. "), schemaVersion
))
2020-05-10 01:36:14 +00:00
def get_model():
if 'db_model' not in g:
2020-05-10 01:36:14 +00:00
connection = current_app.config['PSYCOPG2_CONNECTION_POOL'].getconn()
cursor = connection.cursor()
g.db_model = DBModel(connection, cursor)
return g.db_model
def close_db(e=None):
db_model = g.pop("db_model", None)
if db_model is not None:
db_model.cursor.close()
current_app.config['PSYCOPG2_CONNECTION_POOL'].putconn(db_model.connection)