state_versions.py

# Copyright IBM Corp. 2025, 2026
# SPDX-License-Identifier: MPL-2.0

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path

from pytfe import TFEClient, TFEConfig
from pytfe.errors import ErrStateVersionUploadNotSupported
from pytfe.models import (
    StateVersionCreateOptions,
    StateVersionCurrentOptions,
    StateVersionListOptions,
    StateVersionOutputsListOptions,
)
from pytfe.models.workspace import WorkspaceLockOptions


def _print_header(title: str):
    print("\n" + "=" * 80)
    print(title)
    print("=" * 80)


def main():
    parser = argparse.ArgumentParser(
        description="State Versions demo for python-tfe SDK"
    )
    parser.add_argument(
        "--address", default=os.getenv("TFE_ADDRESS", "https://app.terraform.io")
    )
    parser.add_argument("--token", default=os.getenv("TFE_TOKEN", ""))
    parser.add_argument("--org", required=True, help="Organization name")
    parser.add_argument("--workspace", required=True, help="Workspace name")
    parser.add_argument("--workspace-id", required=True, help="Workspace ID")
    parser.add_argument("--download", help="Path to save downloaded current state")
    parser.add_argument("--upload", help="Path to a .tfstate (or JSON state) to upload")
    parser.add_argument("--page-size", type=int, default=10)
    parser.add_argument(
        "--rollback-to",
        help="State version id to roll the workspace back to. The workspace "
        "will be locked, rolled back, then unlocked.",
    )
    parser.add_argument(
        "--rollback-dry-run",
        action="store_true",
        help="With --rollback-to, print the plan without performing the rollback.",
    )
    args = parser.parse_args()

    cfg = TFEConfig(address=args.address, token=args.token)
    client = TFEClient(cfg)

    options = StateVersionListOptions(
        page_size=args.page_size,
        organization=args.org,
        workspace=args.workspace,
    )

    sv_list = list(client.state_versions.list(options))
    print(f"Total state versions: {len(sv_list)}")
    print()

    for sv in sv_list:
        print(f"- {sv.id} | status={sv.status} | created_at={sv.created_at}")

    # 1) List all state versions across org and workspace filters
    _print_header("Org-scoped listing via /api/v2/state-versions (first page)")
    all_sv = client.state_versions.list(
        StateVersionListOptions(
            organization=args.org, workspace=args.workspace, page_size=args.page_size
        )
    )
    for sv in all_sv:
        print(f"- {sv.id} | status={sv.status} | created_at={sv.created_at}")

    # 2) Read the current state version (with outputs included if you want)
    _print_header("Reading current state version")
    current = client.state_versions.read_current_with_options(
        args.workspace_id, StateVersionCurrentOptions(include=["outputs"])
    )
    print(
        f"Current SV: {current.id} status={current.status} durl={current.hosted_state_download_url}"
    )

    # 3) (Optional) Download the current state (optional)
    if args.download:
        _print_header(f"Downloading current state to: {args.download}")
        raw = client.state_versions.download(current.id)
        Path(args.download).write_bytes(raw)
        print(f"Wrote {len(raw)} bytes to {args.download}")

    # 4) List outputs for the current state version (paged)
    _print_header("Listing outputs (current state version)")
    outs = list(
        client.state_versions.list_outputs(
            current.id, options=StateVersionOutputsListOptions(page_size=50)
        )
    )
    if not outs:
        print("No outputs found.")
    for o in outs:
        # Sensitive outputs will have value = None
        print(f"- {o.name}: sensitive={o.sensitive} type={o.type} value={o.value}")

    if args.workspace_id:
        # 4b) List outputs for the current state version via workspace endpoint
        _print_header("Listing outputs via workspace endpoint")
        outs2 = list(
            client.state_version_outputs.read_current(
                args.workspace_id, options=StateVersionOutputsListOptions(page_size=50)
            )
        )
        if not outs2:
            print("No outputs found.")
        for o in outs2:
            print(f"- {o.name}: sensitive={o.sensitive} type={o.type} value={o.value}")

    # 5) (Optional) Upload a new state file
    if args.upload:
        _print_header(f"Uploading new state from: {args.upload}")
        try:
            payload = Path(args.upload).read_bytes()
            state_obj = json.loads(payload.decode("utf-8"))
            serial = int(state_obj["serial"])
            lineage = state_obj.get("lineage")
            md5 = hashlib.md5(payload).hexdigest()  # nosec B324
            locked_workspace = False

            try:
                client.workspaces.lock(
                    args.workspace_id,
                    WorkspaceLockOptions(
                        reason="python-tfe state_versions upload example"
                    ),
                )
                locked_workspace = True
            except Exception:
                # Continue in case the workspace is already locked by the caller.
                pass

            # If your server supports signed uploads, this will:
            #   a) create SV (to get upload URL)
            #   b) PUT bytes to the signed URL
            #   c) read back the SV to return a hydrated object
            try:
                new_sv = client.state_versions.upload(
                    args.workspace_id,
                    raw_state=payload,
                    options=StateVersionCreateOptions(
                        serial=serial,
                        md5=md5,
                        lineage=lineage,
                    ),
                )
            finally:
                if locked_workspace:
                    client.workspaces.unlock(args.workspace_id)
            print(f"Uploaded new SV: {new_sv.id} status={new_sv.status}")
        except FileNotFoundError:
            print(f"Upload file not found: {args.upload}")
        except (KeyError, ValueError, json.JSONDecodeError):
            print(
                "Upload input must be a valid Terraform state JSON containing at least a serial value."
            )
        except ErrStateVersionUploadNotSupported as e:
            # Some older/self-hosted versions don’t support direct upload
            print(f"Upload not supported on this server: {e}")

    # 6) (Optional) Roll back to a previous state version
    if args.rollback_to:
        _print_header(
            f"Rolling {args.workspace_id} back to state version {args.rollback_to}"
        )
        if args.rollback_dry_run:
            print("--rollback-dry-run set; not locking or rolling back")
        else:
            print("locking workspace ...")
            client.workspaces.lock(
                args.workspace_id,
                WorkspaceLockOptions(reason="python-tfe rollback demo"),
            )
            try:
                new_sv = client.state_versions.rollback(
                    args.workspace_id, args.rollback_to
                )
                print(
                    f"rollback succeeded — new state version: {new_sv.id} "
                    f"(serial={new_sv.serial})"
                )
            finally:
                print("unlocking workspace ...")
                client.workspaces.unlock(args.workspace_id)


if __name__ == "__main__":
    main()