94 lines
		
	
	
		
			2.9 KiB
		
	
	
	
		
			Python
		
	
	
		
			Executable File
		
	
	
	
	
			
		
		
	
	
			94 lines
		
	
	
		
			2.9 KiB
		
	
	
	
		
			Python
		
	
	
		
			Executable File
		
	
	
	
	
| #!/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 pathlib import Path
 | |
| from typing import Any
 | |
| 
 | |
| from psycopg2 import connect
 | |
| from structlog.stdlib import get_logger
 | |
| 
 | |
| from authentik.lib.config import CONFIG
 | |
| 
 | |
| LOGGER = get_logger()
 | |
| ADV_LOCK_UID = 1000
 | |
| LOCKED = False
 | |
| 
 | |
| 
 | |
| class BaseMigration:
 | |
|     """Base System Migration"""
 | |
| 
 | |
|     cur: Any
 | |
|     con: Any
 | |
| 
 | |
|     def __init__(self, cur: Any, con: Any):
 | |
|         self.cur = cur
 | |
|         self.con = con
 | |
| 
 | |
|     def needs_migration(self) -> bool:
 | |
|         """Return true if Migration needs to be run"""
 | |
|         return False
 | |
| 
 | |
|     def run(self):
 | |
|         """Run the actual migration"""
 | |
| 
 | |
| 
 | |
| def wait_for_lock():
 | |
|     """lock an advisory lock to prevent multiple instances from migrating at once"""
 | |
|     LOGGER.info("waiting to acquire database lock")
 | |
|     curr.execute("SELECT pg_advisory_lock(%s)", (ADV_LOCK_UID,))
 | |
|     # pylint: disable=global-statement
 | |
|     global LOCKED
 | |
|     LOCKED = True
 | |
| 
 | |
| 
 | |
| def release_lock():
 | |
|     """Release database lock"""
 | |
|     if not LOCKED:
 | |
|         return
 | |
|     curr.execute("SELECT pg_advisory_unlock(%s)", (ADV_LOCK_UID,))
 | |
| 
 | |
| 
 | |
| if __name__ == "__main__":
 | |
| 
 | |
|     conn = connect(
 | |
|         dbname=CONFIG.y("postgresql.name"),
 | |
|         user=CONFIG.y("postgresql.user"),
 | |
|         password=CONFIG.y("postgresql.password"),
 | |
|         host=CONFIG.y("postgresql.host"),
 | |
|         port=int(CONFIG.y("postgresql.port")),
 | |
|     )
 | |
|     curr = conn.cursor()
 | |
|     try:
 | |
|         for migration in Path(__file__).parent.absolute().glob("system_migrations/*.py"):
 | |
|             spec = spec_from_file_location("lifecycle.system_migrations", migration)
 | |
|             mod = module_from_spec(spec)
 | |
|             # pyright: reportGeneralTypeIssues=false
 | |
|             spec.loader.exec_module(mod)
 | |
| 
 | |
|             for name, sub in getmembers(mod, isclass):
 | |
|                 if name != "Migration":
 | |
|                     continue
 | |
|                 migration = sub(curr, conn)
 | |
|                 if migration.needs_migration():
 | |
|                     wait_for_lock()
 | |
|                     LOGGER.info("Migration needs to be applied", migration=sub)
 | |
|                     migration.run()
 | |
|                     LOGGER.info("Migration finished applying", migration=sub)
 | |
|                     release_lock()
 | |
|         LOGGER.info("applying django migrations")
 | |
|         os.environ.setdefault("DJANGO_SETTINGS_MODULE", "authentik.root.settings")
 | |
|         wait_for_lock()
 | |
|         try:
 | |
|             from django.core.management import execute_from_command_line
 | |
|         except ImportError as exc:
 | |
|             raise ImportError(
 | |
|                 "Couldn't import Django. Are you sure it's installed and "
 | |
|                 "available on your PYTHONPATH environment variable? Did you "
 | |
|                 "forget to activate a virtual environment?"
 | |
|             ) from exc
 | |
|         execute_from_command_line(["", "migrate"])
 | |
|     finally:
 | |
|         release_lock()
 | 
