diff --git a/lifecycle/migrate.py b/lifecycle/migrate.py index 8eee74adb..68d45dbc3 100755 --- a/lifecycle/migrate.py +++ b/lifecycle/migrate.py @@ -1,12 +1,12 @@ #!/usr/bin/env python """System Migration handler""" -import os from importlib.util import module_from_spec, spec_from_file_location from inspect import getmembers, isclass +from os import environ, system from pathlib import Path from typing import Any -from psycopg import connect +from psycopg import Connection, Cursor, connect from structlog.stdlib import get_logger from authentik.lib.config import CONFIG @@ -16,16 +16,33 @@ ADV_LOCK_UID = 1000 LOCKED = False +class CommandError(Exception): + """Error raised when a system_crit command fails""" + + class BaseMigration: """Base System Migration""" - cur: Any - con: Any + cur: Cursor + con: Connection def __init__(self, cur: Any, con: Any): self.cur = cur self.con = con + def system_crit(self, command: str): + """Run system command""" + LOGGER.debug("Running system_crit command", command=command) + retval = system(command) # nosec + if retval != 0: + raise CommandError("Migration error") + + def fake_migration(self, *app_migration: tuple[str, str]): + """Fake apply a list of migrations, arguments are + expected to be tuples of (app_label, migration_name)""" + for app, _migration in app_migration: + self.system_crit(f"./manage.py migrate {app} {_migration} --fake") + def needs_migration(self) -> bool: """Return true if Migration needs to be run""" return False @@ -82,7 +99,7 @@ if __name__ == "__main__": LOGGER.info("Migration finished applying", migration=sub) release_lock() LOGGER.info("applying django migrations") - os.environ.setdefault("DJANGO_SETTINGS_MODULE", "authentik.root.settings") + environ.setdefault("DJANGO_SETTINGS_MODULE", "authentik.root.settings") wait_for_lock() try: from django.core.management import execute_from_command_line diff --git a/lifecycle/system_migrations/install_id.py b/lifecycle/system_migrations/install_id.py index 28867d8fa..2545bd910 100644 --- a/lifecycle/system_migrations/install_id.py +++ b/lifecycle/system_migrations/install_id.py @@ -4,11 +4,9 @@ from uuid import uuid4 from authentik.lib.config import CONFIG from lifecycle.migrate import BaseMigration -SQL_STATEMENT = """BEGIN TRANSACTION; -CREATE TABLE IF NOT EXISTS authentik_install_id ( +SQL_STATEMENT = """CREATE TABLE IF NOT EXISTS authentik_install_id ( id TEXT NOT NULL -); -COMMIT;""" +);""" class Migration(BaseMigration): @@ -19,19 +17,18 @@ class Migration(BaseMigration): return not bool(self.cur.rowcount) def upgrade(self, migrate=False): - self.cur.execute(SQL_STATEMENT) - self.con.commit() - if migrate: - # If we already have migrations in the database, assume we're upgrading an existing install - # and set the install id to the secret key - self.cur.execute( - "INSERT INTO authentik_install_id (id) VALUES (%s)", (CONFIG.get("secret_key"),) - ) - else: - # Otherwise assume a new install, generate an install ID based on a UUID - install_id = str(uuid4()) - self.cur.execute("INSERT INTO authentik_install_id (id) VALUES (%s)", (install_id,)) - self.con.commit() + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT) + if migrate: + # If we already have migrations in the database, assume we're upgrading an existing install + # and set the install id to the secret key + self.cur.execute( + "INSERT INTO authentik_install_id (id) VALUES (%s)", (CONFIG.get("secret_key"),) + ) + else: + # Otherwise assume a new install, generate an install ID based on a UUID + install_id = str(uuid4()) + self.cur.execute("INSERT INTO authentik_install_id (id) VALUES (%s)", (install_id,)) def run(self): self.cur.execute( diff --git a/lifecycle/system_migrations/otp_merge.py b/lifecycle/system_migrations/otp_merge.py index c3908bffa..7561f0df7 100644 --- a/lifecycle/system_migrations/otp_merge.py +++ b/lifecycle/system_migrations/otp_merge.py @@ -1,10 +1,7 @@ # flake8: noqa -from os import system - from lifecycle.migrate import BaseMigration SQL_STATEMENT = """ -BEGIN TRANSACTION; DELETE FROM django_migrations WHERE app = 'otp_static'; DELETE FROM django_migrations WHERE app = 'otp_totp'; -- Rename tables (static) @@ -15,7 +12,7 @@ ALTER SEQUENCE otp_static_staticdevice_id_seq RENAME TO authentik_stages_authent -- Rename tables (totp) ALTER TABLE otp_totp_totpdevice RENAME TO authentik_stages_authenticator_totp_totpdevice; ALTER SEQUENCE otp_totp_totpdevice_id_seq RENAME TO authentik_stages_authenticator_totp_totpdevice_id_seq; -COMMIT;""" +""" class Migration(BaseMigration): @@ -25,23 +22,24 @@ class Migration(BaseMigration): ) return bool(self.cur.rowcount) - def system_crit(self, command): - retval = system(command) # nosec - if retval != 0: - raise Exception("Migration error") - def run(self): - self.cur.execute(SQL_STATEMENT) - self.con.commit() - self.system_crit( - "./manage.py migrate authentik_stages_authenticator_static 0008_initial --fake" - ) - self.system_crit( - "./manage.py migrate authentik_stages_authenticator_static 0009_throttling --fake" - ) - self.system_crit( - "./manage.py migrate authentik_stages_authenticator_totp 0008_initial --fake" - ) - self.system_crit( - "./manage.py migrate authentik_stages_authenticator_totp 0009_auto_20190420_0723 --fake" - ) + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT) + self.fake_migration( + ( + "authentik_stages_authenticator_static", + "0008_initial", + ), + ( + "authentik_stages_authenticator_static", + "0009_throttling", + ), + ( + "authentik_stages_authenticator_totp", + "0008_initial", + ), + ( + "authentik_stages_authenticator_totp", + "0009_auto_20190420_0723", + ), + ) diff --git a/lifecycle/system_migrations/to_0_10.py b/lifecycle/system_migrations/to_0_10.py index 77ff3f69c..84ab45b39 100644 --- a/lifecycle/system_migrations/to_0_10.py +++ b/lifecycle/system_migrations/to_0_10.py @@ -1,10 +1,7 @@ # flake8: noqa -from os import system - from lifecycle.migrate import BaseMigration SQL_STATEMENT = """ -BEGIN TRANSACTION; DELETE FROM django_migrations WHERE app = 'passbook_stages_prompt'; DROP TABLE passbook_stages_prompt_prompt cascade; DROP TABLE passbook_stages_prompt_promptstage cascade; @@ -25,7 +22,7 @@ DELETE FROM django_migrations WHERE app = 'passbook_flows' AND name = '0008_defa DELETE FROM django_migrations WHERE app = 'passbook_flows' AND name = '0009_source_flows'; DELETE FROM django_migrations WHERE app = 'passbook_flows' AND name = '0010_provider_flows'; DELETE FROM django_migrations WHERE app = 'passbook_stages_password' AND name = '0002_passwordstage_change_flow'; -COMMIT;""" +""" class Migration(BaseMigration): @@ -35,17 +32,14 @@ class Migration(BaseMigration): ) return bool(self.cur.rowcount) - def system_crit(self, command): - retval = system(command) # nosec - if retval != 0: - raise Exception("Migration error") - def run(self): - self.cur.execute(SQL_STATEMENT) - self.con.commit() - self.system_crit("./manage.py migrate passbook_stages_prompt") - self.system_crit("./manage.py migrate passbook_flows 0008_default_flows --fake") - self.system_crit("./manage.py migrate passbook_flows 0009_source_flows --fake") - self.system_crit("./manage.py migrate passbook_flows 0010_provider_flows --fake") - self.system_crit("./manage.py migrate passbook_flows") - self.system_crit("./manage.py migrate passbook_stages_password --fake") + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT) + self.system_crit("./manage.py migrate passbook_stages_prompt") + self.fake_migration( + ("passbook_flows", "0008_default_flows"), + ("passbook_flows", "0009_source_flows"), + ("passbook_flows", "0010_provider_flows"), + ) + self.system_crit("./manage.py migrate passbook_flows") + self.fake_migration(("passbook_stages_password", "")) diff --git a/lifecycle/system_migrations/to_0_13_authentik.py b/lifecycle/system_migrations/to_0_13_authentik.py index ff2088349..8ba702132 100644 --- a/lifecycle/system_migrations/to_0_13_authentik.py +++ b/lifecycle/system_migrations/to_0_13_authentik.py @@ -4,7 +4,7 @@ from redis import Redis from authentik.lib.config import CONFIG from lifecycle.migrate import BaseMigration -SQL_STATEMENT = """BEGIN TRANSACTION; +SQL_STATEMENT = """ ALTER TABLE passbook_audit_event RENAME TO authentik_audit_event; ALTER TABLE passbook_core_application RENAME TO authentik_core_application; ALTER TABLE passbook_core_group RENAME TO authentik_core_group; @@ -92,8 +92,7 @@ ALTER SEQUENCE passbook_stages_prompt_promptstage_validation_policies_id_seq REN UPDATE django_migrations SET app = replace(app, 'passbook', 'authentik'); UPDATE django_content_type SET app_label = replace(app_label, 'passbook', 'authentik'); - -END TRANSACTION;""" +""" class Migration(BaseMigration): @@ -104,18 +103,18 @@ class Migration(BaseMigration): return bool(self.cur.rowcount) def run(self): - self.cur.execute(SQL_STATEMENT) - self.con.commit() - # We also need to clean the cache to make sure no pickeled objects still exist - for db in [ - CONFIG.get("redis.message_queue_db"), - CONFIG.get("redis.cache_db"), - CONFIG.get("redis.ws_db"), - ]: - redis = Redis( - host=CONFIG.get("redis.host"), - port=6379, - db=db, - password=CONFIG.get("redis.password"), - ) - redis.flushall() + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT) + # We also need to clean the cache to make sure no pickeled objects still exist + for db in [ + CONFIG.get("redis.message_queue_db"), + CONFIG.get("redis.cache_db"), + CONFIG.get("redis.ws_db"), + ]: + redis = Redis( + host=CONFIG.get("redis.host"), + port=6379, + db=db, + password=CONFIG.get("redis.password"), + ) + redis.flushall() diff --git a/lifecycle/system_migrations/to_0_14_events..py b/lifecycle/system_migrations/to_0_14_events..py index 4745b8853..b1a0cc727 100644 --- a/lifecycle/system_migrations/to_0_14_events..py +++ b/lifecycle/system_migrations/to_0_14_events..py @@ -1,12 +1,9 @@ # flake8: noqa from lifecycle.migrate import BaseMigration -SQL_STATEMENT = """BEGIN TRANSACTION; -ALTER TABLE authentik_audit_event RENAME TO authentik_events_event; +SQL_STATEMENT = """ALTER TABLE authentik_audit_event RENAME TO authentik_events_event; UPDATE django_migrations SET app = replace(app, 'authentik_audit', 'authentik_events'); -UPDATE django_content_type SET app_label = replace(app_label, 'authentik_audit', 'authentik_events'); - -END TRANSACTION;""" +UPDATE django_content_type SET app_label = replace(app_label, 'authentik_audit', 'authentik_events');""" class Migration(BaseMigration): @@ -17,5 +14,5 @@ class Migration(BaseMigration): return bool(self.cur.rowcount) def run(self): - self.cur.execute(SQL_STATEMENT) - self.con.commit() + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT) diff --git a/lifecycle/system_migrations/to_2021_3_authenticator.py b/lifecycle/system_migrations/to_2021_3_authenticator.py index 3dd895bdb..3b633fef1 100644 --- a/lifecycle/system_migrations/to_2021_3_authenticator.py +++ b/lifecycle/system_migrations/to_2021_3_authenticator.py @@ -1,7 +1,7 @@ # flake8: noqa from lifecycle.migrate import BaseMigration -SQL_STATEMENT = """BEGIN TRANSACTION; +SQL_STATEMENT = """ ALTER TABLE authentik_stages_otp_static_otpstaticstage RENAME TO authentik_stages_authenticator_static_otpstaticstage; UPDATE django_migrations SET app = replace(app, 'authentik_stages_otp_static', 'authentik_stages_authenticator_static'); UPDATE django_content_type SET app_label = replace(app_label, 'authentik_stages_otp_static', 'authentik_stages_authenticator_static'); @@ -13,8 +13,7 @@ UPDATE django_content_type SET app_label = replace(app_label, 'authentik_stages_ ALTER TABLE authentik_stages_otp_validate_otpvalidatestage RENAME TO authentik_stages_authenticator_validate_otpvalidatestage; UPDATE django_migrations SET app = replace(app, 'authentik_stages_otp_validate', 'authentik_stages_authenticator_validate'); UPDATE django_content_type SET app_label = replace(app_label, 'authentik_stages_otp_validate', 'authentik_stages_authenticator_validate'); - -END TRANSACTION;""" +""" class Migration(BaseMigration): @@ -26,5 +25,5 @@ class Migration(BaseMigration): return bool(self.cur.rowcount) def run(self): - self.cur.execute(SQL_STATEMENT) - self.con.commit() + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT) diff --git a/lifecycle/system_migrations/to_2023_1_hibp_remove.py b/lifecycle/system_migrations/to_2023_1_hibp_remove.py index 4c8b9e292..c43f6bb85 100644 --- a/lifecycle/system_migrations/to_2023_1_hibp_remove.py +++ b/lifecycle/system_migrations/to_2023_1_hibp_remove.py @@ -1,10 +1,8 @@ # flake8: noqa from lifecycle.migrate import BaseMigration -SQL_STATEMENT = """BEGIN TRANSACTION; -DROP TABLE "authentik_policies_hibp_haveibeenpwendpolicy"; -DELETE FROM django_migrations WHERE app = 'authentik_policies_hibp'; -END TRANSACTION;""" +SQL_STATEMENT = """DROP TABLE "authentik_policies_hibp_haveibeenpwendpolicy"; +DELETE FROM django_migrations WHERE app = 'authentik_policies_hibp';""" class Migration(BaseMigration): @@ -16,5 +14,5 @@ class Migration(BaseMigration): return bool(self.cur.rowcount) def run(self): - self.cur.execute(SQL_STATEMENT) - self.con.commit() + with self.con.transaction(): + self.cur.execute(SQL_STATEMENT)