from collections import OrderedDict from typing import Optional from importlib import import_module from duckdb import DuckDBPyConnection MAX_MIGRATION = 8 MIGRATIONS = OrderedDict() for i in range(1, MAX_MIGRATION + 1): MIGRATIONS[i] = import_module(f".migration_{i}", __package__).migration def get_version(db: DuckDBPyConnection) -> int: """Get the current database version. Args: db (DuckDBPyConnection): Database connection. Returns: int: Current database version. """ try: return int(db.execute("SELECT MAX(id) FROM migrations").fetchone()[0]) except: return 0 def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_version: Optional[int] = None) -> None: """Migrate the database to a specific version. Args: db (DuckDBPyConnection): Database connection. from_version (Optional[int]): Version to migrate from. If None, the current version is used. to_version (Optional[int]): Version to migrate to. If None, the latest version is used. """ if from_version is None: from_version = get_version(db) + 1 if to_version is None: to_version = max(MIGRATIONS.keys()) if from_version > to_version: raise ValueError("Cannot migrate from a higher version to a lower version.") print(f"Migrating from version {from_version} to {to_version}...") for version in range(from_version, to_version): if version in MIGRATIONS: print(f"Running migration {version}...") MIGRATIONS[version](db) return from_version, to_version