57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
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 1
|
|
|
|
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)
|
|
|
|
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.")
|
|
|
|
# dont migrate if already at the target version
|
|
if from_version == to_version:
|
|
return from_version, to_version
|
|
|
|
print(f"Migrating from version {from_version} to {to_version}...")
|
|
for version in range(from_version, to_version + 1):
|
|
if version in MIGRATIONS:
|
|
print(f"Running migration {version}...")
|
|
MIGRATIONS[version](db)
|
|
|
|
return from_version, to_version
|