350 lines
14 KiB
Python
350 lines
14 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'
|
|
|
|
#: The last revision the squashed baseline actually covers.
|
|
#:
|
|
#: Everything AFTER this point is existence-guarded (INFORMATION_SCHEMA checks,
|
|
#: or an idempotent ENUM MODIFY), so it can be replayed over the baseline
|
|
#: safely. Everything BEFORE it is not — several pre-phase33 migrations would
|
|
#: collide with the full-schema baseline, which is why bootstrap stamps past
|
|
#: them instead of running them.
|
|
#:
|
|
#: If the baseline is ever re-squashed to a later point, move this with it.
|
|
BASELINE_COVERS_THROUGH = 'phase32_device_token_columns'
|
|
|
|
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 bring it to head.
|
|
|
|
Three steps, and the middle one is the load-bearing part:
|
|
|
|
1. `upgrade(baseline)` — the squashed baseline builds the bulk of the
|
|
schema in one go.
|
|
2. `stamp(BASELINE_COVERS_THROUGH)` — declare the DB to be at the last
|
|
revision the baseline actually covers, WITHOUT running the pre-phase33
|
|
migrations. Those are not idempotent and would collide with the
|
|
baseline; skipping them is the whole reason bootstrap exists.
|
|
3. `upgrade(head)` — replay the guarded tail (phase33 onward). Every one
|
|
of those checks INFORMATION_SCHEMA before touching anything (the two
|
|
ENUM widenings use an idempotent MODIFY), so this adds exactly what
|
|
the baseline lacks and skips the rest.
|
|
|
|
This used to stop after step 1 and `stamp('head')` instead — claiming the
|
|
database was current when it was missing every column and table added since
|
|
phase32. Because the baseline was last refreshed around phase33, a tenant
|
|
provisioned that way lacked `users.mfa_enabled` and `users.ui_theme`, and
|
|
SQLAlchemy emits every mapped column in its SELECT — so the new workspace
|
|
could not even log in. It failed at the first query, not at some optional
|
|
feature, and only for freshly provisioned tenants (tenant-zero was adopted
|
|
in place with a real schema), which is what kept it hidden.
|
|
|
|
Use upgrade_tenant() for ongoing incremental migrations of a tenant that
|
|
already exists.
|
|
|
|
Returns the head revision the database ends up at.
|
|
"""
|
|
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) # 1. baseline schema
|
|
command.stamp(cfg, BASELINE_COVERS_THROUGH) # 2. skip the unguarded past
|
|
command.upgrade(cfg, 'head') # 3. replay the guarded tail
|
|
applied = current_revision(db_uri)
|
|
|
|
# A bootstrap that does not end at head has silently produced a broken
|
|
# tenant — exactly the failure this sequence exists to prevent. Say so
|
|
# here rather than letting it surface as a missing column later.
|
|
expected = chain_head(script_location, version_locations)
|
|
if applied != expected:
|
|
raise RuntimeError(
|
|
f'bootstrap ended at {applied!r}, expected head {expected!r} — '
|
|
f'the tenant database is incomplete')
|
|
|
|
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 '
|
|
f'{BASELINE_COVERS_THROUGH}, 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 _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 _cmd_stamp(args):
|
|
"""Stamp tenant DB(s) to a specific revision without running migrations.
|
|
|
|
Use this to rewind a mistakenly-stamped tenant so a subsequent `upgrade`
|
|
will actually replay the missing migration. Example:
|
|
|
|
python -m control.tenant_migrate stamp --tenant gov phase32_device_token_columns
|
|
python -m control.tenant_migrate upgrade --tenant gov
|
|
"""
|
|
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:
|
|
cfg = _make_config(t.db_uri)
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
command.stamp(cfg, args.revision)
|
|
rev = current_revision(t.db_uri)
|
|
print(f' [ok] {t.slug} (id={t.id}) stamped -> {rev}')
|
|
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 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)
|
|
|
|
p_stamp = sub.add_parser('stamp', help='Stamp tenant DB(s) to a specific revision (no migrations run)')
|
|
p_stamp.add_argument('--tenant', required=True, help="'all', or a tenant id or slug")
|
|
p_stamp.add_argument('revision', help='Alembic revision id (e.g. phase32_device_token_columns)')
|
|
p_stamp.set_defaults(func=_cmd_stamp)
|
|
|
|
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()
|