Schema Versioning and Rollback¶

version = '20260109_120000'  # Timestamp-based identifier

from surql.migration.versioning import SchemaSnapshot

snapshot = SchemaSnapshot(
  version='20260109_120000',
  created_at=datetime.now(UTC),
  tables={'user': {...}, 'post': {...}},
  edges={'likes': {...}},
  indexes={'user': [...]},
  checksum='sha256_hash',
  migration_count=5,
)

from surql.migration.versioning import VersionGraph

graph = VersionGraph()
graph.add_version(migration1)
graph.add_version(migration2, parent='20260108_120000')

from surql.migration.rollback import create_rollback_plan

plan = await create_rollback_plan(
  client,
  migrations,
  target_version='20260108_120000',
)

from surql.migration.versioning import create_snapshot
from surql.connection.client import get_client

async def capture_schema():
  async with get_client(config) as client:
    snapshot = await create_snapshot(
      client=client,
      version='20260109_120000',
      migration_count=5,
    )

    print(f'Snapshot created:')
    print(f'  Version: {snapshot.version}')
    print(f'  Tables: {len(snapshot.tables)}')
    print(f'  Edges: {len(snapshot.edges)}')
    print(f'  Checksum: {snapshot.checksum[:16]}...')

from surql.migration.versioning import store_snapshot

async def save_snapshot(snapshot):
  async with get_client(config) as client:
    await store_snapshot(client, snapshot)
    print('Snapshot stored successfully')

from surql.migration.versioning import load_snapshot

async def retrieve_snapshot(version):
  async with get_client(config) as client:
    snapshot = await load_snapshot(client, version)

    if snapshot:
      print(f'Loaded snapshot for version {version}')
      print(f'  Created at: {snapshot.created_at}')
      print(f'  Migration count: {snapshot.migration_count}')
    else:
      print(f'Snapshot not found for version {version}')

    return snapshot

from surql.migration.versioning import list_snapshots

async def show_all_snapshots():
  async with get_client(config) as client:
    snapshots = await list_snapshots(client)

    print(f'Found {len(snapshots)} snapshot(s):')
    for snapshot in snapshots:
      print(f'  {snapshot.version} - {snapshot.created_at}')
      print(f'    Tables: {len(snapshot.tables)}, Migrations: {snapshot.migration_count}')

from surql.migration.versioning import compare_snapshots

async def compare_versions(version1, version2):
  async with get_client(config) as client:
    snapshot1 = await load_snapshot(client, version1)
    snapshot2 = await load_snapshot(client, version2)

    if not snapshot1 or not snapshot2:
      print('One or both snapshots not found')
      return

    diff = compare_snapshots(snapshot1, snapshot2)

    print(f'Comparing {version1} → {version2}')

    if diff['checksum_match']:
      print('Schemas are identical')
      return

    if diff['tables_added']:
      print(f"Tables added: {', '.join(diff['tables_added'])}")

    if diff['tables_removed']:
      print(f"Tables removed: {', '.join(diff['tables_removed'])}")

    if diff['tables_modified']:
      print(f"Tables modified: {', '.join(diff['tables_modified'])}")

    if diff['edges_added']:
      print(f"Edges added: {', '.join(diff['edges_added'])}")

    if diff['edges_removed']:
      print(f"Edges removed: {', '.join(diff['edges_removed'])}")

from surql.migration.versioning import VersionGraph
from surql.migration.discovery import discover_migrations
from pathlib import Path

def build_version_graph():
  migrations = discover_migrations(Path('migrations'))

  graph = VersionGraph()

  for i, migration in enumerate(migrations):
    parent = migrations[i - 1].version if i > 0 else None
    graph.add_version(migration, parent=parent)

  print(f'Graph contains {len(graph.get_all_versions())} versions')
  return graph

def query_graph(graph):
  versions = graph.get_all_versions()

  if len(versions) >= 2:
    first = versions[0]
    last = versions[-1]

    # Get path between versions
    path = graph.get_path(first, last)
    if path:
      print(f'Path from {first} to {last}:')
      for v in path:
        print(f'  → {v}')

    # Get ancestors
    ancestors = graph.get_ancestors(last)
    print(f'\nAncestors of {last}:')
    for v in ancestors:
      print(f'  ← {v}')

    # Get descendants
    descendants = graph.get_descendants(first)
    print(f'\nDescendants of {first}:')
    for v in descendants:
      print(f'  → {v}')

def plan_using_graph(graph, current_version, target_version):
  # Get path from current to target
  path = graph.get_path(current_version, target_version)

  if not path:
    print(f'No path found from {current_version} to {target_version}')
    return

  # Migrations to rollback are those in the path (excluding target)
  to_rollback = [v for v in path if v != target_version]
  to_rollback.reverse()  # Newest first

  print(f'Will rollback {len(to_rollback)} migration(s):')
  for v in to_rollback:
    print(f'  - {v}')

from surql.migration.rollback import create_rollback_plan
from surql.migration.discovery import discover_migrations

async def plan_rollback():
  migrations = discover_migrations(Path('migrations'))

  async with get_client(config) as client:
    plan = await create_rollback_plan(
      client=client,
      migrations=migrations,
      target_version='20260108_120000',
    )

    print('Rollback Plan:')
    print(f'  From: {plan.from_version}')
    print(f'  To: {plan.to_version}')
    print(f'  Migrations to rollback: {plan.migration_count}')
    print(f'  Overall safety: {plan.overall_safety.value.upper()}')

    if plan.issues:
      print(f'\nSafety Issues ({len(plan.issues)}):')
      for issue in plan.issues:
        print(f'  [{issue.safety.value.upper()}] {issue.migration}')
        print(f'    {issue.description}')
        if issue.recommendation:
          print(f'    → {issue.recommendation}')

    if plan.requires_approval:
      print('\n⚠️  This rollback requires manual approval')

    return plan

from surql.migration.rollback import analyze_rollback_safety

async def check_safety(target_version):
  migrations = discover_migrations(Path('migrations'))

  async with get_client(config) as client:
    issues = await analyze_rollback_safety(
      client,
      migrations,
      target_version,
    )

    if not issues:
      print('✓ Rollback is safe - no issues detected')
    else:
      print(f'Found {len(issues)} potential issue(s):')
      for issue in issues:
        print(f'  {issue.safety.value}: {issue.description}')

# Migration added an index
def up():
  return ['DEFINE INDEX email_idx ON TABLE user COLUMNS email;']

def down():
  return ['REMOVE INDEX email_idx ON TABLE user;']

# Rollback: RollbackSafety.SAFE

# Migration added a field
def up():
  return ['DEFINE FIELD bio ON TABLE user TYPE string;']

def down():
  return ['REMOVE FIELD bio ON TABLE user;']

# Rollback: RollbackSafety.DATA_LOSS
# Reason: bio field data will be lost

# Migration created a table
def up():
  return ['DEFINE TABLE user SCHEMAFULL;', ...]

def down():
  return ['REMOVE TABLE user;']

# Rollback: RollbackSafety.UNSAFE
# Reason: All user table data will be lost

from surql.migration.rollback import execute_rollback

async def rollback_safely():
  migrations = discover_migrations(Path('migrations'))

  async with get_client(config) as client:
    # Create plan
    plan = await create_rollback_plan(
      client,
      migrations,
      target_version='20260108_120000',
    )

    # Check safety
    if plan.overall_safety == RollbackSafety.SAFE:
      print('Executing safe rollback...')

      result = await execute_rollback(client, plan)

      if result.success:
        print(f'✓ Successfully rolled back {result.rolled_back_count} migration(s)')
        print(f'  Duration: {result.actual_duration_ms}ms')
      else:
        print(f'✗ Rollback failed after {result.rolled_back_count} migration(s)')
        for error in result.errors:
          print(f'  Error: {error}')
    else:
      print(f'Rollback is {plan.overall_safety.value} - review issues first')

async def force_rollback():
  migrations = discover_migrations(Path('migrations'))

  async with get_client(config) as client:
    plan = await create_rollback_plan(
      client,
      migrations,
      target_version='20260107_120000',
    )

    # Display warnings
    print(f'Safety: {plan.overall_safety.value.upper()}')

    if plan.has_data_loss:
      print('\n⚠️  WARNING: This rollback may cause data loss!')
      print('Issues:')
      for issue in plan.issues:
        print(f'  - {issue.description}')
        if issue.affected_data:
          print(f'    Affected: {issue.affected_data}')

      # Simulate user confirmation
      confirm = input('\nProceed with rollback? [yes/NO]: ')

      if confirm.lower() == 'yes':
        print('\nProceeding with forced rollback...')
        result = await execute_rollback(client, plan, force=True)

        if result.success:
          print('✓ Rollback completed')
        else:
          print(f'✗ Rollback failed: {result.errors}')
      else:
        print('Rollback cancelled')

async def preview_rollback():
  migrations = discover_migrations(Path('migrations'))

  async with get_client(config) as client:
    plan = await create_rollback_plan(
      client,
      migrations,
      target_version='20260108_120000',
    )

    print('=== Rollback Preview ===')
    print(f'From: {plan.from_version}')
    print(f'To: {plan.to_version}')
    print(f'Safety: {plan.overall_safety.value}')
    print()

    print(f'Migrations to rollback ({plan.migration_count}):')
    for migration in plan.migrations:
      print(f'  {migration.version}: {migration.description}')

      # Show down() SQL
      try:
        statements = migration.down()
        for stmt in statements:
          print(f'    {stmt[:60]}...' if len(stmt) > 60 else f'    {stmt}')
      except Exception as e:
        print(f'    Error getting down migration: {e}')

    print('\nDry run - no changes made')

from surql.migration.history import enable_auto_snapshots, disable_auto_snapshots

# Enable automatic snapshots
enable_auto_snapshots()
print('Automatic snapshots enabled')

# Future migrations will create snapshots automatically
# await apply_migrations(...)

# Disable when done
disable_auto_snapshots()
print('Automatic snapshots disabled')

from surql.migration.history import enable_auto_snapshots
from surql.migration.executor import apply_migrations

async def migrate_with_snapshots():
  # Enable auto-snapshots
  enable_auto_snapshots()

  async with get_client(config) as client:
    # Apply migrations
    await apply_migrations(client, migrations)

    # Snapshots were created automatically for each migration
    snapshots = await list_snapshots(client)
    print(f'{len(snapshots)} snapshots created')

async def major_migration():
  async with get_client(config) as client:
    # Snapshot before
    snapshot_before = await create_snapshot(client, 'before_major_change', 10)
    await store_snapshot(client, snapshot_before)

    # Apply migrations
    await apply_migrations(client, migrations)

    # Snapshot after
    snapshot_after = await create_snapshot(client, 'after_major_change', 15)
    await store_snapshot(client, snapshot_after)

async def safe_unsafe_rollback():
  async with get_client(config) as client:
    # Create rollback plan
    plan = await create_rollback_plan(client, migrations, target_version)

    if plan.overall_safety == RollbackSafety.UNSAFE:
      # Identify affected tables
      affected_tables = set()
      for issue in plan.issues:
        if issue.affected_data:
          # Extract table name from affected_data
          if 'table' in issue.affected_data.lower():
            table_name = issue.affected_data.split()[0]
            affected_tables.add(table_name)

      # Backup affected tables
      for table in affected_tables:
        print(f'Backing up {table}...')
        records = await client.select(table)
        with open(f'backup_{table}.json', 'w') as f:
          json.dump(records, f)

      # Now safe to rollback
      result = await execute_rollback(client, plan, force=True)

async def production_migration():
  # Enable auto-snapshots for safety
  enable_auto_snapshots()

  try:
    async with get_client(config) as client:
      await apply_migrations(client, migrations)
  finally:
    disable_auto_snapshots()

"""
Migration: Add user profile fields

ROLLBACK WARNING: Rolling back this migration will lose all user
profile data stored in the bio and avatar_url fields.

Before rollback:
1. Export user profiles: SELECT id, bio, avatar_url FROM user
2. Save to backup file
3. Verify backup integrity
4. Proceed with rollback
"""

def up():
  return [
    'DEFINE FIELD bio ON TABLE user TYPE string;',
    'DEFINE FIELD avatar_url ON TABLE user TYPE string;',
  ]

def down():
  return [
    'REMOVE FIELD avatar_url ON TABLE user;',
    'REMOVE FIELD bio ON TABLE user;',
  ]

import time

async def monitored_rollback(plan):
  async with get_client(config) as client:
    start = time.time()
    result = await execute_rollback(client, plan)
    duration = time.time() - start

    # Log results
    logger.info(
      'rollback_completed',
      success=result.success,
      rolled_back=result.rolled_back_count,
      planned=plan.migration_count,
      duration_seconds=duration,
      errors=result.errors,
    )

    return result

async def cleanup_old_snapshots(keep_count=10):
  """Keep only the most recent snapshots."""
  async with get_client(config) as client:
    snapshots = await list_snapshots(client)

    if len(snapshots) > keep_count:
      to_delete = snapshots[:-keep_count]  # All except last N

      for snapshot in to_delete:
        await client.execute(
          'DELETE _schema_snapshot WHERE version = $version',
          {'version': snapshot.version},
        )
        print(f'Deleted snapshot: {snapshot.version}')

async def debug_rollback_plan():
  from surql.migration.history import get_applied_migrations

  async with get_client(config) as client:
    applied = await get_applied_migrations(client)

    if applied:
      current = applied[-1].version
      print(f'Current version: {current}')
      print(f'Applied migrations: {[m.version for m in applied]}')
    else:
      print('No migrations applied')

async def recover_from_failed_rollback():
  async with get_client(config) as client:
    # Check current state
    history = await get_applied_migrations(client)

    if history:
      current = history[-1].version
      print(f'Rollback stopped at: {current}')

      # Option 1: Retry rollback from current position
      plan = await create_rollback_plan(client, migrations, target_version)
      result = await execute_rollback(client, plan)

      # Option 2: Roll forward to consistent state
      # await apply_migrations(client, pending_migrations)

async def verify_schema_integrity():
  async with get_client(config) as client:
    # Create new snapshot
    current_snapshot = await create_snapshot(client, 'current', 0)

    # Load expected snapshot
    expected_snapshot = await load_snapshot(client, expected_version)

    if current_snapshot.checksum == expected_snapshot.checksum:
      print('✓ Schema matches expected state')
    else:
      print('✗ Schema diverged from expected state')

      # Compare to find differences
      diff = compare_snapshots(expected_snapshot, current_snapshot)
      print('Differences:', diff)

from pathlib import Path
from surql.migration.discovery import discover_migrations
from surql.migration.rollback import create_rollback_plan, execute_rollback
from surql.migration.versioning import create_snapshot, store_snapshot

async def complete_rollback_workflow():
  """Complete workflow: snapshot, plan, review, execute."""

  async with get_client(config) as client:
    # Step 1: Create pre-rollback snapshot
    print('Step 1: Creating snapshot before rollback')
    snapshot = await create_snapshot(client, 'pre_rollback', 5)
    await store_snapshot(client, snapshot)
    print(f'  Snapshot created: {snapshot.checksum[:16]}...')

    # Step 2: Create rollback plan
    print('\nStep 2: Creating rollback plan')
    migrations = discover_migrations(Path('migrations'))
    plan = await create_rollback_plan(client, migrations, '20260108_120000')

    print(f'  From: {plan.from_version}')
    print(f'  To: {plan.to_version}')
    print(f'  Migrations: {plan.migration_count}')
    print(f'  Safety: {plan.overall_safety.value}')

    # Step 3: Review issues
    print('\nStep 3: Reviewing safety issues')
    if plan.issues:
      for issue in plan.issues:
        print(f'  [{issue.safety.value}] {issue.description}')
    else:
      print('  No issues found')

    # Step 4: Execute if safe
    print('\nStep 4: Executing rollback')
    if plan.is_safe:
      result = await execute_rollback(client, plan)

      if result.success:
        print(f'  ✓ Rolled back {result.rolled_back_count} migrations')
        print(f'  Duration: {result.actual_duration_ms}ms')
      else:
        print(f'  ✗ Rollback failed: {result.errors}')
    else:
      print(f'  ⚠️  Rollback is {plan.overall_safety.value} - requires review')

async def test_migration_reversibility():
  """Test that migrations can be safely rolled back."""

  migrations = discover_migrations(Path('migrations'))

  async with get_client(config) as client:
    for migration in migrations:
      print(f'Testing {migration.version}...')

      # Apply migration
      await execute_migration(client, migration, MigrationDirection.UP)

      # Create snapshot
      snapshot_after = await create_snapshot(client, f'after_{migration.version}', 1)

      # Rollback
      await execute_migration(client, migration, MigrationDirection.DOWN)

      # Verify state
      snapshot_rolled_back = await create_snapshot(client, f'rolled_back_{migration.version}', 0)

      # Compare (should match initial state)
      # Note: This is simplified - real test would compare with before snapshot
      print(f'  ✓ {migration.version} is reversible')

from datetime import datetime, timedelta

async def manage_snapshots():
  """Manage snapshot lifecycle - create, clean old, export."""

  async with get_client(config) as client:
    # Create current snapshot
    current = await create_snapshot(client, datetime.now().strftime('%Y%m%d_%H%M%S'), 0)
    await store_snapshot(client, current)
    print(f'Created snapshot: {current.version}')

    # List all snapshots
    all_snapshots = await list_snapshots(client)
    print(f'Total snapshots: {len(all_snapshots)}')

    # Clean old snapshots (keep last 30 days)
    cutoff = datetime.now() - timedelta(days=30)
    old_snapshots = [s for s in all_snapshots if s.created_at < cutoff]

    for snapshot in old_snapshots:
      await client.execute(
        'DELETE _schema_snapshot WHERE version = $version',
        {'version': snapshot.version},
      )
      print(f'Deleted old snapshot: {snapshot.version}')

    print(f'Cleaned {len(old_snapshots)} old snapshot(s)')