pr refactor and other changes

Author: jsong468

build_scripts/migrate_prod_memory_schema.py

@@ -9,33 +9,26 @@
 doc/contributing/10_release_process.md) or by a CD pipeline — never by normal
 application startup.
 
+It constructs an AzureSQLMemory with skip_schema_migration=True (bypassing the
+runtime guard), then explicitly calls _run_schema_migration to upgrade to head.
+The environment checks ensure this only runs from a release branch.
+
 Safety rails:
-- Requires an explicit --target-revision argument (no blind "upgrade head").
 - Validates the environment (release branch, clean working tree, no .dev version).
-- Validates the target revision exists in the local migration graph.
-- Reports current DB revision before and after migration.
-- Runs schema check after upgrade to confirm models match.
 - Interactive confirmation when running in a terminal.
 - Exits non-zero on any failure.
 
 Usage:
-    python build_scripts/migrate_prod_memory_schema.py \
-        --target-revision c3d5e7f9a1b2
+    python build_scripts/migrate_prod_memory_schema.py
 
-The script reads the production connection string from the
-AZURE_SQL_DB_CONNECTION_STRING environment variable (same as AzureSQLMemory).
+The script reads the production connection string from
+AZURE_SQL_DB_CONNECTION_STRING_PROD (loaded from ~/.pyrit/.env).
 """
 
-import argparse
-import os
 import subprocess
 import sys
 
 import dotenv
-from alembic import command
-from alembic.script import ScriptDirectory
-from alembic.util.exc import AutogenerateDiffsDetected
-from sqlalchemy import create_engine, text
 
 from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH
 
@@ -51,77 +44,17 @@
 _YELLOW = "\033[93m"
 _RESET = "\033[0m"
 
-_CONNECTION_STRING_ENV = "AZURE_SQL_DB_CONNECTION_STRING_PROD"
-
 
 def _print_error(message: str) -> None:
     """Print an error message in red to stderr."""
     print(f"{_RED}ERROR: {message}{_RESET}", file=sys.stderr)
 
 
-def _print_warning(message: str) -> None:
-    """Print a warning message in yellow."""
-    print(f"{_YELLOW}WARNING: {message}{_RESET}")
-
-
 def _print_success(message: str) -> None:
     """Print a success message in green."""
     print(f"{_GREEN}{message}{_RESET}")
 
 
-def _get_current_revision(*, engine) -> str | None:
-    """
-    Read the current Alembic revision from the database.
-
-    Returns None if no version table exists (fresh database).
-    """
-    from sqlalchemy import inspect as sa_inspect
-
-    from pyrit.memory.migration import PYRIT_MEMORY_ALEMBIC_VERSION_TABLE
-
-    inspector = sa_inspect(engine)
-    if PYRIT_MEMORY_ALEMBIC_VERSION_TABLE not in inspector.get_table_names():
-        return None
-
-    with engine.connect() as conn:
-        result = conn.execute(text(f"SELECT version_num FROM {PYRIT_MEMORY_ALEMBIC_VERSION_TABLE}"))
-        row = result.fetchone()
-        return row[0] if row else None
-
-
-def _validate_revision_exists(*, target_revision: str) -> bool:
-    """Check that the target revision exists in the local migration script directory."""
-
-    from pathlib import Path
-
-    script_location = Path(__file__).parent.parent / "pyrit" / "memory" / "alembic"
-    # Build a minimal config to get the ScriptDirectory
-    from alembic.config import Config
-
-    config = Config()
-    config.set_main_option("script_location", str(script_location))
-    script_dir = ScriptDirectory.from_config(config)
-
-    try:
-        script_dir.get_revision(target_revision)
-        return True
-    except Exception:
-        return False
-
-
-def _get_all_revision_ids() -> list[str]:
-    """Get all revision IDs from the local migration graph."""
-    from pathlib import Path
-
-    from alembic.config import Config
-
-    script_location = Path(__file__).parent.parent / "pyrit" / "memory" / "alembic"
-    config = Config()
-    config.set_main_option("script_location", str(script_location))
-    script_dir = ScriptDirectory.from_config(config)
-    return [rev.revision for rev in script_dir.walk_revisions()]
-
-
 def _check_release_environment() -> list[str]:
     """
     Validate that the script is running in a proper release environment.
@@ -130,7 +63,6 @@ def _check_release_environment() -> list[str]:
     """
     issues: list[str] = []
 
-    # Check 1: Running from a release branch (releases/vX.Y.Z)
     try:
         branch = subprocess.check_output(
             ["git", "rev-parse", "--abbrev-ref", "HEAD"],
@@ -145,7 +77,6 @@ def _check_release_environment() -> list[str]:
     except (subprocess.CalledProcessError, FileNotFoundError):
         issues.append("Could not determine current Git branch.")
 
-    # Check 2: Clean working tree (no uncommitted changes to memory files)
     try:
         dirty_files = subprocess.check_output(
             ["git", "status", "--porcelain", "--", "pyrit/memory/"],
@@ -161,7 +92,6 @@ def _check_release_environment() -> list[str]:
     except (subprocess.CalledProcessError, FileNotFoundError):
         issues.append("Could not check Git working tree status.")
 
-    # Check 3: Not a .dev version (should be a release version)
     try:
         from pyrit import __version__
 
@@ -176,120 +106,21 @@ def _check_release_environment() -> list[str]:
     return issues
 
 
-def _run_migration(*, connection_string: str, target_revision: str) -> int:
-    """
-    Execute the migration against the target database.
-
-    Returns 0 on success, 1 on failure.
-    """
-    from pathlib import Path
-
-    from alembic.config import Config
-
-    from pyrit.memory.migration import check_schema_migrations
-
-    engine = create_engine(connection_string)
-
-    # Step 1: Report current state
-    current_rev = _get_current_revision(engine=engine)
-    print(f"Current database revision: {current_rev or '(none — fresh database)'}")
-    print(f"Target revision:           {target_revision}")
-    print()
-
-    if current_rev == target_revision:
-        _print_success("Database is already at the target revision. Nothing to do.")
-        engine.dispose()
-        return 0
-
-    # Step 2: Run upgrade to specific revision (not head)
-    print("Applying migrations...")
-    try:
-        script_location = Path(__file__).parent.parent / "pyrit" / "memory" / "alembic"
-        with engine.begin() as connection:
-            config = Config()
-            config.set_main_option("script_location", str(script_location))
-            config.attributes["connection"] = connection
-
-            # Stamp unversioned schemas if needed
-            from pyrit.memory.migration import _validate_and_stamp_unversioned_memory_schema
-
-            _validate_and_stamp_unversioned_memory_schema(config=config, connection=connection)
-            command.upgrade(config, target_revision)
-    except Exception as e:
-        _print_error(f"Migration failed: {e}")
-        engine.dispose()
-        return 1
-
-    # Step 3: Verify new revision
-    new_rev = _get_current_revision(engine=engine)
-    print(f"Database revision after migration: {new_rev}")
-
-    if new_rev != target_revision:
-        _print_error(f"Expected revision {target_revision}, but database is at {new_rev}.")
-        engine.dispose()
-        return 1
-
-    # Step 4: Schema check — verify models match
-    print("Verifying schema matches models...")
-    try:
-        check_schema_migrations(engine=engine)
-    except AutogenerateDiffsDetected as e:
-        _print_error(f"Schema check failed after migration: {e}")
-        _print_error("The models in this codebase do not match the database schema.")
-        engine.dispose()
-        return 1
-
-    _print_success("Migration completed and verified successfully.")
-    engine.dispose()
-    return 0
-
+def main() -> int:
+    """Entry point for production schema migration."""
+    import argparse
 
-def _build_parser() -> argparse.ArgumentParser:
-    """Build the CLI argument parser."""
     parser = argparse.ArgumentParser(
-        description=(
-            "Apply Alembic schema migrations to a production database. "
-            "Validates release environment and requires an explicit target revision."
-        ),
-    )
-    parser.add_argument(
-        "--target-revision",
-        required=True,
-        help="The exact Alembic revision ID to upgrade to (e.g., 'c3d5e7f9a1b2').",
-    )
-    parser.add_argument(
-        "--connection-string-env",
-        default=_CONNECTION_STRING_ENV,
-        help=f"Environment variable containing the connection string. Default: {_CONNECTION_STRING_ENV}",
+        description="Apply Alembic schema migrations to the production database.",
     )
     parser.add_argument(
         "--skip-environment-check",
         action="store_true",
         help="Skip release environment checks (branch, clean tree, version). Use only in CI with caution.",
     )
-    return parser
-
+    args = parser.parse_args()
 
-def main() -> int:
-    """Entry point for production schema migration."""
-    args = _build_parser().parse_args()
-
-    # Safety rail 1: Require connection string
-    connection_string = os.environ.get(args.connection_string_env)
-    if not connection_string:
-        _print_error(f"Environment variable '{args.connection_string_env}' is not set.")
-        return 1
-
-    # Safety rail 2: Validate target revision exists in local migration graph
-    if not _validate_revision_exists(target_revision=args.target_revision):
-        available = _get_all_revision_ids()
-        _print_error(
-            f"Target revision '{args.target_revision}' not found in the local migration graph.\n"
-            f"Available revisions: {', '.join(available)}"
-        )
-        return 1
-
-    # Safety rail 3: Verify release environment (branch, clean tree, version)
+    # Safety rail: Verify release environment
     if not args.skip_environment_check:
         issues = _check_release_environment()
         if issues:
@@ -299,18 +130,39 @@ def main() -> int:
             _print_error("Fix the above issues or pass --skip-environment-check (CI only).")
             return 1
     else:
-        _print_warning("Skipping release environment checks (--skip-environment-check).")
+        print(f"{_YELLOW}WARNING: Skipping release environment checks.{_RESET}")
 
-    # Confirmation prompt (unless running in CI where stdin is not interactive)
+    # Interactive confirmation
     if sys.stdin.isatty():
-        print(f"About to migrate production database to revision: {args.target_revision}")
-        print(f"Connection string source: ${args.connection_string_env}")
+        print("About to migrate production database schema to head.")
         response = input("Type 'yes' to proceed: ")
         if response.strip().lower() != "yes":
             print("Aborted.")
             return 1
 
-    return _run_migration(connection_string=connection_string, target_revision=args.target_revision)
+    # Construct AzureSQLMemory with skip_schema_migration=True to bypass the runtime guard,
+    # then explicitly run migration.
+    import os
+
+    from pyrit.memory import AzureSQLMemory
+
+    prod_conn = os.environ.get(AzureSQLMemory.AZURE_SQL_DB_CONNECTION_STRING_PROD)
+    if not prod_conn:
+        _print_error(f"Environment variable '{AzureSQLMemory.AZURE_SQL_DB_CONNECTION_STRING_PROD}' is not set.")
+        return 1
+
+    try:
+        memory = AzureSQLMemory(
+            connection_string=prod_conn,
+            skip_schema_migration=True,
+        )
+        print("Running schema migration...")
+        memory._run_schema_migration()
+        _print_success("Production schema migration completed and verified successfully.")
+        return 0
+    except Exception as e:
+        _print_error(f"Migration failed: {e}")
+        return 1
 
 
 if __name__ == "__main__":

doc/contributing/10_release_process.md

@@ -220,36 +220,27 @@ uv sync --frozen
 python -c "import pyrit; print(pyrit.__version__)"  # verify: x.y.z (no .dev0)
 ```
 
-**Identify the target revision** — the Alembic head on this branch. We use an explicit
-revision (not `head`) so the migration is deterministic and tied to this exact release.
-
-```bash
-python build_scripts/memory_migrations.py head
-```
-
-This prints the revision ID (e.g., `c3d5e7f9a1b2`) to use as `<revision_id>` below.
-
 **Run the migration** (reads `AZURE_SQL_DB_CONNECTION_STRING_PROD` from `~/.pyrit/.env`):
 
 ```bash
-python build_scripts/migrate_prod_memory_schema.py --target-revision <revision_id>
+python build_scripts/migrate_prod_memory_schema.py
 ```
 
 The script validates the environment (release branch, clean tree, no `.dev` version),
-confirms the target revision exists, applies migrations, and verifies the schema matches
-models. It exits non-zero on any failure, and migrations roll back automatically.
+constructs an `AzureSQLMemory` pointed at prod, and runs `_run_schema_migration()` which
+upgrades to head and verifies the schema matches models. Since you're on the release branch,
+head is the release revision.
 
-**Verify prod is usable after migration.** This connects to prod using the check-only
-path (no schema modification) and confirms compatibility:
+**Verify prod is usable after migration:**
 
 ```bash
 python -c "from pyrit.memory import AzureSQLMemory; AzureSQLMemory()"
 ```
 
 If it exits without error, prod is ready.
 
-If no schema changes landed in this release, the script reports "already at target revision"
-and exits cleanly. Still run it as confirmation.
+If no schema changes landed in this release, `_run_schema_migration` is a no-op.
+Still run it as confirmation.
 
 **Rollback policy:** forward-fix only. Ship a new corrective migration rather than downgrading,
 since `downgrade()` risks data loss.

pyrit/memory/azure_sql_memory.py

@@ -117,15 +117,28 @@ def __init__(
         prod_connection_string = default_values.get_non_required_value(
             env_var_name=self.AZURE_SQL_DB_CONNECTION_STRING_PROD
         )
-        if not skip_schema_migration:
-            if self._connection_string == prod_connection_string:
-                # Production guard: verify schema compatibility without modifying the database.
-                # Logs a warning on mismatch but does not block startup, so developers on
-                # newer code can still query prod data.
+
+        is_prod = bool(prod_connection_string) and self._connection_string == prod_connection_string
+        should_migrate = not is_prod and not skip_schema_migration
+
+        if should_migrate:
+            # Non-production: run schema migration (upgrade + check).
+            self._run_schema_migration(silent=silent)
+        else:
+            # Production or skip_schema_migration=True: verify schema compatibility
+            # without modifying the database. Logs a warning on mismatch but does not
+            # block startup, so developers on newer code can still query data.
+            from alembic.util.exc import AutogenerateDiffsDetected, CommandError
+
+            try:
                 self._check_schema_migration(silent=silent)
-            else:
-                # For non-production databases, run normal schema migration which will create/update tables as needed.
-                self._run_schema_migration(silent=silent)
+            except (AutogenerateDiffsDetected, CommandError) as e:
+                logger.warning(
+                    "Schema mismatch detected. "
+                    "Your code models differ from the database schema. "
+                    "This may cause errors if your code references columns or tables that don't exist. "
+                    f"Schema was NOT modified. Details: {e}"
+                )
 
         super().__init__()