validate_env.py

#!/usr/bin/env python3
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Environment Variables Validation Script

Validates all required environment variables (playground and deployment).

Usage:
    python validate_env.py [--env-file=.env]
"""

import argparse
import os
import sys
from typing import Dict, List

# All required variables (var_name -> description). Matches .env.
REQUIRED_VARS: Dict[str, str] = {
    "APP_ID_GITHUB": "GitHub App ID",
    "INSTALLATION_ID_GITHUB": "GitHub App installation ID",
    "SECRET_MANAGER_PROJECT_ID": "Project ID or number for Secret Manager",
    "SECRET_MANAGER_SECRET_NAME": "Secret name for GitHub private key",
    "SECRET_MANAGER_TOKEN_NAME": "Secret name for GitHub token",
    "VERTEX_AI_LOCATION": "Vertex AI region (e.g. us-central1)",
    "VERTEX_AI_PROJECT_ID": "Vertex AI project number",
    "VERTEX_AI_RESOURCE_ID": "Vertex AI Agent Engine resource ID",
    "PROJECT_ID": "Google Cloud Project ID",
    "DEPLOYMENT_REGION": "GCP region for deployment",
    "MODEL_LOCATION": "GCP region for Gemini models",
    "LLM_LOCATION": "Location for LLM operations (e.g. global)",
    "DATA_STORE_REGION": "Vertex AI Search data store region",
    "DATA_STORE_ID": "Vertex AI Search data store ID",
    "STAGING_BUCKET": "GCS bucket URL (e.g. gs://code_review)",
    "LOGS_BUCKET_NAME": "GCS bucket name for logs",
    "DATASET_ID": "BigQuery dataset ID for telemetry",
    "TABLE_ID": "BigQuery table ID for events",
    "MODEL_NAME": "Gemini model name (e.g. gemini-2.5-flash)",
    "AGENT_DISPLAY_NAME": "Display name for the agent engine",
    "AGENT_DESCRIPTION": "Description of the agent",
    # Optional agent deploy vars (defaults in deploy.py): AGENT_SOURCE_PACKAGES=./agent,
    # AGENT_ENTRYPOINT_MODULE=agent_engine_app, AGENT_ENTRYPOINT_OBJECT=agent_engine,
    # AGENT_REQUIREMENTS_FILE=agent/app_utils/.requirements.txt,
    # AGENT_MIN_INSTANCES=1, AGENT_MAX_INSTANCES=10, AGENT_CPU=4, AGENT_MEMORY=8Gi,
    # AGENT_CONTAINER_CONCURRENCY=9, AGENT_NUM_WORKERS=1
}


def load_env_file(env_file: str = ".env") -> None:
    """Load environment variables from a .env file if it exists."""
    if not os.path.exists(env_file):
        return
    with open(env_file) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, value = line.split("=", 1)
                key = key.strip()
                value = value.strip()
                if value.startswith('"') and value.endswith('"'):
                    value = value[1:-1]
                elif value.startswith("'") and value.endswith("'"):
                    value = value[1:-1]
                if key not in os.environ and value:
                    os.environ[key] = value


def validate() -> tuple[bool, List[str]]:
    """Validate all required variables. Returns (success, list of missing var descriptions)."""
    missing = []
    for var_name, description in REQUIRED_VARS.items():
        value = os.getenv(var_name)
        if value is None or (isinstance(value, str) and value.strip() == ""):
            missing.append(f"  - {var_name}: {description}")
    return len(missing) == 0, missing


def main() -> int:
    parser = argparse.ArgumentParser(description="Validate required environment variables")
    parser.add_argument("--env-file", default=".env", help="Path to .env file (default: .env)")
    args = parser.parse_args()

    load_env_file(args.env_file)

    print("=" * 80)
    print("Environment Variables Validation")
    print("=" * 80)
    print()

    success, missing = validate()
    if success:
        print("✅ SUCCESS: All required environment variables are set!")
        return 0
    print("❌ ERROR: Missing required environment variables:")
    print()
    for var in missing:
        print(var)
    print()
    print("Set these in your environment or .env file.")
    return 1


if __name__ == "__main__":
    sys.exit(main())