starrocks-br 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

starrocks_br/cli.py

@@ -15,6 +15,40 @@ from . import schema
 from . import logger
 
 
+def _handle_snapshot_exists_error(error_details: dict, label: str, config: str, repository: str, backup_type: str, group: str, baseline_backup: str = None) -> None:
+    """Handle snapshot_exists error by providing helpful guidance to the user.
+    
+    Args:
+        error_details: Error details dict containing error_type and snapshot_name
+        label: The backup label that was generated
+        config: Path to config file
+        repository: Repository name
+        backup_type: Type of backup ('incremental' or 'full')
+        group: Inventory group name
+        baseline_backup: Optional baseline backup label (for incremental backups)
+    """
+    snapshot_name = error_details.get('snapshot_name', label)
+    logger.error(f"Snapshot '{snapshot_name}' already exists in the repository.")
+    logger.info("")
+    logger.info("This typically happens when:")
+    logger.info("  • The CLI lost connectivity during a previous backup operation")
+    logger.info("  • The backup completed on the server, but backup_history wasn't updated")
+    logger.info("")
+    logger.info("To resolve this, retry the backup with a custom label using --name:")
+    
+    if backup_type == 'incremental':
+        retry_cmd = f"  starrocks-br backup incremental --config {config} --group {group} --name {snapshot_name}_retry"
+        if baseline_backup:
+            retry_cmd += f" --baseline-backup {baseline_backup}"
+        logger.info(retry_cmd)
+    else:
+        logger.info(f"  starrocks-br backup full --config {config} --group {group} --name {snapshot_name}_retry")
+    
+    logger.info("")
+    logger.tip("You can verify the existing backup by checking the repository or running:")
+    logger.tip(f"  SHOW SNAPSHOT ON {repository} WHERE Snapshot = '{snapshot_name}'")
+
+
 @click.group()
 def cli():
     """StarRocks Backup & Restore automation tool."""
@@ -43,7 +77,8 @@ def init(config):
             port=cfg['port'],
             user=cfg['user'],
             password=os.getenv('STARROCKS_PASSWORD'),
-            database=cfg['database']
+            database=cfg['database'],
+            tls_config=cfg.get('tls'),
         )
         
         with database:
@@ -102,7 +137,8 @@ def backup_incremental(config, baseline_backup, group, name):
             port=cfg['port'],
             user=cfg['user'],
             password=os.getenv('STARROCKS_PASSWORD'),
-            database=cfg['database']
+            database=cfg['database'],
+            tls_config=cfg.get('tls'),
         )
         
         with database:
@@ -174,6 +210,13 @@ def backup_incremental(config, baseline_backup, group, name):
                 logger.success(f"Backup completed successfully: {result['final_status']['state']}")
                 sys.exit(0)
             else:
+                error_details = result.get('error_details')
+                if error_details and error_details.get('error_type') == 'snapshot_exists':
+                    _handle_snapshot_exists_error(
+                        error_details, label, config, cfg['repository'], 'incremental', group, baseline_backup
+                    )
+                    sys.exit(1)
+                
                 state = result.get('final_status', {}).get('state', 'UNKNOWN')
                 if state == "LOST":
                     logger.critical("Backup tracking lost!")
@@ -215,7 +258,8 @@ def backup_full(config, group, name):
             port=cfg['port'],
             user=cfg['user'],
             password=os.getenv('STARROCKS_PASSWORD'),
-            database=cfg['database']
+            database=cfg['database'],
+            tls_config=cfg.get('tls'),
         )
         
         with database:
@@ -274,6 +318,13 @@ def backup_full(config, group, name):
                 logger.success(f"Backup completed successfully: {result['final_status']['state']}")
                 sys.exit(0)
             else:
+                error_details = result.get('error_details')
+                if error_details and error_details.get('error_type') == 'snapshot_exists':
+                    _handle_snapshot_exists_error(
+                        error_details, label, config, cfg['repository'], 'full', group
+                    )
+                    sys.exit(1)
+                
                 state = result.get('final_status', {}).get('state', 'UNKNOWN')
                 if state == "LOST":
                     logger.critical("Backup tracking lost!")
@@ -300,8 +351,10 @@ def backup_full(config, group, name):
 @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('--table', help='Optional table name to restore (table name only, database comes from config). Cannot be used with --group.')
 @click.option('--rename-suffix', default='_restored', help='Suffix for temporary tables during restore (default: _restored)')
-def restore_command(config, target_label, group, rename_suffix):
+@click.option('--yes', is_flag=True, help='Skip confirmation prompt and proceed automatically')
+def restore_command(config, target_label, group, table, rename_suffix, yes):
     """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:
@@ -311,9 +364,23 @@ def restore_command(config, target_label, group, rename_suffix):
     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
+    Flow: load config → check health → ensure repository → find restore pair → get tables from backup → execute restore flow
     """
     try:
+        if group and table:
+            logger.error("Cannot specify both --group and --table. Use --table for single table restore or --group for inventory group restore.")
+            sys.exit(1)
+        
+        if table:
+            table = table.strip()
+            if not table:
+                logger.error("Table name cannot be empty")
+                sys.exit(1)
+            
+            if '.' in table:
+                logger.error("Table name must not include database prefix. Use 'table_name' not 'database.table_name'. Database comes from config file.")
+                sys.exit(1)
+        
         cfg = config_module.load_config(config)
         config_module.validate_config(cfg)
         
@@ -322,7 +389,8 @@ def restore_command(config, target_label, group, rename_suffix):
             port=cfg['port'],
             user=cfg['user'],
             password=os.getenv('STARROCKS_PASSWORD'),
-            database=cfg['database']
+            database=cfg['database'],
+            tls_config=cfg.get('tls'),
         )
         
         with database:
@@ -332,6 +400,17 @@ def restore_command(config, target_label, group, rename_suffix):
                 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")
+            
             logger.info(f"Finding restore sequence for target backup: {target_label}")
             
             try:
@@ -342,11 +421,24 @@ def restore_command(config, target_label, group, rename_suffix):
                 sys.exit(1)
             
             logger.info("Determining tables to restore from backup manifest...")
-            tables_to_restore = restore.get_tables_from_backup(database, target_label, group)
+            
+            try:
+                tables_to_restore = restore.get_tables_from_backup(
+                    database, 
+                    target_label, 
+                    group=group, 
+                    table=table, 
+                    database=cfg['database'] if table else None
+                )
+            except ValueError as e:
+                logger.error(str(e))
+                sys.exit(1)
             
             if not tables_to_restore:
                 if group:
                     logger.warning(f"No tables found in backup '{target_label}' for group '{group}'")
+                elif table:
+                    logger.warning(f"No tables found in backup '{target_label}' for table '{table}'")
                 else:
                     logger.warning(f"No tables found in backup '{target_label}'")
                 sys.exit(1)
@@ -359,7 +451,8 @@ def restore_command(config, target_label, group, rename_suffix):
                 cfg['repository'],
                 restore_pair,
                 tables_to_restore,
-                rename_suffix
+                rename_suffix,
+                skip_confirmation=yes
             )
             
             if result['success']:
@@ -375,6 +468,9 @@ def restore_command(config, target_label, group, rename_suffix):
     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)

starrocks_br/config.py

@@ -1,8 +1,8 @@
 import yaml
-from typing import Dict
+from typing import Any, Dict, Optional
 
 
-def load_config(config_path: str) -> Dict:
+def load_config(config_path: str) -> Dict[str, Any]:
     """Load and parse YAML configuration file.
     
     Args:
@@ -24,7 +24,7 @@ def load_config(config_path: str) -> Dict:
     return config
 
 
-def validate_config(config: Dict) -> None:
+def validate_config(config: Dict[str, Any]) -> None:
     """Validate that config contains required fields.
     
     Args:
@@ -39,3 +39,26 @@ def validate_config(config: Dict) -> None:
         if field not in config:
             raise ValueError(f"Missing required config field: {field}")
 
+    _validate_tls_section(config.get('tls'))
+
+
+def _validate_tls_section(tls_config) -> None:
+    if tls_config is None:
+        return
+
+    if not isinstance(tls_config, dict):
+        raise ValueError("TLS configuration must be a dictionary")
+
+    enabled = bool(tls_config.get('enabled', False))
+
+    if enabled and not tls_config.get('ca_cert'):
+        raise ValueError("TLS configuration requires 'ca_cert' when 'enabled' is true")
+
+    if 'verify_server_cert' in tls_config and not isinstance(tls_config['verify_server_cert'], bool):
+        raise ValueError("TLS configuration field 'verify_server_cert' must be a boolean if provided")
+
+    if 'tls_versions' in tls_config:
+        tls_versions = tls_config['tls_versions']
+        if not isinstance(tls_versions, list) or not all(isinstance(version, str) for version in tls_versions):
+            raise ValueError("TLS configuration field 'tls_versions' must be a list of strings if provided")
+

starrocks_br/db.py

@@ -1,11 +1,19 @@
 import mysql.connector
-from typing import List
+from typing import Any, Dict, List, Optional
 
 
 class StarRocksDB:
     """Database connection wrapper for StarRocks."""
     
-    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        user: str,
+        password: str,
+        database: str,
+        tls_config: Optional[Dict[str, Any]] = None,
+    ):
         """Initialize database connection.
         
         Args:
@@ -21,16 +29,44 @@ class StarRocksDB:
         self.password = password
         self.database = database
         self._connection = None
+        self.tls_config = tls_config or {}
+        self._timezone: Optional[str] = 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
-        )
+        conn_args: Dict[str, Any] = {
+            'host': self.host,
+            'port': self.port,
+            'user': self.user,
+            'password': self.password,
+            'database': self.database,
+        }
+
+        if self.tls_config.get('enabled'):
+            ssl_args: Dict[str, Any] = {
+                'ssl_ca': self.tls_config.get('ca_cert'),
+                'ssl_cert': self.tls_config.get('client_cert'),
+                'ssl_key': self.tls_config.get('client_key'),
+                'ssl_verify_cert': self.tls_config.get('verify_server_cert', True),
+            }
+
+            tls_versions = self.tls_config.get('tls_versions', ['TLSv1.2', 'TLSv1.3'])
+            if tls_versions:
+                ssl_args['tls_versions'] = tls_versions
+
+            conn_args.update({key: value for key, value in ssl_args.items() if value is not None})
+
+        try:
+            self._connection = mysql.connector.connect(**conn_args)
+        except mysql.connector.Error as e:
+            if self.tls_config.get('enabled') and "SSL is required" in str(e):
+                raise mysql.connector.Error(
+                    f"TLS is enabled in configuration but StarRocks server doesn't support it. "
+                    f"Error: {e}. "
+                    f"To fix this, you need to enable TLS/SSL in your StarRocks server configuration. "
+                    f"Alternatively, set 'enabled: false' in the tls section of your config file."
+                ) from e
+            raise
     
     def close(self) -> None:
         """Close database connection."""
@@ -85,4 +121,34 @@ class StarRocksDB:
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Context manager exit."""
         self.close()
+    
+    @property
+    def timezone(self) -> str:
+        """Get the StarRocks cluster timezone.
+        
+        Queries the cluster timezone on first access and caches it for subsequent use.
+        If the query fails (e.g., database unavailable, connection error, permissions),
+        defaults to 'UTC' to ensure the property always returns a valid timezone string.
+        
+        Returns:
+            Timezone string (e.g., 'Asia/Shanghai', 'UTC', '+08:00')
+            Defaults to 'UTC' if query fails or returns no results.
+        """
+        if self._timezone is None:
+            try:
+                query = "SHOW VARIABLES LIKE 'time_zone'"
+                rows = self.query(query)
+                
+                if not rows:
+                    self._timezone = "UTC"
+                else:
+                    row = rows[0]
+                    if isinstance(row, dict):
+                        self._timezone = row.get("Value", "UTC")
+                    else:
+                        self._timezone = row[1] if len(row) > 1 else "UTC"
+            except Exception:
+                self._timezone = "UTC"
+        
+        return self._timezone
 

starrocks_br/executor.py

@@ -1,23 +1,68 @@
 import time
 import datetime
-from typing import Dict, Literal, Optional
-from . import history, concurrency, logger
+import re
+from typing import Dict, Literal, Optional, Tuple
+from . import history, concurrency, logger, timezone
 
-MAX_POLLS = 21600 # 6 hours
+MAX_POLLS = 86400 # 1 day
 
-def submit_backup_command(db, backup_command: str) -> tuple[bool, Optional[str]]:
+def submit_backup_command(db, backup_command: str) -> Tuple[bool, Optional[str], Optional[Dict[str, str]]]:
     """Submit a backup command to StarRocks.
     
-    Returns (success, error_message).
+    Returns (success, error_message, error_details).
+    error_details is a dict with keys like 'error_type' and 'snapshot_name' for specific error cases.
     """
     try:
         db.execute(backup_command.strip())
-        return True, None
+        return True, None, None
     except Exception as e:
-        error_msg = f"Failed to submit backup command: {type(e).__name__}: {str(e)}"
+        error_str = str(e)
+        error_type = type(e).__name__
+        
+        snapshot_exists_match = _check_snapshot_exists_error(e, error_str)
+        if snapshot_exists_match:
+            snapshot_name = snapshot_exists_match
+            error_details = {
+                'error_type': 'snapshot_exists',
+                'snapshot_name': snapshot_name
+            }
+            error_msg = f"Snapshot '{snapshot_name}' already exists in repository"
+            logger.error(error_msg)
+            logger.error(f"backup_command: {backup_command}")
+            return False, error_msg, error_details
+        
+        error_msg = f"Failed to submit backup command: {error_type}: {error_str}"
         logger.error(error_msg)
         logger.error(f"backup_command: {backup_command}")
-        return False, error_msg
+        return False, error_msg, None
+
+
+def _check_snapshot_exists_error(exception: Exception, error_str: str) -> Optional[str]:
+    """Check if the error is a 'snapshot already exists' error and extract snapshot name.
+    
+    Args:
+        exception: The exception that was raised
+        error_str: String representation of the error
+        
+    Returns:
+        Snapshot name if this is a snapshot exists error, None otherwise
+    """
+    snapshot_name_pattern = r"Snapshot with name '([^']+)' already exist"
+    error_lower = error_str.lower()
+    
+    is_snapshot_exists_error = (
+        "already exist" in error_lower or 
+        "already exists" in error_lower or
+        ("5064" in error_str and "already exist" in error_lower) or
+        (hasattr(exception, 'errno') and exception.errno == 5064)
+    )
+    
+    if is_snapshot_exists_error:
+        match = re.search(snapshot_name_pattern, error_str, re.IGNORECASE)
+        if match:
+            return match.group(1)
+    
+    return None
 
 
 def poll_backup_status(db, label: str, database: str, max_polls: int = MAX_POLLS, poll_interval: float = 1.0) -> Dict[str, str]:
@@ -117,15 +162,19 @@ def execute_backup(
     if not database:
         database = _extract_database_from_command(backup_command)
 
-    started_at = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    cluster_tz = db.timezone
+    started_at = timezone.get_current_time_in_cluster_tz(cluster_tz)
     
-    success, submit_error = submit_backup_command(db, backup_command)
+    success, submit_error, error_details = submit_backup_command(db, backup_command)
     if not success:
-        return {
+        result = {
             "success": False,
             "final_status": None,
             "error_message": submit_error or "Failed to submit backup command (unknown error)"
         }
+        if error_details:
+            result["error_details"] = error_details
+        return result
     
     try:
         final_status = poll_backup_status(db, label, database, max_polls, poll_interval)
@@ -133,6 +182,7 @@ def execute_backup(
         success = final_status["state"] == "FINISHED"
 
         try:
+            finished_at = timezone.get_current_time_in_cluster_tz(cluster_tz)
             history.log_backup(
                 db,
                 {
@@ -141,7 +191,7 @@ def execute_backup(
                     "status": final_status["state"],
                     "repository": repository,
                     "started_at": started_at,
-                    "finished_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                    "finished_at": finished_at,
                     "error_message": None if success else (final_status["state"] or ""),
                 },
             )

starrocks_br/planner.py

@@ -1,6 +1,8 @@
 from typing import List, Dict, Optional
+import datetime
+import hashlib
 
-from starrocks_br import logger
+from starrocks_br import logger, timezone
 
 
 def find_latest_full_backup(db, database: str) -> Optional[Dict[str, str]]:
@@ -11,7 +13,8 @@ def find_latest_full_backup(db, database: str) -> Optional[Dict[str, str]]:
         database: Database name to search for
         
     Returns:
-        Dictionary with keys: label, backup_type, finished_at, or None if no full backup found
+        Dictionary with keys: label, backup_type, finished_at, or None if no full backup found.
+        The finished_at value is returned as a string in the cluster timezone format.
     """
     query = f"""
     SELECT label, backup_type, finished_at
@@ -29,10 +32,18 @@ def find_latest_full_backup(db, database: str) -> Optional[Dict[str, str]]:
         return None
     
     row = rows[0]
+    finished_at = row[2]
+    
+    if isinstance(finished_at, datetime.datetime):
+        cluster_tz = db.timezone
+        finished_at = finished_at.strftime("%Y-%m-%d %H:%M:%S")
+    elif not isinstance(finished_at, str):
+        finished_at = str(finished_at)
+    
     return {
         "label": row[0],
         "backup_type": row[1],
-        "finished_at": row[2]
+        "finished_at": finished_at
     }
 
 
@@ -66,6 +77,8 @@ def find_recent_partitions(db, database: str, baseline_backup_label: Optional[st
     Returns list of dictionaries with keys: database, table, partition_name.
     Only partitions of tables within the specified database are returned.
     """
+    cluster_tz = db.timezone
+    
     if baseline_backup_label:
         baseline_query = f"""
         SELECT finished_at
@@ -76,17 +89,21 @@ def find_recent_partitions(db, database: str, baseline_backup_label: Optional[st
         baseline_rows = db.query(baseline_query)
         if not baseline_rows:
             raise ValueError(f"Baseline backup '{baseline_backup_label}' not found or not successful")
-        baseline_time = baseline_rows[0][0]
+        baseline_time_raw = baseline_rows[0][0]
     else:
         latest_backup = find_latest_full_backup(db, database)
         if not latest_backup:
             raise ValueError(f"No successful full backup found for database '{database}'. Run a full database backup first.")
-        baseline_time = latest_backup['finished_at']
+        baseline_time_raw = latest_backup['finished_at']
     
-    if isinstance(baseline_time, str):
-        threshold_str = baseline_time
+    if isinstance(baseline_time_raw, datetime.datetime):
+        baseline_time_str = baseline_time_raw.strftime("%Y-%m-%d %H:%M:%S")
+    elif isinstance(baseline_time_raw, str):
+        baseline_time_str = baseline_time_raw
     else:
-        threshold_str = baseline_time.strftime("%Y-%m-%d %H:%M:%S")
+        baseline_time_str = str(baseline_time_raw)
+    
+    baseline_dt = timezone.parse_datetime_with_tz(baseline_time_str, cluster_tz)
     
     group_tables = find_tables_by_group(db, group_name)
 
@@ -129,12 +146,16 @@ def find_recent_partitions(db, database: str, baseline_backup_label: Optional[st
             partition_name = row[1]
             visible_version_time = row[3]
 
-            if isinstance(visible_version_time, str):
-                version_time_str = visible_version_time
+            if isinstance(visible_version_time, datetime.datetime):
+                visible_version_time_str = visible_version_time.strftime("%Y-%m-%d %H:%M:%S")
+            elif isinstance(visible_version_time, str):
+                visible_version_time_str = visible_version_time
             else:
-                version_time_str = visible_version_time.strftime("%Y-%m-%d %H:%M:%S")
+                visible_version_time_str = str(visible_version_time)
+            
+            visible_version_dt = timezone.parse_datetime_with_tz(visible_version_time_str, cluster_tz)
             
-            if version_time_str > threshold_str:
+            if visible_version_dt > baseline_dt:
                 recent_partitions.append({
                     'database': db_name,
                     'table': table_name,
@@ -212,7 +233,7 @@ def build_full_backup_command(db, group_name: str, repository: str, label: str,
 
 def record_backup_partitions(db, label: str, partitions: List[Dict[str, str]]) -> None:
     """Record partition metadata for a backup in ops.backup_partitions table.
-    
+
     Args:
         db: Database connection
         label: Backup label
@@ -220,12 +241,15 @@ def record_backup_partitions(db, label: str, partitions: List[Dict[str, str]]) -
     """
     if not partitions:
         return
-    
+
     for partition in partitions:
+        composite_key = f"{label}|{partition['database']}|{partition['table']}|{partition['partition_name']}"
+        key_hash = hashlib.md5(composite_key.encode('utf-8')).hexdigest()
+
         db.execute(f"""
-            INSERT INTO ops.backup_partitions 
-            (label, database_name, table_name, partition_name)
-            VALUES ('{label}', '{partition['database']}', '{partition['table']}', '{partition['partition_name']}')
+            INSERT INTO ops.backup_partitions
+            (key_hash, label, database_name, table_name, partition_name)
+            VALUES ('{key_hash}', '{label}', '{partition['database']}', '{partition['table']}', '{partition['partition_name']}')
         """)