Northern-Thailand-Ping-River-Monitor/scripts/setup_postgres.py

#!/usr/bin/env python3
"""
PostgreSQL setup script for Northern Thailand Ping River Monitor
This script helps you configure and test your PostgreSQL connection
"""

import os
import sys
import logging
from typing import Optional
from urllib.parse import urlparse

def setup_logging():
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

def test_postgres_connection(connection_string: str) -> bool:
    """Test connection to PostgreSQL database"""
    try:
        from sqlalchemy import create_engine, text

        # Test connection
        engine = create_engine(connection_string, pool_pre_ping=True)
        with engine.connect() as conn:
            result = conn.execute(text("SELECT version()"))
            version = result.fetchone()[0]
            logging.info(f"✅ Connected to PostgreSQL successfully!")
            logging.info(f"Database version: {version}")
            return True

    except ImportError:
        logging.error("❌ psycopg2-binary not installed. Run: uv add psycopg2-binary")
        return False
    except Exception as e:
        logging.error(f"❌ Connection failed: {e}")
        return False

def parse_connection_string(connection_string: str) -> dict:
    """Parse PostgreSQL connection string into components"""
    try:
        parsed = urlparse(connection_string)
        return {
            'host': parsed.hostname,
            'port': parsed.port or 5432,
            'database': parsed.path[1:] if parsed.path else None,
            'username': parsed.username,
            'password': parsed.password,
        }
    except Exception as e:
        logging.error(f"Failed to parse connection string: {e}")
        return {}

def create_database_if_not_exists(connection_string: str, database_name: str) -> bool:
    """Create database if it doesn't exist"""
    try:
        from sqlalchemy import create_engine, text

        # Connect to default postgres database to create our database
        parsed = urlparse(connection_string)
        admin_connection = connection_string.replace(f"/{parsed.path[1:]}", "/postgres")

        engine = create_engine(admin_connection, pool_pre_ping=True)

        with engine.connect() as conn:
            # Check if database exists
            result = conn.execute(text(
                "SELECT 1 FROM pg_database WHERE datname = :db_name"
            ), {"db_name": database_name})

            if result.fetchone():
                logging.info(f"✅ Database '{database_name}' already exists")
                return True
            else:
                # Create database
                conn.execute(text("COMMIT"))  # End transaction
                conn.execute(text(f'CREATE DATABASE "{database_name}"'))
                logging.info(f"✅ Created database '{database_name}'")
                return True

    except Exception as e:
        logging.error(f"❌ Failed to create database: {e}")
        return False

def initialize_tables(connection_string: str) -> bool:
    """Initialize database tables"""
    try:
        # Import the database adapter to create tables
        sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
        from database_adapters import SQLAdapter

        adapter = SQLAdapter(connection_string=connection_string, db_type='postgresql')
        if adapter.connect():
            logging.info("✅ Database tables initialized successfully")
            return True
        else:
            logging.error("❌ Failed to initialize tables")
            return False

    except Exception as e:
        logging.error(f"❌ Failed to initialize tables: {e}")
        return False

def interactive_setup():
    """Interactive setup wizard"""
    print("🐘 PostgreSQL Setup Wizard for Ping River Monitor")
    print("=" * 50)

    # Get connection details
    host = input("PostgreSQL host (e.g., 192.168.1.100): ").strip()
    port = input("PostgreSQL port [5432]: ").strip() or "5432"
    database = input("Database name [water_monitoring]: ").strip() or "water_monitoring"
    username = input("Username: ").strip()
    password = input("Password: ").strip()

    # Optional SSL
    use_ssl = input("Use SSL connection? (y/N): ").strip().lower() == 'y'
    ssl_params = "?sslmode=require" if use_ssl else ""

    connection_string = f"postgresql://{username}:{password}@{host}:{port}/{database}{ssl_params}"

    print(f"\nGenerated connection string:")
    print(f"POSTGRES_CONNECTION_STRING={connection_string}")

    return connection_string

def main():
    setup_logging()

    print("🚀 Northern Thailand Ping River Monitor - PostgreSQL Setup")
    print("=" * 60)

    # Check if connection string is provided via environment
    connection_string = os.getenv('POSTGRES_CONNECTION_STRING')

    if not connection_string:
        print("No POSTGRES_CONNECTION_STRING found in environment.")
        print("Starting interactive setup...\n")
        connection_string = interactive_setup()

        # Suggest adding to .env file
        print(f"\n💡 Add this to your .env file:")
        print(f"DB_TYPE=postgresql")
        print(f"POSTGRES_CONNECTION_STRING={connection_string}")

    # Parse connection details
    config = parse_connection_string(connection_string)
    if not config.get('host'):
        logging.error("Invalid connection string format")
        return False

    print(f"\n🔗 Connecting to PostgreSQL at {config['host']}:{config['port']}")

    # Test connection
    if not test_postgres_connection(connection_string):
        return False

    # Try to create database
    database_name = config.get('database', 'water_monitoring')
    if database_name:
        create_database_if_not_exists(connection_string, database_name)

    # Initialize tables
    if not initialize_tables(connection_string):
        return False

    print("\n🎉 PostgreSQL setup completed successfully!")
    print("\nNext steps:")
    print("1. Update your .env file with the connection string")
    print("2. Run: make run-test")
    print("3. Run: make run-api")

    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)