starrocks-br 0.1.0__py3-none-any.whl

starrocks_br/__init__.py

@@ -0,0 +1 @@
+__all__ = ["cli", "config"]

starrocks_br/cli.py

@@ -0,0 +1,385 @@
+import click
+import os
+import sys
+from datetime import datetime
+from . import config as config_module
+from . import db
+from . import health
+from . import repository
+from . import concurrency
+from . import planner
+from . import labels
+from . import executor
+from . import restore
+from . import schema
+from . import logger
+
+
+@click.group()
+def cli():
+    """StarRocks Backup & Restore automation tool."""
+    pass
+
+
+@cli.command('init')
+@click.option('--config', required=True, help='Path to config YAML file')
+def init(config):
+    """Initialize ops database and control tables.
+    
+    Creates the ops database with required tables:
+    - ops.table_inventory: Inventory groups mapping to databases/tables
+    - ops.backup_history: Backup operation history
+    - ops.restore_history: Restore operation history
+    - ops.run_status: Job concurrency control
+    
+    Run this once before using backup/restore commands.
+    """
+    try:
+        cfg = config_module.load_config(config)
+        config_module.validate_config(cfg)
+        
+        database = db.StarRocksDB(
+            host=cfg['host'],
+            port=cfg['port'],
+            user=cfg['user'],
+            password=os.getenv('STARROCKS_PASSWORD'),
+            database=cfg['database']
+        )
+        
+        with database:
+            logger.info("Initializing ops schema...")
+            schema.initialize_ops_schema(database)
+            logger.info("")
+            logger.info("Next steps:")
+            logger.info("1. Insert your table inventory records:")
+            logger.info("   INSERT INTO ops.table_inventory")
+            logger.info("   (inventory_group, database_name, table_name)")
+            logger.info("   VALUES ('my_daily_incremental', 'your_db', 'your_fact_table');")
+            logger.info("   VALUES ('my_full_database_backup', 'your_db', '*');")
+            logger.info("   VALUES ('my_full_dimension_tables', 'your_db', 'dim_customers');")
+            logger.info("   VALUES ('my_full_dimension_tables', 'your_db', 'dim_products');")
+            logger.info("")
+            logger.info("2. Run your first backup:")
+            logger.info("   starrocks-br backup incremental --group my_daily_incremental --config config.yaml")
+            
+    except FileNotFoundError as e:
+        logger.error(f"Config file not found: {e}")
+        sys.exit(1)
+    except ValueError as e:
+        logger.error(f"Configuration error: {e}")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Failed to initialize schema: {e}")
+        sys.exit(1)
+
+
+@cli.group()
+def backup():
+    """Backup commands."""
+    pass
+
+
+@backup.command('incremental')
+@click.option('--config', required=True, help='Path to config YAML file')
+@click.option('--baseline-backup', help='Specific backup label to use as baseline (optional). If not provided, uses the latest successful full backup.')
+@click.option('--group', required=True, help='Inventory group to backup from table_inventory. Supports wildcard \'*\'.')
+@click.option('--name', help='Optional logical name (label) for the backup. Supports -v#r placeholder for auto-versioning.')
+def backup_incremental(config, baseline_backup, group, name):
+    """Run incremental backup of partitions changed since the latest full backup.
+    
+    By default, uses the latest successful full backup as baseline.
+    Optionally specify a specific backup label to use as baseline.
+    
+    Flow: load config → check health → ensure repository → reserve job slot →
+    find baseline backup → find recent partitions → generate label → build backup command → execute backup
+    """
+    try:
+        cfg = config_module.load_config(config)
+        config_module.validate_config(cfg)
+        
+        database = db.StarRocksDB(
+            host=cfg['host'],
+            port=cfg['port'],
+            user=cfg['user'],
+            password=os.getenv('STARROCKS_PASSWORD'),
+            database=cfg['database']
+        )
+        
+        with database:
+            was_created = schema.ensure_ops_schema(database)
+            if was_created:
+                logger.warning("ops schema was auto-created. Please run 'starrocks-br init' after populating config.")
+                logger.warning("Remember to populate ops.table_inventory with your backup groups!")
+                sys.exit(1) # Exit if schema was just created, requires user action
+            
+            healthy, message = health.check_cluster_health(database)
+            if not healthy:
+                logger.error(f"Cluster health check failed: {message}")
+                sys.exit(1)
+            
+            logger.success(f"Cluster health: {message}")
+            
+            repository.ensure_repository(database, cfg['repository'])
+            
+            logger.success(f"Repository '{cfg['repository']}' verified")
+            
+            label = labels.determine_backup_label(
+                db=database,
+                backup_type='incremental',
+                database_name=cfg['database'],
+                custom_name=name
+            )
+            
+            logger.success(f"Generated label: {label}")
+            
+            if baseline_backup:
+                logger.success(f"Using specified baseline backup: {baseline_backup}")
+            else:
+                latest_backup = planner.find_latest_full_backup(database, cfg['database'])
+                if latest_backup:
+                    logger.success(f"Using latest full backup as baseline: {latest_backup['label']} ({latest_backup['backup_type']})")
+                else:
+                    logger.warning("No full backup found - this will be the first incremental backup")
+            
+            partitions = planner.find_recent_partitions(
+                database, cfg['database'], baseline_backup_label=baseline_backup, group_name=group
+            )
+            
+            if not partitions:
+                logger.warning("No partitions found to backup")
+                sys.exit(1)
+            
+            logger.success(f"Found {len(partitions)} partition(s) to backup")
+            
+            backup_command = planner.build_incremental_backup_command(
+                partitions, cfg['repository'], label, cfg['database']
+            )
+            
+            planner.record_backup_partitions(database, label, partitions)
+            
+            concurrency.reserve_job_slot(database, scope='backup', label=label)
+            
+            logger.success(f"Job slot reserved")
+            logger.info(f"Starting incremental backup for group '{group}'...")
+            result = executor.execute_backup(
+                database,
+                backup_command,
+                repository=cfg['repository'],
+                backup_type='incremental',
+                scope='backup',
+                database=cfg['database']
+            )
+            
+            if result['success']:
+                logger.success(f"Backup completed successfully: {result['final_status']['state']}")
+                sys.exit(0)
+            else:
+                state = result.get('final_status', {}).get('state', 'UNKNOWN')
+                if state == "LOST":
+                    logger.critical("Backup tracking lost!")
+                    logger.warning("Another backup operation started during ours.")
+                    logger.tip("Enable ops.run_status concurrency checks to prevent this.")
+                logger.error(f"{result['error_message']}")
+                sys.exit(1)
+                
+    except FileNotFoundError as e:
+        logger.error(f"Config file not found: {e}")
+        sys.exit(1)
+    except ValueError as e:
+        logger.error(f"Configuration error: {e}")
+        sys.exit(1)
+    except RuntimeError as e:
+        logger.error(f"{e}")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        sys.exit(1)
+
+
+@backup.command('full')
+@click.option('--config', required=True, help='Path to config YAML file')
+@click.option('--group', required=True, help='Inventory group to backup from table_inventory. Supports wildcard \'*\'.')
+@click.option('--name', help='Optional logical name (label) for the backup. Supports -v#r placeholder for auto-versioning.')
+def backup_full(config, group, name):
+    """Run a full backup for a specified inventory group.
+    
+    Flow: load config → check health → ensure repository → reserve job slot →
+    find tables by group → generate label → build backup command → execute backup
+    """
+    try:
+        cfg = config_module.load_config(config)
+        config_module.validate_config(cfg)
+        
+        database = db.StarRocksDB(
+            host=cfg['host'],
+            port=cfg['port'],
+            user=cfg['user'],
+            password=os.getenv('STARROCKS_PASSWORD'),
+            database=cfg['database']
+        )
+        
+        with database:
+            was_created = schema.ensure_ops_schema(database)
+            if was_created:
+                logger.warning("ops schema was auto-created. Please run 'starrocks-br init' after populating config.")
+                logger.warning("Remember to populate ops.table_inventory with your backup groups!")
+                sys.exit(1) # Exit if schema was just created, requires user action
+            
+            healthy, message = health.check_cluster_health(database)
+            if not healthy:
+                logger.error(f"Cluster health check failed: {message}")
+                sys.exit(1)
+            
+            logger.success(f"Cluster health: {message}")
+            
+            repository.ensure_repository(database, cfg['repository'])
+            
+            logger.success(f"Repository '{cfg['repository']}' verified")
+            
+            label = labels.determine_backup_label(
+                db=database,
+                backup_type='full',
+                database_name=cfg['database'],
+                custom_name=name
+            )
+            
+            logger.success(f"Generated label: {label}")
+            
+            backup_command = planner.build_full_backup_command(
+                database, group, cfg['repository'], label, cfg['database']
+            )
+            
+            if not backup_command:
+                logger.warning(f"No tables found in group '{group}' for database '{cfg['database']}' to backup")
+                sys.exit(1)
+            
+            tables = planner.find_tables_by_group(database, group)
+            all_partitions = planner.get_all_partitions_for_tables(database, cfg['database'], tables)
+            planner.record_backup_partitions(database, label, all_partitions)
+            
+            concurrency.reserve_job_slot(database, scope='backup', label=label)
+            
+            logger.success(f"Job slot reserved")
+            logger.info(f"Starting full backup for group '{group}'...")
+            result = executor.execute_backup(
+                database,
+                backup_command,
+                repository=cfg['repository'],
+                backup_type='full',
+                scope='backup',
+                database=cfg['database']
+            )
+            
+            if result['success']:
+                logger.success(f"Backup completed successfully: {result['final_status']['state']}")
+                sys.exit(0)
+            else:
+                state = result.get('final_status', {}).get('state', 'UNKNOWN')
+                if state == "LOST":
+                    logger.critical("Backup tracking lost!")
+                    logger.warning("Another backup operation started during ours.")
+                    logger.tip("Enable ops.run_status concurrency checks to prevent this.")
+                logger.error(f"{result['error_message']}")
+                sys.exit(1)
+                
+    except (FileNotFoundError, ValueError, RuntimeError, Exception) as e:
+        if isinstance(e, FileNotFoundError):
+            logger.error(f"Config file not found: {e}")
+        elif isinstance(e, ValueError):
+            logger.error(f"Configuration error: {e}")
+        elif isinstance(e, RuntimeError):
+            logger.error(f"{e}")
+        else:
+            logger.error(f"Unexpected error: {e}")
+        sys.exit(1)
+
+
+
+
+@cli.command('restore')
+@click.option('--config', required=True, help='Path to config YAML file')
+@click.option('--target-label', required=True, help='Backup label to restore to')
+@click.option('--group', help='Optional inventory group to filter tables to restore')
+@click.option('--rename-suffix', default='_restored', help='Suffix for temporary tables during restore (default: _restored)')
+def restore_command(config, target_label, group, rename_suffix):
+    """Restore data to a specific point in time using intelligent backup chain resolution.
+    
+    This command automatically determines the correct sequence of backups needed for restore:
+    - For full backups: restores directly from the target backup
+    - For incremental backups: restores the base full backup first, then applies the incremental
+    
+    The restore process uses temporary tables with the specified suffix for safety, then performs
+    an atomic rename to make the restored data live.
+    
+    Flow: load config → find restore pair → get tables from backup → execute restore flow
+    """
+    try:
+        cfg = config_module.load_config(config)
+        config_module.validate_config(cfg)
+        
+        database = db.StarRocksDB(
+            host=cfg['host'],
+            port=cfg['port'],
+            user=cfg['user'],
+            password=os.getenv('STARROCKS_PASSWORD'),
+            database=cfg['database']
+        )
+        
+        with database:
+            was_created = schema.ensure_ops_schema(database)
+            if was_created:
+                logger.warning("ops schema was auto-created. Please run 'starrocks-br init' after populating config.")
+                logger.warning("Remember to populate ops.table_inventory with your backup groups!")
+                sys.exit(1) # Exit if schema was just created, requires user action
+            
+            logger.info(f"Finding restore sequence for target backup: {target_label}")
+            
+            try:
+                restore_pair = restore.find_restore_pair(database, target_label)
+                logger.success(f"Found restore sequence: {' -> '.join(restore_pair)}")
+            except ValueError as e:
+                logger.error(f"Failed to find restore sequence: {e}")
+                sys.exit(1)
+            
+            logger.info("Determining tables to restore from backup manifest...")
+            tables_to_restore = restore.get_tables_from_backup(database, target_label, group)
+            
+            if not tables_to_restore:
+                if group:
+                    logger.warning(f"No tables found in backup '{target_label}' for group '{group}'")
+                else:
+                    logger.warning(f"No tables found in backup '{target_label}'")
+                sys.exit(1)
+            
+            logger.success(f"Found {len(tables_to_restore)} table(s) to restore: {', '.join(tables_to_restore)}")
+            
+            logger.info("Starting restore flow...")
+            result = restore.execute_restore_flow(
+                database,
+                cfg['repository'],
+                restore_pair,
+                tables_to_restore,
+                rename_suffix
+            )
+            
+            if result['success']:
+                logger.success(result['message'])
+                sys.exit(0)
+            else:
+                logger.error(f"Restore failed: {result['error_message']}")
+                sys.exit(1)
+                
+    except FileNotFoundError as e:
+        logger.error(f"Config file not found: {e}")
+        sys.exit(1)
+    except ValueError as e:
+        logger.error(f"Configuration error: {e}")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    cli()
+

starrocks_br/concurrency.py

@@ -0,0 +1,177 @@
+from typing import Literal, List, Tuple
+from . import logger
+
+
+def reserve_job_slot(db, scope: str, label: str) -> None:
+    """Reserve a job slot in ops.run_status to prevent overlapping jobs.
+
+    We consider any row with state='ACTIVE' for the same scope as a conflict.
+    However, we implement self-healing logic to automatically clean up stale locks.
+    """
+    active_jobs = _get_active_jobs_for_scope(db, scope)
+    
+    if not active_jobs:
+        _insert_new_job(db, scope, label)
+        return
+    
+    _handle_active_job_conflicts(db, scope, active_jobs)
+    
+    _insert_new_job(db, scope, label)
+
+
+def _get_active_jobs_for_scope(db, scope: str) -> List[Tuple[str, str, str]]:
+    """Get all active jobs for the given scope."""
+    rows = db.query("SELECT scope, label, state FROM ops.run_status WHERE state = 'ACTIVE'")
+    return [row for row in rows if row[0] == scope]
+
+
+def _handle_active_job_conflicts(db, scope: str, active_jobs: List[Tuple[str, str, str]]) -> None:
+    """Handle conflicts with active jobs, cleaning up stale ones where possible."""
+    for active_scope, active_label, _ in active_jobs:
+        if _can_heal_stale_job(active_scope, active_label, db):
+            _cleanup_stale_job(db, active_scope, active_label)
+            logger.success(f"Cleaned up stale backup job: {active_label}")
+        else:
+            _raise_concurrency_conflict(scope, active_jobs)
+
+
+def _can_heal_stale_job(scope: str, label: str, db) -> bool:
+    """Check if a stale job can be healed (only for backup jobs)."""
+    if scope != 'backup':
+        return False
+    
+    return _is_backup_job_stale(db, label)
+
+
+def _raise_concurrency_conflict(scope: str, active_jobs: List[Tuple[str, str, str]]) -> None:
+    """Raise a concurrency conflict error with helpful message."""
+    active_job_strings = [f"{job[0]}:{job[1]}" for job in active_jobs]
+    active_labels = [job[1] for job in active_jobs]
+    
+    raise RuntimeError(
+        f"Concurrency conflict: Another '{scope}' job is already ACTIVE: {', '.join(active_job_strings)}. "
+        f"Wait for it to complete or cancel it via: UPDATE ops.run_status SET state='CANCELLED' "
+        f"WHERE label='{active_labels[0]}' AND state='ACTIVE'"
+    )
+
+
+def _insert_new_job(db, scope: str, label: str) -> None:
+    """Insert a new active job record."""
+    sql = (
+        "INSERT INTO ops.run_status (scope, label, state, started_at) "
+        "VALUES ('%s','%s','ACTIVE', NOW())" % (scope, label)
+    )
+    db.execute(sql)
+
+
+def _is_backup_job_stale(db, label: str) -> bool:
+    """Check if a backup job is stale by querying StarRocks SHOW BACKUP.
+    
+    Returns True if the job is stale (not actually running), False if it's still active.
+    """
+    try:
+        user_databases = _get_user_databases(db)
+        
+        for database_name in user_databases:
+            job_status = _check_backup_job_in_database(db, database_name, label)
+            
+            if job_status is None:
+                continue
+                
+            if job_status == "active":
+                return False
+            elif job_status == "stale":
+                return True
+                
+        return True
+        
+    except Exception as e:
+        logger.error(f"Error checking backup job status: {e}")
+        return False
+
+
+def _get_user_databases(db) -> List[str]:
+    """Get list of user databases (excluding system databases)."""
+    SYSTEM_DATABASES = {'information_schema', 'mysql', 'sys', 'ops'}
+    
+    databases = db.query("SHOW DATABASES")
+    return [
+        _extract_database_name(db_row) 
+        for db_row in databases 
+        if _extract_database_name(db_row) not in SYSTEM_DATABASES
+    ]
+
+
+def _extract_database_name(db_row) -> str:
+    """Extract database name from database query result."""
+    if isinstance(db_row, (list, tuple)):
+        return db_row[0]
+    return db_row.get('Database', '')
+
+
+def _check_backup_job_in_database(db, database_name: str, label: str) -> str:
+    """Check if backup job exists in specific database and return its status.
+    
+    Returns:
+        'active' if job is still running
+        'stale' if job is in terminal state  
+        None if job not found in this database
+    """
+    try:
+        show_backup_query = f"SHOW BACKUP FROM {database_name}"
+        backup_rows = db.query(show_backup_query)
+        
+        if not backup_rows:
+            return None
+            
+        result = backup_rows[0]
+        snapshot_name, state = _extract_backup_info(result)
+        
+        if snapshot_name != label:
+            return None
+            
+        if state in ["FINISHED", "CANCELLED", "FAILED"]:
+            return "stale"
+        else:
+            return "active"
+            
+    except Exception:
+        return None
+
+
+def _extract_backup_info(result) -> Tuple[str, str]:
+    """Extract snapshot name and state from SHOW BACKUP result."""
+    if isinstance(result, dict):
+        snapshot_name = result.get("SnapshotName", "")
+        state = result.get("State", "UNKNOWN")
+    else:
+        snapshot_name = result[1] if len(result) > 1 else ""
+        state = result[3] if len(result) > 3 else "UNKNOWN"
+    
+    return snapshot_name, state
+
+
+def _cleanup_stale_job(db, scope: str, label: str) -> None:
+    """Clean up a stale job by updating its state to CANCELLED."""
+    sql = (
+        "UPDATE ops.run_status SET state='CANCELLED', finished_at=NOW() "
+        "WHERE scope='%s' AND label='%s' AND state='ACTIVE'" % (scope, label)
+    )
+    db.execute(sql)
+
+
+def complete_job_slot(
+    db, 
+    scope: str, 
+    label: str, 
+    final_state: Literal['FINISHED', 'FAILED', 'CANCELLED']
+) -> None:
+    """Complete job slot and persist final state.
+
+    Simple approach: update the same row by scope/label.
+    """
+    sql = (
+        "UPDATE ops.run_status SET state='%s', finished_at=NOW() WHERE scope='%s' AND label='%s'"
+        % (final_state, scope, label)
+    )
+    db.execute(sql)

starrocks_br/config.py

@@ -0,0 +1,41 @@
+import yaml
+from typing import Dict
+
+
+def load_config(config_path: str) -> Dict:
+    """Load and parse YAML configuration file.
+    
+    Args:
+        config_path: Path to the YAML config file
+        
+    Returns:
+        Dictionary containing configuration
+        
+    Raises:
+        FileNotFoundError: If config file doesn't exist
+        yaml.YAMLError: If config file is not valid YAML
+    """
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    if not isinstance(config, dict):
+        raise ValueError("Config must be a dictionary")
+    
+    return config
+
+
+def validate_config(config: Dict) -> None:
+    """Validate that config contains required fields.
+    
+    Args:
+        config: Configuration dictionary
+        
+    Raises:
+        ValueError: If required fields are missing
+    """
+    required_fields = ['host', 'port', 'user', 'database', 'repository']
+    
+    for field in required_fields:
+        if field not in config:
+            raise ValueError(f"Missing required config field: {field}")
+

starrocks_br/db.py

@@ -0,0 +1,88 @@
+import mysql.connector
+from typing import List
+
+
+class StarRocksDB:
+    """Database connection wrapper for StarRocks."""
+    
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        """Initialize database connection.
+        
+        Args:
+            host: Database host
+            port: Database port
+            user: Database user
+            password: Database password
+            database: Default database name
+        """
+        self.host = host
+        self.port = port
+        self.user = user
+        self.password = password
+        self.database = database
+        self._connection = None
+    
+    def connect(self) -> None:
+        """Establish database connection."""
+        self._connection = mysql.connector.connect(
+            host=self.host,
+            port=self.port,
+            user=self.user,
+            password=self.password,
+            database=self.database
+        )
+    
+    def close(self) -> None:
+        """Close database connection."""
+        if self._connection:
+            self._connection.close()
+            self._connection = None
+    
+    def execute(self, sql: str) -> None:
+        """Execute a SQL statement that doesn't return results.
+        
+        Args:
+            sql: SQL statement to execute
+        """
+        if not self._connection:
+            self.connect()
+        
+        cursor = self._connection.cursor()
+        try:
+            cursor.execute(sql)
+            self._connection.commit()
+        finally:
+            cursor.close()
+    
+    def query(self, sql: str, params: tuple = None) -> List[tuple]:
+        """Execute a SQL query and return results.
+        
+        Args:
+            sql: SQL query to execute
+            params: Optional tuple of parameters for parameterized queries
+            
+        Returns:
+            List of tuples containing query results
+        """
+        if not self._connection:
+            self.connect()
+        
+        cursor = self._connection.cursor()
+        try:
+            if params:
+                cursor.execute(sql, params)
+            else:
+                cursor.execute(sql)
+            return cursor.fetchall()
+        finally:
+            cursor.close()
+    
+    def __enter__(self):
+        """Context manager entry."""
+        self.connect()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()
+