python_arango_ogm.db.pao_migrator

 1import importlib
 2import os
 3from pathlib import Path
 4
 5import arango
 6from arango import AQLQueryExecuteError
 7
 8from python_arango_ogm.db.pao_migration_model import PAOMigrationModel
 9from python_arango_ogm.db.pao_database import PAODatabase
10
11
12class PAOMigrator():
13    def __init__(self, pao_db:PAODatabase, target_path: str = '.'):
14        self.pao_db = pao_db
15        self.app_package = os.getenv('PAO_APP_PACKAGE')
16        if not self.app_package:
17            raise ValueError('PAO_APP_PACKAGE needs to be defined in environment or in a .env file.')
18        self.migration_package = self.app_package + '.migrations'
19
20        app_root = self.app_package.replace('.', '/')
21        p = Path(target_path).joinpath(app_root)
22        self.migration_pathname = p.joinpath("migrations")
23        if not self.migration_pathname.exists():
24            raise RuntimeError(f"Migration path does not exist: {self.migration_pathname}; please run make-migrations.")
25
26
27    def list_migrations(self):
28        try:
29            result = list(PAOMigrationModel.all({"migration_filename": "ASC"}, marshall=False))
30        except AQLQueryExecuteError as e:
31            result = None
32        return result
33
34    def migrate_down(self):
35        migrations = self.list_migrations()
36        if len(migrations):
37            m = migrations[-1]
38            migration = importlib.import_module(f"{self.migration_package}.{m['migration_filename']}")
39            migration.down(self.pao_db.db)
40            PAOMigrationModel.remove_by_key(m['_key'])
41
42    def apply_migrations(self):
43        migrations = [f.stem for f in self.migration_pathname.iterdir() if not f.is_dir() and f.suffix == '.py']
44        print(f"Applying migrations from {self.migration_pathname}", migrations)
45        for m in sorted(migrations):
46            try:
47                if self.__find_migration_record(m):
48                    print(f"Migration already applied; skipping [{m}]")
49                    continue
50            except arango.exceptions.AQLQueryExecuteError:
51                # Collection not there at all:
52                print(f"Collection not found for [{m}]...running migration")
53
54            print(f"Applying migration {m}", type(m))
55            migration = importlib.import_module(f"{self.migration_package}.{m}")
56            # migration = importlib.import_module(m, self.migration_package)
57            migration.up(self.pao_db.db)
58            self.__create_migration_record(m)
59
60    def __find_migration_record(self, migration_filename):
61        return self.pao_db.find_by_attributes("pao_migrations", {
62            "migration_filename": migration_filename
63        })
64
65    def __create_migration_record(self, migration_filename):
66        migration_number, migration_name = migration_filename.split('_', 1)
67        PAOMigrationModel.insert({
68            "migration_number": f"`{int(migration_number)}`",
69            "migration_name": migration_name,
70            "migration_filename": migration_filename,
71        })