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

starrocks_br/restore.py

@@ -3,7 +3,7 @@ import datetime
 from typing import Dict, List, Optional
 from . import history, concurrency, logger
 
-MAX_POLLS = 21600 # 6 hours
+MAX_POLLS = 86400 # 1 day
 
 def get_snapshot_timestamp(db, repo_name: str, snapshot_name: str) -> str:
     """Get the backup timestamp for a specific snapshot from the repository.
@@ -281,17 +281,30 @@ def find_restore_pair(db, target_label: str) -> List[str]:
     raise ValueError(f"Unknown backup type '{target_info['backup_type']}' for label '{target_label}'")
 
 
-def get_tables_from_backup(db, label: str, group: Optional[str] = None) -> List[str]:
+def get_tables_from_backup(db, label: str, group: Optional[str] = None, table: Optional[str] = None, database: Optional[str] = None) -> List[str]:
     """Get list of tables to restore from backup manifest.
     
     Args:
         db: Database connection
         label: Backup label
         group: Optional inventory group to filter tables
+        table: Optional table name to filter (single table, database comes from database parameter)
+        database: Database name (required if table is specified)
         
     Returns:
-        List of table names to restore
+        List of table names to restore (format: database.table)
+        
+    Raises:
+        ValueError: If both group and table are specified
+        ValueError: If table is specified but database is not provided
+        ValueError: If table is specified but not found in backup
     """
+    if group and table:
+        raise ValueError("Cannot specify both --group and --table. Use --table for single table restore or --group for inventory group restore.")
+    
+    if table and not database:
+        raise ValueError("database parameter is required when table is specified")
+    
     query = f"""
     SELECT DISTINCT database_name, table_name
     FROM ops.backup_partitions
@@ -305,6 +318,15 @@ def get_tables_from_backup(db, label: str, group: Optional[str] = None) -> List[
     
     tables = [f"{row[0]}.{row[1]}" for row in rows]
     
+    if table:
+        target_table = f"{database}.{table}"
+        filtered_tables = [t for t in tables if t == target_table]
+        
+        if not filtered_tables:
+            raise ValueError(f"Table '{table}' not found in backup '{label}' for database '{database}'")
+        
+        return filtered_tables
+    
     if group:
         group_query = f"""
         SELECT database_name, table_name
@@ -316,14 +338,26 @@ def get_tables_from_backup(db, label: str, group: Optional[str] = None) -> List[
         if not group_rows:
             return []
         
-        group_tables = {f"{row[0]}.{row[1]}" for row in group_rows}
+        group_tables = set()
+        for row in group_rows:
+            database_name, table_name = row[0], row[1]
+            if table_name == '*':
+                show_tables_query = f"SHOW TABLES FROM {database_name}"
+                try:
+                    tables_rows = db.query(show_tables_query)
+                    for table_row in tables_rows:
+                        group_tables.add(f"{database_name}.{table_row[0]}")
+                except Exception:
+                    continue
+            else:
+                group_tables.add(f"{database_name}.{table_name}")
         
         tables = [table for table in tables if table in group_tables]
     
     return tables
 
 
-def execute_restore_flow(db, repo_name: str, restore_pair: List[str], tables_to_restore: List[str], rename_suffix: str = "_restored") -> Dict:
+def execute_restore_flow(db, repo_name: str, restore_pair: List[str], tables_to_restore: List[str], rename_suffix: str = "_restored", skip_confirmation: bool = False) -> Dict:
     """Execute the complete restore flow with safety measures.
     
     Args:
@@ -332,6 +366,7 @@ def execute_restore_flow(db, repo_name: str, restore_pair: List[str], tables_to_
         restore_pair: List of backup labels in restore order
         tables_to_restore: List of tables to restore (format: database.table)
         rename_suffix: Suffix for temporary tables
+        skip_confirmation: If True, skip interactive confirmation prompt
         
     Returns:
         Dictionary with success status and details
@@ -358,12 +393,15 @@ def execute_restore_flow(db, repo_name: str, restore_pair: List[str], tables_to_
     logger.info("This will restore data to temporary tables and then perform atomic rename.")
     logger.warning("WARNING: This operation will replace existing tables!")
     
-    confirmation = input("\nDo you want to proceed? [Y/n]: ").strip()
-    if confirmation.lower() != 'y':
-        return {
-            "success": False,
-            "error_message": "Restore operation cancelled by user"
-        }
+    if not skip_confirmation:
+        confirmation = input("\nDo you want to proceed? [Y/n]: ").strip()
+        if confirmation.lower() != 'y':
+            return {
+                "success": False,
+                "error_message": "Restore operation cancelled by user"
+            }
+    else:
+        logger.info("Proceeding automatically (--yes flag provided)")
     
     try:
         database_name = tables_to_restore[0].split('.')[0]
@@ -470,6 +508,19 @@ def _build_restore_command_without_rename(backup_label: str, repo_name: str, tab
     PROPERTIES ("backup_timestamp" = "{backup_timestamp}")"""
 
 
+def _generate_timestamped_backup_name(table_name: str) -> str:
+    """Generate a timestamped backup table name.
+    
+    Args:
+        table_name: Original table name
+        
+    Returns:
+        Timestamped backup name in format: {table_name}_backup_YYYYMMDD_HHMMSS
+    """
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    return f"{table_name}_backup_{timestamp}"
+
+
 def _perform_atomic_rename(db, tables: List[str], rename_suffix: str) -> Dict:
     """Perform atomic rename of temporary tables to make them live."""
     try:
@@ -477,8 +528,9 @@ def _perform_atomic_rename(db, tables: List[str], rename_suffix: str) -> Dict:
         for table in tables:
             database, table_name = table.split('.', 1)
             temp_table_name = f"{table_name}{rename_suffix}"
+            backup_table_name = _generate_timestamped_backup_name(table_name)
             
-            rename_statements.append(f"ALTER TABLE {database}.{table_name} RENAME {table_name}_backup")
+            rename_statements.append(f"ALTER TABLE {database}.{table_name} RENAME {backup_table_name}")
             rename_statements.append(f"ALTER TABLE {database}.{temp_table_name} RENAME {table_name}")
         
         for statement in rename_statements:

starrocks_br/schema.py

@@ -71,8 +71,9 @@ def get_table_inventory_schema() -> str:
         created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
         updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
     )
-    PRIMARY KEY (inventory_group, database_name, table_name)
+    UNIQUE KEY (inventory_group, database_name, table_name)
     COMMENT "Inventory groups mapping to databases/tables (supports '*' wildcard)"
+    DISTRIBUTED BY HASH(inventory_group)
     """
 
 
@@ -120,7 +121,7 @@ def get_run_status_schema() -> str:
     CREATE TABLE IF NOT EXISTS ops.run_status (
         scope STRING NOT NULL COMMENT "Job scope: backup or restore",
         label STRING NOT NULL COMMENT "Job label or identifier",
-        state STRING NOT NULL DEFAULT "ACTIVE" COMMENT "Job state: ACTIVE or COMPLETED",
+        state STRING NOT NULL DEFAULT "ACTIVE" COMMENT "Job state: ACTIVE, FINISHED, FAILED, or CANCELLED",
         started_at DATETIME DEFAULT CURRENT_TIMESTAMP COMMENT "Job start timestamp",
         finished_at DATETIME COMMENT "Job completion timestamp"
     )
@@ -133,12 +134,14 @@ def get_backup_partitions_schema() -> str:
     """Get CREATE TABLE statement for backup_partitions."""
     return """
     CREATE TABLE IF NOT EXISTS ops.backup_partitions (
+        key_hash STRING NOT NULL COMMENT "MD5 hash of composite key (label, database_name, table_name, partition_name)",
         label STRING NOT NULL COMMENT "The backup label this partition belongs to. FK to ops.backup_history.label.",
         database_name STRING NOT NULL COMMENT "The name of the database the partition belongs to.",
         table_name STRING NOT NULL COMMENT "The name of the table the partition belongs to.",
         partition_name STRING NOT NULL COMMENT "The name of the specific partition.",
         created_at DATETIME DEFAULT CURRENT_TIMESTAMP COMMENT "Timestamp when this record was created."
     )
-    PRIMARY KEY (label, database_name, table_name, partition_name)
+    PRIMARY KEY (key_hash)
     COMMENT "Tracks every partition included in a backup snapshot."
+    DISTRIBUTED BY HASH(key_hash)
     """

starrocks_br/timezone.py

@@ -0,0 +1,125 @@
+import datetime
+from typing import Union
+from zoneinfo import ZoneInfo
+
+
+def get_current_time_in_cluster_tz(cluster_tz: str) -> str:
+    """Get current time formatted in cluster timezone.
+    
+    Args:
+        cluster_tz: Timezone string (e.g., 'Asia/Shanghai', 'UTC', '+08:00')
+        
+    Returns:
+        Formatted datetime string in 'YYYY-MM-DD HH:MM:SS' format in the cluster timezone
+    """
+    tz = _get_timezone(cluster_tz)
+    now = datetime.datetime.now(tz)
+    return now.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def parse_datetime_with_tz(dt_str: str, tz: str) -> datetime.datetime:
+    """Parse datetime string assuming the given timezone.
+    
+    Args:
+        dt_str: Datetime string in 'YYYY-MM-DD HH:MM:SS' format
+        tz: Timezone string (e.g., 'Asia/Shanghai', 'UTC', '+08:00')
+        
+    Returns:
+        Timezone-aware datetime object
+    """
+    timezone = _get_timezone(tz)
+    
+    dt = datetime.datetime.strptime(dt_str, "%Y-%m-%d %H:%M:%S")
+    dt = dt.replace(tzinfo=timezone)
+    
+    return dt
+
+
+def normalize_datetime_to_tz(dt: datetime.datetime, target_tz: str) -> datetime.datetime:
+    """Convert datetime to target timezone.
+    
+    Args:
+        dt: Datetime object (timezone-aware or naive)
+        target_tz: Target timezone string (e.g., 'Asia/Shanghai', 'UTC', '+08:00')
+        
+    Returns:
+        Timezone-aware datetime object in the target timezone
+    """
+    timezone = _get_timezone(target_tz)
+    
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=datetime.timezone.utc)
+    
+    dt = dt.astimezone(timezone)
+    
+    return dt
+
+
+def _get_timezone(tz_str: str) -> Union[ZoneInfo, datetime.timezone]:
+    """Get timezone object from timezone string.
+    
+    Handles both named timezones (e.g., 'Asia/Shanghai') and offset strings (e.g., '+08:00', '-05:00').
+    
+    Args:
+        tz_str: Timezone string
+        
+    Returns:
+        ZoneInfo or timezone object
+    """
+    tz_str = tz_str.strip()
+    
+    if tz_str.upper() == "UTC" or tz_str == "+00:00" or tz_str == "-00:00":
+        return ZoneInfo("UTC")
+    
+    if tz_str.startswith(("+", "-")):
+        try:
+            hours, minutes = _parse_offset(tz_str)
+            offset = datetime.timedelta(hours=hours, minutes=minutes)
+            return datetime.timezone(offset)
+        except ValueError:
+            return ZoneInfo("UTC")
+    
+    try:
+        return ZoneInfo(tz_str)
+    except Exception:
+        return ZoneInfo("UTC")
+
+
+def _parse_offset(offset_str: str) -> tuple[int, int]:
+    """Parse timezone offset string to hours and minutes.
+    
+    Args:
+        offset_str: Offset string in format '+HH:MM' or '-HH:MM'
+        
+    Returns:
+        Tuple of (hours, minutes)
+        
+    Raises:
+        ValueError: If offset string is invalid, including:
+            - String length < 6 characters
+            - Invalid format (missing colon, invalid characters)
+            - Hours >= 24 or < 0
+            - Minutes >= 60 or < 0
+    """
+    if len(offset_str) < 6:
+        raise ValueError(f"Invalid offset format: {offset_str}")
+    
+    if offset_str[3] != ':':
+        raise ValueError(f"Invalid offset format: {offset_str} (missing colon)")
+    
+    sign = 1 if offset_str[0] == '+' else -1
+    
+    try:
+        hours = int(offset_str[1:3])
+        minutes = int(offset_str[4:6])
+    except ValueError as e:
+        raise ValueError(f"Invalid offset format: {offset_str} (non-numeric values)") from e
+    
+    if hours < 0 or hours >= 24:
+        raise ValueError(f"Invalid offset format: {offset_str} (hours must be 00-23)")
+    
+    if minutes < 0 or minutes >= 60:
+        raise ValueError(f"Invalid offset format: {offset_str} (minutes must be 00-59)")
+    
+    return sign * hours, sign * minutes
+