Better support env specific sql

README.md

@@ -146,6 +146,32 @@ semantics. `pgdevkit.migrate.list_migration_files`/`pending_migrations` and
 `pgdevkit.fetch_missing.find_missing_objects` takes neither, for the reason
 above.
 
+## Environment-tagged files (`<name>.<env>.sql`)
+
+A file whose name ends `.<env>.sql` (e.g. `grants.prod.sql`,
+`seed.staging.sql`) is only in scope when targeting that environment; a
+plain `<name>.sql` file is untagged and always in scope, regardless of
+environment. `.init.sql` (see `docs/database-layout.md`) is reserved and is
+never treated as an environment tag.
+
+- `pgdb testdb up`/`pgdb testdb reset` accept `--env` (default
+  `local_test`) — so an untagged `grants.sql` always applies, but
+  `grants.prod.sql` is skipped unless run with `--env prod`.
+- `pgdb migrate check`/`pgdb migrate apply` accept `--env` too, but it's
+  optional with **no** default: omit it and every file is a candidate
+  regardless of its tag (unchanged, today's behavior); pass it to restrict
+  to files tagged for that environment plus untagged ones.
+
+```bash
+pgdb testdb up --env prod   # apply prod-tagged files too, against the local test container
+pgdb migrate apply path/to/database/_migration_scripts --url ... --env prod
+```
+
+`pgdevkit.envtag` exposes the same logic for scripting: `file_env` reads a
+file's tag, `env_allowed` applies the filtering semantics above, and
+`strip_env_suffix` returns a tagged file's logical name (e.g.
+`grants.prod.sql` -> `"grants"`).
+
 ## `pgdb testdb`
 
 Manages a single shared, Podman-backed Postgres container for local tests
@@ -174,7 +200,8 @@ def ensure_test_postgres():
         os.environ[k] = v
 ```
 
-CLI: `pgdb testdb up|reset|run-sql|status|shell|clean`.
+CLI: `pgdb testdb up|reset|run-sql|status|shell|clean`. `up`/`reset` accept
+`--env` (default `local_test`) — see "Environment-tagged files" above.
 
 `up`/`reset` accept `--area`/`--exclude-area` and `--schema`/`--exclude-schema`
 (see "Area and schema filtering" above) to scope which `database/` files get

docs/database-layout.md

@@ -71,10 +71,18 @@ One object per file: `tables/user.sql`, `views/all_edits.sql`,
 |---|---|
 | `<name>.sql` | The object's live definition (`CREATE TABLE`, `CREATE OR REPLACE VIEW`, ...) |
 | `<name>.test_data.json` | Seed rows for a table — a JSON array of row objects, loaded after the table is created |
-| `<name>.init.sql` | One-time setup for an object (e.g. a backfill), run once, kept separate from the reusable definition |
-| `<name>.prod.sql` / `.prod` anywhere in the name | Production-only (real permission grants, real user accounts) — skipped by `pgdb testdb` |
+| `<name>.init.sql` | One-time setup for an object (e.g. a backfill), run once, kept separate from the reusable definition — `init` is reserved and is never treated as an environment tag |
+| `<name>.<env>.sql` | Only applied when targeting environment `<env>` (any name you like — `prod`, `staging`, ...); a file with no such suffix is untagged and always applies, regardless of environment |
 | `all.sql` | Generated concatenation of the whole tree — not hand-edited, not committed |
 
+`pgdb testdb up`/`pgdb testdb reset` apply the `--env` they're given (default
+`local_test`) — so an untagged `grants.sql` always applies, but
+`grants.prod.sql` is skipped unless you pass `--env prod`. `pgdb migrate
+check`/`pgdb migrate apply` accept the same `--env`, but it's optional with no
+default: omit it and every file is a candidate regardless of its tag (today's
+behavior); pass it to restrict to files tagged for that environment plus
+untagged ones.
+
 ---
 
 ## Migrations
@@ -156,5 +164,5 @@ leading sort number.
 - [ ] Object-type folder (`tables`, `views`, ...) matches the apply-order table above — that's what governs ordering, not the layer's leading number
 - [ ] One-off changes go in `migrations/`, dated, never edited after applying
 - [ ] The live `.sql` file is updated in the same change as any migration touching that object
-- [ ] `.prod` files are production-only and skipped by `pgdb testdb`
+- [ ] `.<env>.sql` files (e.g. `.prod.sql`) are skipped by `pgdb testdb` unless it's run with a matching `--env`
 - [ ] Every table (and non-obvious column) has a `COMMENT ON`, placed in the object's own `.sql` file

pgdevkit/cli.py

@@ -41,6 +41,17 @@
     help="Skip files referencing this DB schema (repeatable); "
     "files with no detectable schema reference are never excluded",
 )
+_TESTDB_ENV_OPTION = typer.Option(
+    "local_test",
+    "--env",
+    help="Environment to apply: skips any <name>.<other-env>.sql file (e.g. grants.prod.sql); untagged files always apply",
+)
+_MIGRATE_ENV_OPTION = typer.Option(
+    None,
+    "--env",
+    help="Restrict to files tagged for this environment (e.g. <name>.prod.sql); untagged files always apply; "
+    "omit to apply every file regardless of its env tag",
+)
 
 
 def _as_set(values: list[str]) -> frozenset[str] | None:
@@ -209,14 +220,15 @@ def fetch_missing(
 
 @testdb_app.command("up")
 def testdb_up(
+    env: str = _TESTDB_ENV_OPTION,
     area: list[str] = _AREA_OPTION,
     exclude_area: list[str] = _EXCLUDE_AREA_OPTION,
     schema: list[str] = _SCHEMA_OPTION,
     exclude_schema: list[str] = _EXCLUDE_SCHEMA_OPTION,
 ) -> None:
     """Ensure the container is running, the workspace DB exists, and schema is applied."""
     testdb.ensure_testdb(
-        areas=_as_set(area), exclude_areas=_as_set(exclude_area), schemas=_as_set(schema),
+        env=env, areas=_as_set(area), exclude_areas=_as_set(exclude_area), schemas=_as_set(schema),
         exclude_schemas=_as_set(exclude_schema),
     )
     info = testdb.status()
@@ -225,14 +237,15 @@ def testdb_up(
 
 @testdb_app.command("reset")
 def testdb_reset(
+    env: str = _TESTDB_ENV_OPTION,
     area: list[str] = _AREA_OPTION,
     exclude_area: list[str] = _EXCLUDE_AREA_OPTION,
     schema: list[str] = _SCHEMA_OPTION,
     exclude_schema: list[str] = _EXCLUDE_SCHEMA_OPTION,
 ) -> None:
     """Drop and recreate only this workspace's database, then reapply schema + seed data."""
     testdb.reset_testdb(
-        areas=_as_set(area), exclude_areas=_as_set(exclude_area), schemas=_as_set(schema),
+        env=env, areas=_as_set(area), exclude_areas=_as_set(exclude_area), schemas=_as_set(schema),
         exclude_schemas=_as_set(exclude_schema),
     )
     info = testdb.status()
@@ -310,6 +323,7 @@ def migrate_check(
     exclude_area: list[str] = _EXCLUDE_AREA_OPTION,
     schema: list[str] = _SCHEMA_OPTION,
     exclude_schema: list[str] = _EXCLUDE_SCHEMA_OPTION,
+    env: str | None = _MIGRATE_ENV_OPTION,
 ) -> None:
     """List which migration files under migrations_dir are applied vs. pending."""
     if not migrations_dir.is_dir():
@@ -324,6 +338,7 @@ def migrate_check(
         exclude_areas=_as_set(exclude_area),
         schemas=_as_set(schema),
         exclude_schemas=_as_set(exclude_schema),
+        env=env,
     )
     try:
         applied = migrate.applied_migrations(conninfo, tracking_table)
@@ -367,6 +382,7 @@ def migrate_apply(
     exclude_area: list[str] = _EXCLUDE_AREA_OPTION,
     schema: list[str] = _SCHEMA_OPTION,
     exclude_schema: list[str] = _EXCLUDE_SCHEMA_OPTION,
+    env: str | None = _MIGRATE_ENV_OPTION,
 ) -> None:
     """Apply pending migration files, in filename order, tracking each in tracking_table."""
     if not migrations_dir.is_dir():
@@ -393,14 +409,15 @@ def migrate_apply(
                 exclude_areas=exclude_areas,
                 schemas=schemas,
                 exclude_schemas=exclude_schemas,
+                env=env,
             )
         except migrate.TrackingTableMissing:
             err_console.print(
                 f"[yellow]⚠[/yellow]  {tracking_table} not found — treating every migration as pending"
             )
             targets = migrate.list_migration_files(
                 migrations_dir, areas=areas, exclude_areas=exclude_areas, schemas=schemas,
-                exclude_schemas=exclude_schemas,
+                exclude_schemas=exclude_schemas, env=env,
             )
 
     if not targets:

pgdevkit/envtag.py

@@ -0,0 +1,54 @@
+"""Optional `<name>.<env>.sql` filename convention: a file whose dot-segment
+immediately before `.sql` names a deployment environment (e.g.
+`grants.prod.sql`, `seed.staging.sql`) is only in scope when the caller is
+targeting that same environment. A plain `<name>.sql` file (no such segment)
+is untagged/common and is always in scope, regardless of which environment is
+requested — mirroring the untagged-file rule for `-- area:` tags in
+areas.py.
+
+This generalizes the older, hardcoded `.prod.sql` convention (still the usual
+name for a production-only file — grants, real user accounts — that
+`pgdb testdb` should never touch); any string can now be used as an
+environment name.
+
+`.init.sql` (one-time setup, see docs/database-layout.md) is reserved and is
+never interpreted as an environment tag.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+_RESERVED_SQL_SUFFIXES = {"init"}
+
+
+def file_env(path: Path) -> str | None:
+    """The environment tag from `path`'s name, or None if it's untagged (or
+    the suffix is a reserved, non-env one like `.init.sql`). Only `.sql`
+    files can carry a tag."""
+    if path.suffix != ".sql":
+        return None
+    stem = path.stem
+    base, dot, suffix = stem.rpartition(".")
+    if not dot or suffix in _RESERVED_SQL_SUFFIXES:
+        return None
+    return suffix
+
+
+def env_allowed(path: Path, env: str | None) -> bool:
+    """Whether `path` is in scope for `env`. `env=None` means no environment
+    filtering was requested, so every file (tagged or not) is in scope."""
+    if env is None:
+        return True
+    tag = file_env(path)
+    return tag is None or tag == env
+
+
+def strip_env_suffix(path: Path) -> str:
+    """`path.stem` with a trailing `.<tag>` removed, so a tagged file
+    resolves to the same logical name as its untagged counterpart would
+    (e.g. `grants.prod.sql` -> "grants", same as `grants.sql`)."""
+    tag = file_env(path)
+    if tag is None:
+        return path.stem
+    return path.stem[: -(len(tag) + 1)]

pgdevkit/migrate.py

@@ -19,6 +19,7 @@
 from psycopg import sql as pg_sql
 
 from .areas import filter_by_area
+from .envtag import env_allowed
 from .schemas import filter_by_schema
 from .sql_text import strip_line_comments
 
@@ -236,10 +237,17 @@ def list_migration_files(
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
     exclude_schemas: frozenset[str] | None = None,
+    env: str | None = None,
 ) -> list[Path]:
+    """Migration files under migrations_dir, restricted by area (see
+    `.areas`), by schema (see `.schemas`), and by environment tag (see
+    `.envtag`) — e.g. a `2026-07-10_backfill.prod.sql` is only included when
+    `env="prod"`. `env=None` (the default) applies no environment filtering
+    at all, so every file is a candidate regardless of its tag."""
     files = sorted(migrations_dir.glob("*.sql"))
     files = filter_by_area(files, only=areas, exclude=exclude_areas)
-    return filter_by_schema(files, only=schemas, exclude=exclude_schemas)
+    files = filter_by_schema(files, only=schemas, exclude=exclude_schemas)
+    return [f for f in files if env_allowed(f, env)]
 
 
 def applied_migrations(conninfo: str, tracking_table: str) -> dict[str, tuple[datetime, str]]:
@@ -264,10 +272,16 @@ def pending_migrations(
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
     exclude_schemas: frozenset[str] | None = None,
+    env: str | None = None,
 ) -> list[Path]:
     applied = applied_migrations(conninfo, tracking_table)
     files = list_migration_files(
-        migrations_dir, areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas
+        migrations_dir,
+        areas=areas,
+        exclude_areas=exclude_areas,
+        schemas=schemas,
+        exclude_schemas=exclude_schemas,
+        env=env,
     )
     return [p for p in files if p.name not in applied]
 

pgdevkit/testdb/api.py

@@ -73,6 +73,7 @@ async def _apply(
     db_name: str,
     force_reset: bool,
     *,
+    env: str = "local_test",
     areas: frozenset[str] | None = None,
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
@@ -84,6 +85,7 @@ async def _apply(
             config.root / config.database_dir,
             extensions=config.extensions,
             force_reset=force_reset,
+            env=env,
             areas=areas,
             exclude_areas=exclude_areas,
             schemas=schemas,
@@ -95,6 +97,7 @@ def ensure_testdb(
     project_root: Path | None = None,
     force_reset: bool = False,
     *,
+    env: str = "local_test",
     areas: frozenset[str] | None = None,
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
@@ -105,6 +108,9 @@ def ensure_testdb(
     vars for this workspace (or the mssql equivalent's env vars, per
     `config.engine`).
 
+    `env` selects which environment-tagged files apply (see pgdevkit.envtag,
+    e.g. a `grants.prod.sql` is skipped unless env="prod").
+
     `areas`/`exclude_areas` and `schemas`/`exclude_schemas` restrict which
     database/ files get applied -- e.g. for a test DB scoped to one area or
     schema. Neither filters what gets *dropped* by force_reset/clean, only
@@ -113,7 +119,7 @@ def ensure_testdb(
     if config.engine == "mssql":
         return _mssql_api().ensure_testdb(
             config, db_name, force_reset,
-            areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
+            env=env, areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
         )
 
     ensure_container()
@@ -124,7 +130,7 @@ async def _run() -> None:
         await _ensure_database(db_name)
         await _apply(
             config, db_name, force_reset,
-            areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
+            env=env, areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
         )
 
     asyncio.run(_run())
@@ -134,6 +140,7 @@ async def _run() -> None:
 def reset_testdb(
     project_root: Path | None = None,
     *,
+    env: str = "local_test",
     areas: frozenset[str] | None = None,
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
@@ -143,7 +150,7 @@ def reset_testdb(
     schema and seed data."""
     return ensure_testdb(
         project_root, force_reset=True,
-        areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
+        env=env, areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
     )
 
 

pgdevkit/testdb/mssql/api.py

@@ -9,6 +9,7 @@
 
 from ...db.mssql_sql import ident, json_encode_value
 from ...dialect import MSSQL
+from ...envtag import strip_env_suffix
 from .. import query
 from ..config import ProjectConfig
 from ..schema import _iter_sql_files, _strip_layer_prefix
@@ -106,6 +107,7 @@ async def _apply(
     db_name: str,
     force_reset: bool,
     *,
+    env: str = "local_test",
     areas: frozenset[str] | None = None,
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
@@ -121,19 +123,21 @@ def _connect() -> Any:
             return
         for file, sql in _iter_sql_files(
             database_dir, MSSQL,
-            areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
+            env=env, areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
         ):
             for batch in query.split_tsql_batches(sql):
 
                 def _exec(batch: str = batch) -> None:
                     conn.cursor().execute(batch)
 
                 await asyncio.to_thread(_exec)
-            json_file = file.with_suffix(".test_data.json")
+            table_stem = strip_env_suffix(file)
+            json_file = file.parent / f"{table_stem}.test_data.json"
             if json_file.exists():
                 schema_name = _strip_layer_prefix(file.parent.parent.name)
-                table_stem = _strip_layer_prefix(file.stem)
-                await _insert_test_data(json_file, f"{schema_name}.{table_stem}", force_reset, conn)
+                await _insert_test_data(
+                    json_file, f"{schema_name}.{_strip_layer_prefix(table_stem)}", force_reset, conn
+                )
     finally:
         await asyncio.to_thread(conn.close)
 
@@ -143,6 +147,7 @@ def ensure_testdb(
     db_name: str,
     force_reset: bool,
     *,
+    env: str = "local_test",
     areas: frozenset[str] | None = None,
     exclude_areas: frozenset[str] | None = None,
     schemas: frozenset[str] | None = None,
@@ -156,7 +161,7 @@ async def _run() -> None:
         await _ensure_database(db_name)
         await _apply(
             config, db_name, force_reset,
-            areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
+            env=env, areas=areas, exclude_areas=exclude_areas, schemas=schemas, exclude_schemas=exclude_schemas,
         )
 
     asyncio.run(_run())