feat: add Pydantic V2 models for --json output (#903)

* feat: add Pydantic V2 models for --json output (STOPS-7491)

Define structured JSON output models for the 16 pgbelt commands
PGBaaS invokes. Simple commands (setup, teardown variants, analyze,
load-constraints, revoke-logins) use CommandResult directly with
steps + detail. Commands with rich per-item output (precheck, status,
sync-sequences, sync-tables, validate-data, create-indexes,
check-connectivity, connections) get dedicated subclasses.

Includes round-trip serialization/deserialization tests for all models.

Made-with: Cursor

* fix: isolate poetry from project venv in Dockerfile

Poetry and its transitive deps (cryptography, pycparser, cffi) were
installed into the same venv as the project. When poetry install ran,
it tried to downgrade those packages to match the lock file, corrupting
dist-info and breaking pip. Install poetry into a separate venv
(/opt/poetry) so it never conflicts with project dependencies.

Made-with: Cursor

Author: vjeeva

Dockerfile

@@ -1,7 +1,8 @@
 ARG PYTHON_VERSION=3.13
 FROM python:${PYTHON_VERSION}-slim
 ENV VIRTUAL_ENV=/opt/venv
-ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+ENV POETRY_HOME=/opt/poetry
+ENV PATH="$POETRY_HOME/bin:$VIRTUAL_ENV/bin:$PATH"
 
 RUN set -e \
     && apt-get -y update
@@ -11,9 +12,10 @@ RUN apt-get -y install \
     gcc
 
 RUN set -e \
-    && python -m venv $VIRTUAL_ENV \
-    && python -m pip install --upgrade pip \
-    && pip install poetry poetry-dynamic-versioning
+    && python -m venv $POETRY_HOME \
+    && $POETRY_HOME/bin/pip install poetry poetry-dynamic-versioning
+
+RUN python -m venv $VIRTUAL_ENV
 
 COPY ./ /opt/pgbelt
 WORKDIR /opt/pgbelt

pgbelt/models/__init__.py

@@ -0,0 +1,62 @@
+"""
+Pydantic V2 models for ``belt --json`` output.
+
+When to use ``CommandResult`` directly vs. create a subclass
+------------------------------------------------------------
+
+Use **CommandResult** directly when the command's output is just pass/fail
+plus optional step tracking and a handful of extras.  Put command-specific
+values in the ``detail`` dict.  Examples: setup, teardown variants, analyze,
+load-constraints, revoke-logins.
+
+Create a **subclass** of ``CommandResult`` when the command produces
+per-item structured data that a consumer would iterate, filter, or
+aggregate -- e.g. per-table validation results, per-index status,
+per-sequence sync detail.  The subclass lives in its own file grouped
+by domain (sync.py, schema.py, etc.).
+
+Breaking changes to any model require a pgbelt version bump because
+PGBaaS deserializes against the same types.
+"""
+
+from pgbelt.models.base import CommandError
+from pgbelt.models.base import CommandResult
+from pgbelt.models.base import StepResult
+from pgbelt.models.base import StepStatus
+from pgbelt.models.connectivity import ConnectivityCheckResult
+from pgbelt.models.connectivity import ConnectivityCheckRow
+from pgbelt.models.connections import ConnectionsResult
+from pgbelt.models.connections import ConnectionsRow
+from pgbelt.models.connections import ConnectionsSide
+from pgbelt.models.preflight import PrecheckResult
+from pgbelt.models.preflight import PrecheckSide
+from pgbelt.models.schema import CreateIndexesResult
+from pgbelt.models.schema import IndexDetail
+from pgbelt.models.status import StatusResult
+from pgbelt.models.status import StatusRow
+from pgbelt.models.sync import SyncSequencesResult
+from pgbelt.models.sync import SyncTablesResult
+from pgbelt.models.sync import ValidateDataResult
+
+__all__ = [
+    # Base -- used directly by simple commands (setup, teardown, analyze, etc.)
+    "CommandError",
+    "CommandResult",
+    "StepResult",
+    "StepStatus",
+    # Rich result models -- commands with structured per-item output
+    "ConnectivityCheckResult",
+    "ConnectivityCheckRow",
+    "ConnectionsResult",
+    "ConnectionsRow",
+    "ConnectionsSide",
+    "PrecheckResult",
+    "PrecheckSide",
+    "CreateIndexesResult",
+    "IndexDetail",
+    "StatusResult",
+    "StatusRow",
+    "SyncSequencesResult",
+    "SyncTablesResult",
+    "ValidateDataResult",
+]

pgbelt/models/base.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from datetime import datetime
+from datetime import timezone
+from enum import Enum
+from typing import Any
+from typing import Optional
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class StepStatus(str, Enum):
+    ok = "ok"
+    skipped = "skipped"
+    failed = "failed"
+
+
+class CommandError(BaseModel):
+    """Structured error information from a failed command."""
+
+    error_type: str
+    message: str
+    detail: Optional[str] = None
+
+
+class StepResult(BaseModel):
+    """Outcome of a single phase/step within a command."""
+
+    name: str
+    status: StepStatus
+    message: Optional[str] = None
+    duration_ms: Optional[int] = None
+
+
+class CommandResult(BaseModel):
+    """
+    Base model for all pgbelt --json output.
+
+    Simple commands (setup, teardown variants, analyze, load-constraints) use
+    this directly -- they only need pass/fail, steps, and maybe a few extras
+    in ``detail``.  Commands with rich structured output (precheck, status,
+    sync-sequences, etc.) extend this with their own fields.
+    """
+
+    db: str
+    dc: str
+    command: str
+    success: bool
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    duration_ms: Optional[int] = None
+    error: Optional[CommandError] = None
+    steps: list[StepResult] = []
+    detail: dict[str, Any] = {}

pgbelt/models/connections.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+
+from pydantic import BaseModel
+
+from pgbelt.models.base import CommandResult
+
+
+class ConnectionsSide(BaseModel):
+    """Connection summary for one side of a database pair."""
+
+    total_connections: int
+    by_user: dict[str, int] = {}
+
+
+class ConnectionsRow(BaseModel):
+    """Connection info for a single database pair."""
+
+    db: str
+    source: ConnectionsSide
+    destination: ConnectionsSide
+
+
+class ConnectionsResult(CommandResult):
+    """JSON output for ``belt connections``."""
+
+    command: str = "connections"
+    exclude_users: list[str] = []
+    exclude_patterns: list[str] = []
+    results: list[ConnectionsRow] = []

pgbelt/models/connectivity.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+
+from pydantic import BaseModel
+from pydantic import computed_field
+
+from pgbelt.models.base import CommandResult
+
+
+class ConnectivityCheckRow(BaseModel):
+    """Connectivity results for a single database pair."""
+
+    db: str
+    src_tcp: bool
+    src_query: bool
+    src_to_dst_dblink: bool
+    dst_tcp: bool
+    dst_query: bool
+    dst_to_src_dblink: bool
+
+    @computed_field
+    @property
+    def all_ok(self) -> bool:
+        return all(
+            [
+                self.src_tcp,
+                self.src_query,
+                self.src_to_dst_dblink,
+                self.dst_tcp,
+                self.dst_query,
+                self.dst_to_src_dblink,
+            ]
+        )
+
+    @computed_field
+    @property
+    def failed_checks(self) -> list[str]:
+        checks = {
+            "src_tcp": self.src_tcp,
+            "src_query": self.src_query,
+            "src_to_dst_dblink": self.src_to_dst_dblink,
+            "dst_tcp": self.dst_tcp,
+            "dst_query": self.dst_query,
+            "dst_to_src_dblink": self.dst_to_src_dblink,
+        }
+        return [name for name, passed in checks.items() if not passed]
+
+
+class ConnectivityCheckResult(CommandResult):
+    """JSON output for ``belt check-connectivity``."""
+
+    command: str = "check-connectivity"
+    results: list[ConnectivityCheckRow] = []
+
+    @computed_field
+    @property
+    def overall_ok(self) -> bool:
+        return all(r.all_ok for r in self.results)

pgbelt/models/preflight.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+from pydantic import computed_field
+
+from pgbelt.models.base import CommandResult
+
+
+class RoleInfo(BaseModel):
+    """Role/user details from pg_roles."""
+
+    rolname: str
+    rolcanlogin: bool
+    rolcreaterole: bool
+    rolinherit: bool
+    rolsuper: bool
+    memberof: list[str] = []
+    can_create: Optional[bool] = None
+
+
+class RelationInfo(BaseModel):
+    """A table or sequence discovered in the schema."""
+
+    name: str
+    schema_name: str
+    owner: str
+    object_type: str
+
+    @computed_field
+    @property
+    def can_replicate(self) -> bool:
+        """True when schema and owner match the targeted migration config."""
+        return True  # Actual check requires runtime context; populated by caller.
+
+
+class TableReplicationInfo(BaseModel):
+    """Table with replication-method classification (source side only)."""
+
+    name: str
+    schema_name: str
+    owner: str
+    has_primary_key: bool
+    replication_method: str  # "pglogical" | "dump_and_load" | "unavailable"
+
+
+class ExtensionInfo(BaseModel):
+    """Installed Postgres extension."""
+
+    extname: str
+    in_other_side: Optional[bool] = None
+
+
+class PrecheckSide(BaseModel):
+    """Precheck data for one side (source or destination) of a migration pair."""
+
+    db: str
+    schema_name: str
+    server_version: str
+    max_replication_slots: str
+    max_worker_processes: str
+    max_wal_senders: str
+    shared_preload_libraries: list[str] = []
+    rds_logical_replication: str
+    root_user: RoleInfo
+    owner_user: RoleInfo
+    tables: list[TableReplicationInfo] = []
+    sequences: list[RelationInfo] = []
+    extensions: list[ExtensionInfo] = []
+
+    @computed_field
+    @property
+    def root_ok(self) -> bool:
+        return (
+            self.root_user.rolcanlogin
+            and self.root_user.rolcreaterole
+            and self.root_user.rolinherit
+            and ("rds_superuser" in self.root_user.memberof or self.root_user.rolsuper)
+        )
+
+    @computed_field
+    @property
+    def shared_preload_ok(self) -> bool:
+        return (
+            "pglogical" in self.shared_preload_libraries
+            and "pg_stat_statements" in self.shared_preload_libraries
+        )
+
+
+class PrecheckResult(CommandResult):
+    """JSON output for ``belt precheck``."""
+
+    command: str = "precheck"
+    src: Optional[PrecheckSide] = None
+    dst: Optional[PrecheckSide] = None
+
+    @computed_field
+    @property
+    def extensions_match(self) -> Optional[bool]:
+        if self.src is None or self.dst is None:
+            return None
+        src_names = {e.extname for e in self.src.extensions}
+        dst_names = {e.extname for e in self.dst.extensions}
+        return src_names == dst_names

pgbelt/models/schema.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from pgbelt.models.base import CommandResult
+
+
+class IndexDetail(BaseModel):
+    """Result for a single CREATE INDEX operation."""
+
+    name: str
+    status: str  # "created" | "skipped_exists" | "failed"
+    duration_ms: Optional[int] = None
+    error: Optional[str] = None
+
+
+class CreateIndexesResult(CommandResult):
+    """JSON output for ``belt create-indexes``."""
+
+    command: str = "create-indexes"
+    indexes_file: Optional[str] = None
+    indexes: list[IndexDetail] = []
+    analyze_ran: bool = False
+
+    @property
+    def created_count(self) -> int:
+        return sum(1 for i in self.indexes if i.status == "created")
+
+    @property
+    def skipped_count(self) -> int:
+        return sum(1 for i in self.indexes if i.status == "skipped_exists")
+
+    @property
+    def failed_count(self) -> int:
+        return sum(1 for i in self.indexes if i.status == "failed")