dotfiles/local-bin/.local/bin/generate_tinyrdm_connections.py

#!/usr/bin/env python3
"""
Generate TinyRDM connections.yaml from sdm status --json output
Organizes Redis connections into Internal, Stage, and Production groups
Excludes activate and readonly connections
"""

import subprocess
import json
import re
import os
import yaml
from collections import defaultdict

def run_sdm_status():
    try:
        result = subprocess.run(['sdm', 'status', '--json'], capture_output=True, text=True)
        if result.returncode != 0:
            print(f"Error running sdm status: {result.stderr}")
            return None
        return json.loads(result.stdout)
    except FileNotFoundError:
        print("Error: sdm command not found")
        return None
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON: {e}")
        return None

def parse_redis_connections(sdm_data):
    connections = []
    excluded = ['activate-cache', 'activate-readonly', 'readonly-redis']
    
    for item in sdm_data:
        if item.get('type') != 'redis':
            continue
            
        name = item.get('name', '')
        if any(x in name for x in excluded):
            continue
        
        address = item.get('address', '')
        if ':' not in address:
            continue
        
        addr, port = address.split(':')
        env_info = parse_connection_name(name)
        
        if env_info:
            connections.append({
                'name': name,
                'addr': addr,
                'port': int(port),
                **env_info
            })
    
    return connections

def parse_connection_name(name):
    pattern = r'oc-([^-]+)-([^-]+)-.*'
    match = re.match(pattern, name)
    
    if not match:
        return None
    
    environment = match.group(1)
    customer = match.group(2)
    
    if environment in ['dev', 'nextrc', 'nextrc2']:
        stage = 'internal'
    elif environment in ['stage', 'stage2', 'uat']:
        stage = 'stage'
    elif environment in ['prod', 'prod2']:
        stage = 'production'
    else:
        stage = 'unknown'
    
    return {
        'environment': environment,
        'customer': customer,
        'stage': stage
    }

def is_sdm_connection(conn_name):
    """Check if connection name matches SDM pattern (oc-*-*-*)"""
    pattern = r'oc-([^-]+)-([^-]+)-.*'
    return re.match(pattern, conn_name) is not None

def load_existing_connections(output_file):
    """Load existing connections, filtering out SDM connections"""
    if not os.path.exists(output_file):
        return {}
    
    try:
        with open(output_file, 'r') as f:
            data = yaml.safe_load(f) or []
        
        preserved_groups = {}
        
        for item in data:
            if not isinstance(item, dict):
                continue
            
            # Handle grouped connections
            if item.get("type") == "group":
                group_name = item.get("name", "")
                connections = item.get("connections", [])
                
                # Filter out SDM connections
                non_sdm_conns = []
                for conn in connections:
                    conn_name = conn.get("name", "")
                    if conn_name and not is_sdm_connection(conn_name):
                        non_sdm_conns.append(conn)
                
                if non_sdm_conns:
                    preserved_groups[group_name] = non_sdm_conns
            
            # Handle ungrouped connections (top-level connections)
            elif item.get("type") != "group" and "addr" in item and "port" in item:
                # This is a standalone connection (not in a group)
                conn_name = item.get("name", "")
                if conn_name and not is_sdm_connection(conn_name):
                    # Put ungrouped connections into a special group
                    if 'ungrouped' not in preserved_groups:
                        preserved_groups['ungrouped'] = []
                    preserved_groups['ungrouped'].append(item)
        
        return preserved_groups
    
    except Exception as e:
        print(f"⚠️  Error reading existing config: {e}")
        return {}

def create_connection_yaml(name, addr, port, stage):
    colors = {
        'internal': '#4ECF60',
        'stage': '#FFA500',
        'production': '#FF0000'
    }
    
    color = colors.get(stage, '#808080')
    
    return f"""    - name: {name}
      last_db: 0
      network: tcp
      addr: {addr}
      port: {port}
      default_filter: '*'
      key_separator: ':'
      conn_timeout: 60
      exec_timeout: 60
      db_filter_type: none
      load_size: 10000
      mark_color: '{color}'
      refresh_interval: 5"""

def group_connections(connections):
    groups = defaultdict(list)
    
    for conn in connections:
        # Group by customer (internal/features -> 'internal', otherwise by customer)
        if conn['customer'] == 'internal' or conn['customer'].startswith('feature'):
            group_name = 'internal'
        elif conn['stage'] in ['stage', 'production']:
            group_name = conn['customer']
        else:
            group_name = 'other'
        
        # Use stage directly for color coding
        conn_yaml = create_connection_yaml(
            conn['name'],
            conn['addr'],
            conn['port'],
            conn['stage']
        )
        
        groups[group_name].append(conn_yaml)
    
    return groups

def format_connection_for_yaml(conn):
    """Format a connection dict as YAML string with proper indentation"""
    yaml_str = yaml.dump(conn, default_flow_style=False, sort_keys=False)
    # Add list marker and proper indentation
    lines = yaml_str.split('\n')
    formatted_lines = []
    for i, line in enumerate(lines):
        if line.strip():
            # First line should start with "    - "
            if i == 0:
                formatted_lines.append(f"    - {line.strip()}")
            else:
                formatted_lines.append(f"      {line}")
    return '\n'.join(formatted_lines)

def generate_yaml_content(groups, preserved_connections):
    """Generate YAML content from SDM groups and preserved connections"""
    yaml_lines = []
    
    # Prefer 'internal' first, then alphabetical by customer
    ordered_groups = []
    if 'internal' in groups:
        ordered_groups.append('internal')
    ordered_groups.extend(sorted([g for g in groups.keys() if g != 'internal']))
    
    # Add preserved groups
    for group_name in preserved_connections.keys():
        if group_name not in ordered_groups:
            ordered_groups.append(group_name)
    
    for group_name in ordered_groups:
        yaml_lines.append(f"- name: {group_name}")
        yaml_lines.append("  last_db: 0")
        yaml_lines.append("  type: group")
        yaml_lines.append("  connections:")
        
        # Add connections (preserved first, then SDM)
        if group_name in preserved_connections:
            for conn_dict in preserved_connections[group_name]:
                yaml_lines.append(format_connection_for_yaml(conn_dict))
        if group_name in groups:
            for conn_yaml in groups[group_name]:
                yaml_lines.append(conn_yaml)
        
        yaml_lines.append("")
    
    return '\n'.join(yaml_lines)

def main():
    print("Generating TinyRDM connections.yaml from sdm status --json...")
    
    output_file = os.path.expanduser('~/.config/TinyRDM/connections.yaml')
    
    # Load existing connections to preserve non-SDM connections
    preserved_connections = load_existing_connections(output_file)
    preserved_count = sum(len(conns) for conns in preserved_connections.values())
    if preserved_count > 0:
        print(f"📁 Found {preserved_count} existing non-SDM connections")
    
    sdm_data = run_sdm_status()
    if not sdm_data:
        return
    
    connections = parse_redis_connections(sdm_data)
    print(f"🔍 Found {len(connections)} Redis connections from SDM")
    
    groups = group_connections(connections)
    for group_name, conns in groups.items():
        print(f"  {group_name}: {len(conns)} connections")
    
    yaml_content = generate_yaml_content(groups, preserved_connections)
    
    try:
        with open(output_file, 'w') as f:
            f.write(yaml_content)
        print(f"✅ Successfully generated {output_file}")
        if preserved_count > 0:
            print(f"💾 Preserved {preserved_count} non-SDM connections")
    except Exception as e:
        print(f"❌ Error writing file: {e}")

if __name__ == '__main__':
    main()