276 lines
9.9 KiB
Python
276 lines
9.9 KiB
Python
"""
|
|
control/tenant_migrate.py
|
|
-------------------------
|
|
Per-tenant migration runner (MT-2).
|
|
|
|
Runs the EXISTING tenant Alembic chain (migrations/versions) against each
|
|
tenant's own database via the standalone migrations_tenant/env.py. Records the
|
|
applied head on tenants.alembic_head and writes a ProvisioningJob row per run.
|
|
|
|
The existing migrations are re-run-safe (INFORMATION_SCHEMA guards, Rule 14),
|
|
so upgrading an already-current tenant (e.g. tenant-zero / LT, whose DB is
|
|
already at head) is a no-op.
|
|
|
|
Usage:
|
|
python -m control.tenant_migrate upgrade --tenant all
|
|
python -m control.tenant_migrate upgrade --tenant lts # id or slug
|
|
python -m control.tenant_migrate current --tenant all
|
|
python -m control.tenant_migrate heads
|
|
|
|
`upgrade_tenant()` is also imported by the MT-3 provisioning service to bring a
|
|
freshly-created tenant DB up to head.
|
|
"""
|
|
|
|
import argparse
|
|
import contextlib
|
|
import io
|
|
import os
|
|
import sys
|
|
from collections import namedtuple
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from alembic.script import ScriptDirectory
|
|
from alembic.runtime.migration import MigrationContext
|
|
from sqlalchemy import create_engine
|
|
|
|
from control.base import control_session
|
|
from control.models import Tenant, ProvisioningJob
|
|
from control.time_utils import now_eastern
|
|
|
|
# Repo layout: control/ and migrations/ are siblings at the repo root.
|
|
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
_DEFAULT_SCRIPT_LOCATION = os.path.join(_REPO_ROOT, 'migrations_tenant')
|
|
_DEFAULT_VERSION_LOCATIONS = os.path.join(_REPO_ROOT, 'migrations', 'versions')
|
|
|
|
# Everything except a deleted tenant should be kept schema-current.
|
|
MIGRATABLE_STATUSES = ('provisioning', 'active', 'suspended')
|
|
|
|
# Squashed baseline that builds the full schema (migrations/versions root).
|
|
BASELINE_REVISION = '0003_add_user_active'
|
|
|
|
TenantRef = namedtuple('TenantRef', ['id', 'slug', 'db_uri'])
|
|
|
|
|
|
def _make_config(db_uri, script_location=None, version_locations=None):
|
|
cfg = Config()
|
|
cfg.set_main_option('script_location', script_location or _DEFAULT_SCRIPT_LOCATION)
|
|
cfg.set_main_option('version_locations', version_locations or _DEFAULT_VERSION_LOCATIONS)
|
|
# Escape % so ConfigParser interpolation doesn't choke on URL-encoded passwords.
|
|
cfg.set_main_option('sqlalchemy.url', db_uri.replace('%', '%%'))
|
|
return cfg
|
|
|
|
|
|
def chain_head(script_location=None, version_locations=None):
|
|
"""Return the head revision id of the tenant migration chain."""
|
|
cfg = _make_config('sqlite://', script_location, version_locations)
|
|
script = ScriptDirectory.from_config(cfg)
|
|
heads = script.get_heads()
|
|
return heads[0] if len(heads) == 1 else ','.join(sorted(heads))
|
|
|
|
|
|
def current_revision(db_uri):
|
|
"""Return the alembic revision currently stamped on the given database."""
|
|
engine = create_engine(db_uri, future=True)
|
|
try:
|
|
with engine.connect() as conn:
|
|
return MigrationContext.configure(conn).get_current_revision()
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
def upgrade_tenant(tenant, script_location=None, version_locations=None, record_job=True):
|
|
"""Upgrade one tenant's database to head.
|
|
|
|
`tenant` must expose `.id` and `.db_uri` (a control ORM Tenant, a TenantRef,
|
|
or any object with those attributes). Returns the applied head revision.
|
|
Records tenants.alembic_head and a ProvisioningJob('migrate') unless
|
|
record_job is False.
|
|
"""
|
|
db_uri = tenant.db_uri
|
|
cfg = _make_config(db_uri, script_location, version_locations)
|
|
|
|
job_id = None
|
|
if record_job:
|
|
with control_session() as s:
|
|
job = ProvisioningJob(tenant_id=tenant.id, action='migrate',
|
|
status='running', created_at=now_eastern())
|
|
s.add(job)
|
|
s.flush()
|
|
job_id = job.id
|
|
|
|
try:
|
|
# Alembic logs progress to stdout; keep the runner output tidy.
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
command.upgrade(cfg, 'head')
|
|
applied = current_revision(db_uri)
|
|
|
|
with control_session() as s:
|
|
t = s.get(Tenant, tenant.id)
|
|
if t is not None:
|
|
t.alembic_head = applied
|
|
if job_id is not None:
|
|
j = s.get(ProvisioningJob, job_id)
|
|
if j is not None:
|
|
j.status = 'ok'
|
|
j.finished_at = now_eastern()
|
|
j.log = f'upgraded to {applied}'
|
|
return applied
|
|
|
|
except Exception as e:
|
|
if job_id is not None:
|
|
with control_session() as s:
|
|
j = s.get(ProvisioningJob, job_id)
|
|
if j is not None:
|
|
j.status = 'failed'
|
|
j.finished_at = now_eastern()
|
|
j.log = f'{type(e).__name__}: {e}'
|
|
raise
|
|
|
|
|
|
def bootstrap_tenant(tenant, script_location=None, version_locations=None,
|
|
baseline_rev=BASELINE_REVISION, record_job=True):
|
|
"""Build a FRESH tenant database, then mark it current.
|
|
|
|
Runs the squashed baseline (full schema) only, then `stamp head` — the
|
|
historical phase migrations are NOT replayed (several are not idempotent and
|
|
would conflict with the full-schema baseline). Use this for brand-new tenant
|
|
databases; use upgrade_tenant() for ongoing incremental migrations.
|
|
|
|
Returns the stamped head revision.
|
|
"""
|
|
db_uri = tenant.db_uri
|
|
cfg = _make_config(db_uri, script_location, version_locations)
|
|
|
|
job_id = None
|
|
if record_job:
|
|
with control_session() as s:
|
|
job = ProvisioningJob(tenant_id=tenant.id, action='migrate',
|
|
status='running', created_at=now_eastern())
|
|
s.add(job)
|
|
s.flush()
|
|
job_id = job.id
|
|
|
|
try:
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
command.upgrade(cfg, baseline_rev) # build full schema (baseline only)
|
|
command.stamp(cfg, 'head') # mark at head without replaying phases
|
|
applied = current_revision(db_uri)
|
|
|
|
with control_session() as s:
|
|
t = s.get(Tenant, tenant.id)
|
|
if t is not None:
|
|
t.alembic_head = applied
|
|
if job_id is not None:
|
|
j = s.get(ProvisioningJob, job_id)
|
|
if j is not None:
|
|
j.status = 'ok'
|
|
j.finished_at = now_eastern()
|
|
j.log = f'bootstrapped (baseline {baseline_rev}) + stamped {applied}'
|
|
return applied
|
|
|
|
except Exception as e:
|
|
if job_id is not None:
|
|
with control_session() as s:
|
|
j = s.get(ProvisioningJob, job_id)
|
|
if j is not None:
|
|
j.status = 'failed'
|
|
j.finished_at = now_eastern()
|
|
j.log = f'{type(e).__name__}: {e}'
|
|
raise
|
|
|
|
def _select_tenants(selector):
|
|
with control_session() as s:
|
|
q = s.query(Tenant).filter(Tenant.status.in_(MIGRATABLE_STATUSES))
|
|
if selector != 'all':
|
|
if selector.isdigit():
|
|
q = q.filter(Tenant.id == int(selector))
|
|
else:
|
|
q = q.filter(Tenant.slug == selector)
|
|
return [TenantRef(t.id, t.slug, t.db_uri) for t in q.order_by(Tenant.id).all()]
|
|
|
|
|
|
def _cmd_upgrade(args):
|
|
tenants = _select_tenants(args.tenant)
|
|
if not tenants:
|
|
print(f'No matching tenants for --tenant {args.tenant}.')
|
|
return 0
|
|
head = chain_head()
|
|
print(f'Target head: {head}')
|
|
failures = 0
|
|
for t in tenants:
|
|
try:
|
|
applied = upgrade_tenant(t)
|
|
print(f' [ok] {t.slug} (id={t.id}) -> {applied}')
|
|
except Exception as e:
|
|
failures += 1
|
|
print(f' [FAIL] {t.slug} (id={t.id}): {type(e).__name__}: {e}')
|
|
print(f'Done. {len(tenants) - failures}/{len(tenants)} upgraded.')
|
|
return 1 if failures else 0
|
|
|
|
|
|
def _cmd_current(args):
|
|
tenants = _select_tenants(args.tenant)
|
|
if not tenants:
|
|
print(f'No matching tenants for --tenant {args.tenant}.')
|
|
return 0
|
|
head = chain_head()
|
|
print(f'Chain head: {head}')
|
|
for t in tenants:
|
|
try:
|
|
rev = current_revision(t.db_uri)
|
|
except Exception as e:
|
|
rev = f'<error: {type(e).__name__}>'
|
|
flag = '' if rev == head else ' (BEHIND)'
|
|
print(f' {t.slug} (id={t.id}): {rev}{flag}')
|
|
return 0
|
|
|
|
|
|
def _cmd_bootstrap(args):
|
|
tenants = _select_tenants(args.tenant)
|
|
if not tenants:
|
|
print(f'No matching tenants for --tenant {args.tenant}.')
|
|
return 0
|
|
failures = 0
|
|
for t in tenants:
|
|
try:
|
|
applied = bootstrap_tenant(t)
|
|
print(f' [ok] {t.slug} (id={t.id}) bootstrapped -> {applied}')
|
|
except Exception as e:
|
|
failures += 1
|
|
print(f' [FAIL] {t.slug} (id={t.id}): {type(e).__name__}: {e}')
|
|
return 1 if failures else 0
|
|
|
|
|
|
def _cmd_heads(_args):
|
|
print(chain_head())
|
|
return 0
|
|
|
|
|
|
def main(argv=None):
|
|
parser = argparse.ArgumentParser(prog='control.tenant_migrate',
|
|
description='Per-tenant migration runner')
|
|
sub = parser.add_subparsers(dest='command', required=True)
|
|
|
|
p_up = sub.add_parser('upgrade', help='Upgrade tenant DB(s) to head')
|
|
p_up.add_argument('--tenant', required=True, help="'all', or a tenant id or slug")
|
|
p_up.set_defaults(func=_cmd_upgrade)
|
|
|
|
p_cur = sub.add_parser('current', help='Show each tenant DB current revision')
|
|
p_cur.add_argument('--tenant', default='all', help="'all', or a tenant id or slug")
|
|
p_cur.set_defaults(func=_cmd_current)
|
|
|
|
p_boot = sub.add_parser('bootstrap',
|
|
help='Build a FRESH tenant DB (baseline + stamp head)')
|
|
p_boot.add_argument('--tenant', required=True, help="'all', or a tenant id or slug")
|
|
p_boot.set_defaults(func=_cmd_bootstrap)
|
|
|
|
sub.add_parser('heads', help='Show the chain head revision').set_defaults(func=_cmd_heads)
|
|
|
|
args = parser.parse_args(argv)
|
|
sys.exit(args.func(args))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|