starrocks-br 0.5.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

starrocks_br/error_handler.py

@@ -14,9 +14,20 @@
 
 import click
 
+from . import config as config_module
 from . import exceptions
 
 
+def _get_ops_database_name(config_path: str | None) -> str:
+    if not config_path:
+        return "ops"
+    try:
+        cfg = config_module.load_config(config_path)
+        return config_module.get_ops_database(cfg)
+    except Exception:
+        return "ops"
+
+
 def display_structured_error(
     title: str,
     reason: str,
@@ -74,13 +85,14 @@ def handle_missing_option_error(exc: exceptions.MissingOptionError, config: str
 def handle_backup_label_not_found_error(
     exc: exceptions.BackupLabelNotFoundError, config: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="RESTORE FAILED",
         reason=f'The backup label "{exc.label}" does not exist in the repository'
         + (f' "{exc.repository}"' if exc.repository else "")
         + ",\nor the backup did not complete successfully.",
         what_to_do=[
-            "List available backups by querying the backup history table:\n     SELECT label, backup_type, status, finished_at FROM ops.backup_history ORDER BY finished_at DESC;",
+            f"List available backups by querying the backup history table:\n     SELECT label, backup_type, status, finished_at FROM {ops_db}.backup_history ORDER BY finished_at DESC;",
             "Check whether the backup completed successfully using StarRocks SQL:"
             + (
                 f"\n     SHOW BACKUP FROM `{exc.repository}`;"
@@ -97,11 +109,12 @@ def handle_backup_label_not_found_error(
 def handle_no_successful_full_backup_found_error(
     exc: exceptions.NoSuccessfulFullBackupFoundError, config: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="RESTORE FAILED",
         reason=f'No successful full backup was found before the incremental backup "{exc.incremental_label}".\nIncremental backups require a base full backup to restore from.',
         what_to_do=[
-            "Verify that a full backup was created before this incremental backup:\n     SELECT label, backup_type, status, finished_at FROM ops.backup_history WHERE backup_type = 'full' AND status = 'FINISHED' ORDER BY finished_at DESC;",
+            f"Verify that a full backup was created before this incremental backup:\n     SELECT label, backup_type, status, finished_at FROM {ops_db}.backup_history WHERE backup_type = 'full' AND status = 'FINISHED' ORDER BY finished_at DESC;",
             "Run a full backup first:\n     starrocks-br backup full --config "
             + (config if config else "<config.yaml>")
             + " --group <group_name>",
@@ -115,13 +128,14 @@ def handle_no_successful_full_backup_found_error(
 def handle_table_not_found_in_backup_error(
     exc: exceptions.TableNotFoundInBackupError, config: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="TABLE NOT FOUND",
         reason=f'Table "{exc.table}" was not found in backup "{exc.label}" for database "{exc.database}".',
         what_to_do=[
             "List all tables in the backup:"
             + (
-                f"\n     SELECT DISTINCT database_name, table_name FROM ops.backup_partitions WHERE label = '{exc.label}';"
+                f"\n     SELECT DISTINCT database_name, table_name FROM {ops_db}.backup_partitions WHERE label = '{exc.label}';"
                 if config
                 else ""
             ),
@@ -201,13 +215,14 @@ def handle_cluster_health_check_failed_error(
 def handle_snapshot_not_found_error(
     exc: exceptions.SnapshotNotFoundError, config: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="SNAPSHOT NOT FOUND",
         reason=f'Snapshot "{exc.snapshot_name}" was not found in repository "{exc.repository}".',
         what_to_do=[
             f"List available snapshots:\n     SHOW SNAPSHOT ON {exc.repository};",
             "Verify the snapshot name spelling is correct",
-            "Ensure the backup completed successfully:\n     SELECT * FROM ops.backup_history WHERE label = '"
+            f"Ensure the backup completed successfully:\n     SELECT * FROM {ops_db}.backup_history WHERE label = '"
             + exc.snapshot_name
             + "';",
         ],
@@ -219,13 +234,14 @@ def handle_snapshot_not_found_error(
 def handle_no_partitions_found_error(
     exc: exceptions.NoPartitionsFoundError, config: str = None, group: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="NO PARTITIONS FOUND",
         reason="No partitions were found to backup"
         + (f" for group '{exc.group_name}'" if exc.group_name else "")
         + ".",
         what_to_do=[
-            "Verify that the inventory group exists in ops.table_inventory:\n     SELECT * FROM ops.table_inventory WHERE inventory_group = "
+            f"Verify that the inventory group exists in {ops_db}.table_inventory:\n     SELECT * FROM {ops_db}.table_inventory WHERE inventory_group = "
             + (f"'{exc.group_name}';" if exc.group_name else "'<your_group>';"),
             "Check that the tables in the group have partitions",
             "Ensure the baseline backup date is correct",
@@ -238,6 +254,7 @@ def handle_no_partitions_found_error(
 def handle_no_tables_found_error(
     exc: exceptions.NoTablesFoundError, config: str = None, target_label: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="NO TABLES FOUND",
         reason="No tables were found"
@@ -250,12 +267,12 @@ def handle_no_tables_found_error(
         )
         + ".",
         what_to_do=[
-            "Verify that tables exist in the backup manifest:\n     SELECT DISTINCT database_name, table_name FROM ops.backup_partitions WHERE label = "
+            f"Verify that tables exist in the backup manifest:\n     SELECT DISTINCT database_name, table_name FROM {ops_db}.backup_partitions WHERE label = "
             + (f"'{exc.label}';" if exc.label else "'<label>';"),
-            "Check that the group name is correct in ops.table_inventory"
+            f"Check that the group name is correct in {ops_db}.table_inventory"
             if exc.group
             else "Verify the backup completed successfully",
-            "List available backups:\n     SELECT label, backup_type, status, finished_at FROM ops.backup_history ORDER BY finished_at DESC;",
+            f"List available backups:\n     SELECT label, backup_type, status, finished_at FROM {ops_db}.backup_history ORDER BY finished_at DESC;",
         ],
         inputs={
             "--target-label": exc.label or target_label,
@@ -282,6 +299,7 @@ def handle_restore_operation_cancelled_error() -> None:
 def handle_concurrency_conflict_error(
     exc: exceptions.ConcurrencyConflictError, config: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     active_job_strings = [f"{job[0]}:{job[1]}" for job in exc.active_jobs]
     first_label = exc.active_labels[0] if exc.active_labels else "unknown"
 
@@ -290,8 +308,8 @@ def handle_concurrency_conflict_error(
         reason=f"Another '{exc.scope}' job is already running.\nOnly one job of the same type can run at a time to prevent conflicts.",
         what_to_do=[
             f"Wait for the active job to complete: {', '.join(active_job_strings)}",
-            f"Check the job status in ops.run_status:\n     SELECT * FROM ops.run_status WHERE label = '{first_label}' AND state = 'ACTIVE';",
-            f"If the job is stuck, cancel it manually:\n     UPDATE ops.run_status SET state = 'CANCELLED' WHERE label = '{first_label}' AND state = 'ACTIVE';",
+            f"Check the job status in {ops_db}.run_status:\n     SELECT * FROM {ops_db}.run_status WHERE label = '{first_label}' AND state = 'ACTIVE';",
+            f"If the job is stuck, cancel it manually:\n     UPDATE {ops_db}.run_status SET state = 'CANCELLED' WHERE label = '{first_label}' AND state = 'ACTIVE';",
             "Verify the job is not actually running in StarRocks before cancelling it",
         ],
         inputs={
@@ -299,13 +317,14 @@ def handle_concurrency_conflict_error(
             "Scope": exc.scope,
             "Active jobs": ", ".join(active_job_strings),
         },
-        help_links=["Check ops.run_status table for job status"],
+        help_links=[f"Check {ops_db}.run_status table for job status"],
     )
 
 
 def handle_no_full_backup_found_error(
     exc: exceptions.NoFullBackupFoundError, config: str = None, group: str = None
 ) -> None:
+    ops_db = _get_ops_database_name(config)
     display_structured_error(
         title="NO FULL BACKUP FOUND",
         reason=f"No successful full backup was found for database '{exc.database}'.\nIncremental backups require a baseline full backup to compare against.",
@@ -313,9 +332,37 @@ def handle_no_full_backup_found_error(
             "Run a full backup first:\n     starrocks-br backup full --config "
             + (config if config else "<config.yaml>")
             + f" --group {group if group else '<group_name>'}",
-            f"Verify no full backups exist for this database:\n     SELECT label, backup_type, status, finished_at FROM ops.backup_history WHERE backup_type = 'full' AND label LIKE '{exc.database}_%' ORDER BY finished_at DESC;",
+            f"Verify no full backups exist for this database:\n     SELECT label, backup_type, status, finished_at FROM {ops_db}.backup_history WHERE backup_type = 'full' AND label LIKE '{exc.database}_%' ORDER BY finished_at DESC;",
             "After the full backup completes successfully, retry the incremental backup",
         ],
         inputs={"Database": exc.database, "--config": config, "--group": group},
         help_links=["starrocks-br backup full --help"],
     )
+
+
+def handle_invalid_tables_in_inventory_error(
+    exc: exceptions.InvalidTablesInInventoryError, config: str = None
+) -> None:
+    ops_db = _get_ops_database_name(config)
+    invalid_tables_str = ", ".join(f"'{t}'" for t in exc.invalid_tables)
+
+    display_structured_error(
+        title="INVALID TABLES IN INVENTORY",
+        reason=f"The following table(s) in the inventory do not exist in database '{exc.database}':\n{invalid_tables_str}\n\nThese tables are referenced in the table inventory but cannot be found in the actual database.",
+        what_to_do=[
+            f"Remove invalid tables from the table inventory:\n     DELETE FROM {ops_db}.table_inventory WHERE database_name = '{exc.database}' AND table_name IN ({invalid_tables_str});",
+            "Verify the table names are spelled correctly in the inventory",
+            f"Check which tables exist in the database:\n     SHOW TABLES FROM `{exc.database}`;",
+            f"Update the inventory with correct table names:\n     UPDATE {ops_db}.table_inventory SET table_name = '<correct_name>' WHERE database_name = '{exc.database}' AND table_name = '<wrong_name>';",
+        ],
+        inputs={
+            "Database": exc.database,
+            "Invalid tables": invalid_tables_str,
+            "Group": exc.group,
+            "--config": config,
+        },
+        help_links=[
+            f"Check {ops_db}.table_inventory for your inventory configuration",
+            "Run 'SHOW TABLES' to see available tables",
+        ],
+    )

starrocks_br/exceptions.py

@@ -122,3 +122,17 @@ class NoFullBackupFoundError(StarRocksBRError):
     def __init__(self, database: str):
         self.database = database
         super().__init__(f"No successful full backup found for database '{database}'")
+
+
+class InvalidTablesInInventoryError(StarRocksBRError):
+    def __init__(self, database: str, invalid_tables: list[str], group: str = None):
+        self.database = database
+        self.invalid_tables = invalid_tables
+        self.group = group
+        tables_str = ", ".join(f"'{t}'" for t in invalid_tables)
+        if group:
+            super().__init__(
+                f"Invalid tables in inventory group '{group}' for database '{database}': {tables_str}"
+            )
+        else:
+            super().__init__(f"Invalid tables for database '{database}': {tables_str}")

starrocks_br/executor.py

@@ -181,6 +181,7 @@ def execute_backup(
     backup_type: Literal["incremental", "full"] = None,
     scope: str = "backup",
     database: str | None = None,
+    ops_database: str = "ops",
 ) -> dict:
     """Execute a complete backup workflow: submit command and monitor progress.
 
@@ -193,6 +194,7 @@ def execute_backup(
         backup_type: Type of backup (for logging)
         scope: Job scope (for concurrency control)
         database: Database name (required for SHOW BACKUP)
+        ops_database: Name of ops database (default: "ops")
 
     Returns dictionary with keys: success, final_status, error_message
     """
@@ -233,13 +235,18 @@ def execute_backup(
                     "finished_at": finished_at,
                     "error_message": None if success else (final_status["state"] or ""),
                 },
+                ops_database=ops_database,
             )
         except Exception:
             pass
 
         try:
             concurrency.complete_job_slot(
-                db, scope=scope, label=label, final_state=final_status["state"]
+                db,
+                scope=scope,
+                label=label,
+                final_state=final_status["state"],
+                ops_database=ops_database,
             )
         except Exception:
             pass
@@ -271,7 +278,7 @@ def _build_error_message(final_status: dict, label: str, database: str) -> str:
             f"Backup tracking lost for '{label}' in database '{database}'. "
             f"Another backup operation overwrote the last backup status visible in SHOW BACKUP. "
             f"This indicates a concurrency issue - only one backup per database should run at a time. "
-            f"Recommendation: Use ops.run_status concurrency control to prevent simultaneous backups, "
+            f"Recommendation: Use run_status concurrency control to prevent simultaneous backups, "
             f"or verify if another tool/user is running backups on this database."
         )
     elif state == "CANCELLED":

starrocks_br/history.py

@@ -15,8 +15,8 @@
 from . import logger
 
 
-def log_backup(db, entry: dict[str, str | None]) -> None:
-    """Write a backup history entry to ops.backup_history.
+def log_backup(db, entry: dict[str, str | None], ops_database: str = "ops") -> None:
+    """Write a backup history entry to the backup_history table.
 
     Expected keys in entry:
       - job_id (optional; auto-generated if missing)
@@ -42,7 +42,7 @@ def log_backup(db, entry: dict[str, str | None]) -> None:
         return "'" + str(val).replace("'", "''") + "'"
 
     sql = f"""
-    INSERT INTO ops.backup_history (
+    INSERT INTO {ops_database}.backup_history (
         label, backup_type, status, repository, started_at, finished_at, error_message
     ) VALUES (
         {esc(label)}, {esc(backup_type)}, {esc(status)}, {esc(repository)},
@@ -57,8 +57,8 @@ def log_backup(db, entry: dict[str, str | None]) -> None:
         raise
 
 
-def log_restore(db, entry: dict[str, str | None]) -> None:
-    """Write a restore history entry to ops.restore_history.
+def log_restore(db, entry: dict[str, str | None], ops_database: str = "ops") -> None:
+    """Write a restore history entry to the restore_history table.
 
     Expected keys in entry:
       - job_id
@@ -87,12 +87,12 @@ def log_restore(db, entry: dict[str, str | None]) -> None:
         return "'" + str(val).replace("'", "''") + "'"
 
     sql = f"""
-    INSERT INTO ops.restore_history (
-        job_id, backup_label, restore_type, status, repository, 
+    INSERT INTO {ops_database}.restore_history (
+        job_id, backup_label, restore_type, status, repository,
         started_at, finished_at, error_message, verification_checksum
     ) VALUES (
-        {esc(job_id)}, {esc(backup_label)}, {esc(restore_type)}, {esc(status)}, 
-        {esc(repository)}, {esc(started_at)}, {esc(finished_at)}, 
+        {esc(job_id)}, {esc(backup_label)}, {esc(restore_type)}, {esc(status)},
+        {esc(repository)}, {esc(started_at)}, {esc(finished_at)},
         {esc(error_message)}, {esc(verification_checksum)}
     )
     """

starrocks_br/labels.py

@@ -21,12 +21,13 @@ def determine_backup_label(
     backup_type: Literal["incremental", "full"],
     database_name: str,
     custom_name: str | None = None,
+    ops_database: str = "ops",
 ) -> str:
     """Determine a unique backup label for the given parameters.
 
     This is the single entry point for all backup label generation. It handles both
     custom names and auto-generated date-based labels, ensuring uniqueness by checking
-    the ops.backup_history table.
+    the backup_history table in the configured ops database.
 
     Args:
         db: Database connection
@@ -34,6 +35,7 @@ def determine_backup_label(
         database_name: Name of the database being backed up
         custom_name: Optional custom name for the backup. If provided, this becomes
                     the base label. If None, generates a date-based label.
+        ops_database: Name of the database containing operational tables. Defaults to "ops".
 
     Returns:
         Unique label string that doesn't conflict with existing backups
@@ -44,9 +46,9 @@ def determine_backup_label(
         today = datetime.now().strftime("%Y%m%d")
         base_label = f"{database_name}_{today}_{backup_type}"
 
-    query = """
+    query = f"""
     SELECT label
-    FROM ops.backup_history
+    FROM {ops_database}.backup_history
     WHERE label LIKE %s
     ORDER BY label
     """

starrocks_br/planner.py

@@ -18,7 +18,7 @@ import hashlib
 from starrocks_br import exceptions, logger, timezone, utils
 
 
-def find_latest_full_backup(db, database: str) -> dict[str, str] | None:
+def find_latest_full_backup(db, database: str, ops_database: str = "ops") -> dict[str, str] | None:
     """Find the latest successful full backup for a database.
 
     Args:
@@ -31,7 +31,7 @@ def find_latest_full_backup(db, database: str) -> dict[str, str] | None:
     """
     query = f"""
     SELECT label, backup_type, finished_at
-    FROM ops.backup_history
+    FROM {ops_database}.backup_history
     WHERE backup_type = 'full'
     AND status = 'FINISHED'
     AND label LIKE {utils.quote_value(f"{database}_%")}
@@ -56,7 +56,7 @@ def find_latest_full_backup(db, database: str) -> dict[str, str] | None:
     return {"label": row[0], "backup_type": row[1], "finished_at": finished_at}
 
 
-def find_tables_by_group(db, group_name: str) -> list[dict[str, str]]:
+def find_tables_by_group(db, group_name: str, ops_database: str = "ops") -> list[dict[str, str]]:
     """Find tables belonging to a specific inventory group.
 
     Returns list of dictionaries with keys: database, table.
@@ -64,7 +64,7 @@ def find_tables_by_group(db, group_name: str) -> list[dict[str, str]]:
     """
     query = f"""
     SELECT database_name, table_name
-    FROM ops.table_inventory
+    FROM {ops_database}.table_inventory
     WHERE inventory_group = {utils.quote_value(group_name)}
     ORDER BY database_name, table_name
     """
@@ -72,8 +72,49 @@ def find_tables_by_group(db, group_name: str) -> list[dict[str, str]]:
     return [{"database": row[0], "table": row[1]} for row in rows]
 
 
+def validate_tables_exist(
+    db, database: str, tables: list[dict[str, str]], group: str = None
+) -> None:
+    """Validate that tables in the inventory actually exist in the database.
+
+    Args:
+        db: Database connection
+        database: Database name to validate tables against
+        tables: List of tables with keys: database, table
+        group: Optional inventory group name for better error messages
+
+    Raises:
+        InvalidTablesInInventoryError: If any tables don't exist in the database
+    """
+    if not tables:
+        return
+
+    db_tables = [t for t in tables if t["database"] == database and t["table"] != "*"]
+
+    if not db_tables:
+        return
+
+    show_tables_query = f"SHOW TABLES FROM {utils.quote_identifier(database)}"
+    existing_tables_rows = db.query(show_tables_query)
+    existing_tables = {row[0] for row in existing_tables_rows}
+
+    invalid_tables = []
+    for table_entry in db_tables:
+        table_name = table_entry["table"]
+        if table_name not in existing_tables:
+            invalid_tables.append(table_name)
+
+    if invalid_tables:
+        raise exceptions.InvalidTablesInInventoryError(database, invalid_tables, group)
+
+
 def find_recent_partitions(
-    db, database: str, baseline_backup_label: str | None = None, *, group_name: str
+    db,
+    database: str,
+    baseline_backup_label: str | None = None,
+    *,
+    group_name: str,
+    ops_database: str = "ops",
 ) -> list[dict[str, str]]:
     """Find partitions updated since baseline for tables in the given inventory group.
 
@@ -91,7 +132,7 @@ def find_recent_partitions(
     if baseline_backup_label:
         baseline_query = f"""
         SELECT finished_at
-        FROM ops.backup_history
+        FROM {ops_database}.backup_history
         WHERE label = {utils.quote_value(baseline_backup_label)}
         AND status = 'FINISHED'
         """
@@ -100,7 +141,7 @@ def find_recent_partitions(
             raise exceptions.BackupLabelNotFoundError(baseline_backup_label)
         baseline_time_raw = baseline_rows[0][0]
     else:
-        latest_backup = find_latest_full_backup(db, database)
+        latest_backup = find_latest_full_backup(db, database, ops_database)
         if not latest_backup:
             raise exceptions.NoFullBackupFoundError(database)
         baseline_time_raw = latest_backup["finished_at"]
@@ -114,7 +155,7 @@ def find_recent_partitions(
 
     baseline_dt = timezone.parse_datetime_with_tz(baseline_time_str, cluster_tz)
 
-    group_tables = find_tables_by_group(db, group_name)
+    group_tables = find_tables_by_group(db, group_name, ops_database)
 
     if not group_tables:
         return []
@@ -218,7 +259,7 @@ def build_incremental_backup_command(
 
 
 def build_full_backup_command(
-    db, group_name: str, repository: str, label: str, database: str
+    db, group_name: str, repository: str, label: str, database: str, ops_database: str = "ops"
 ) -> str:
     """Build BACKUP command for an inventory group.
 
@@ -226,7 +267,7 @@ def build_full_backup_command(
     simple BACKUP DATABASE command. Otherwise, generate ON (TABLE ...) list for
     the specific tables within the database.
     """
-    tables = find_tables_by_group(db, group_name)
+    tables = find_tables_by_group(db, group_name, ops_database)
 
     db_entries = [t for t in tables if t["database"] == database]
     if not db_entries:
@@ -245,8 +286,10 @@ def build_full_backup_command(
     ON ({on_clause})"""
 
 
-def record_backup_partitions(db, label: str, partitions: list[dict[str, str]]) -> None:
-    """Record partition metadata for a backup in ops.backup_partitions table.
+def record_backup_partitions(
+    db, label: str, partitions: list[dict[str, str]], ops_database: str = "ops"
+) -> None:
+    """Record partition metadata for a backup in the backup_partitions table.
 
     Args:
         db: Database connection
@@ -263,7 +306,7 @@ def record_backup_partitions(db, label: str, partitions: list[dict[str, str]]) -
         key_hash = hashlib.md5(composite_key.encode("utf-8")).hexdigest()
 
         db.execute(f"""
-            INSERT INTO ops.backup_partitions
+            INSERT INTO {ops_database}.backup_partitions
             (key_hash, label, database_name, table_name, partition_name)
             VALUES ({utils.quote_value(key_hash)}, {utils.quote_value(label)}, {utils.quote_value(partition["database"])}, {utils.quote_value(partition["table"])}, {utils.quote_value(partition["partition_name"])})
         """)

starrocks_br/repository.py

@@ -30,7 +30,7 @@ def ensure_repository(db, name: str) -> None:
         raise RuntimeError(
             f"Repository '{name}' not found. Please create it first using:\n"
             f"  CREATE REPOSITORY {name} WITH BROKER ON LOCATION '...' PROPERTIES(...)\n"
-            f"For examples, see: https://docs.starrocks.io/docs/sql-reference/sql-statements/data-definition/backup_restore/CREATE_REPOSITORY/"
+            f"For examples, see: https://docs.starrocks.io/docs/sql-reference/sql-statements/backup_restore/CREATE_REPOSITORY/"
         )
 
     # SHOW REPOSITORIES returns: RepoId, RepoName, CreateTime, IsReadOnly, Location, Broker, ErrMsg