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

starrocks_br/__init__.py

@@ -1 +1,15 @@
+# Copyright 2025 deep-bi
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 __all__ = ["cli", "config"]

starrocks_br/cli.py

@@ -1,3 +1,17 @@
+# Copyright 2025 deep-bi
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 import sys
 
@@ -84,13 +98,13 @@ def cli(ctx, verbose):
 @cli.command("init")
 @click.option("--config", required=True, help="Path to config YAML file")
 def init(config):
-    """Initialize ops database and control tables.
+    """Initialize operations 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
+    Creates the operations database (default: 'ops') with required tables:
+    - table_inventory: Inventory groups mapping to databases/tables
+    - backup_history: Backup operation history
+    - restore_history: Restore operation history
+    - run_status: Job concurrency control
 
     Run this once before using backup/restore commands.
     """
@@ -98,6 +112,9 @@ def init(config):
         cfg = config_module.load_config(config)
         config_module.validate_config(cfg)
 
+        ops_database = config_module.get_ops_database(cfg)
+        table_inventory_entries = config_module.get_table_inventory_entries(cfg)
+
         database = db.StarRocksDB(
             host=cfg["host"],
             port=cfg["port"],
@@ -107,23 +124,43 @@ def init(config):
             tls_config=cfg.get("tls"),
         )
 
+        ops_database = config_module.get_ops_database(cfg)
+
         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("Validating repository...")
+            repository.ensure_repository(database, cfg["repository"])
             logger.info("")
-            logger.info("2. Run your first backup:")
-            logger.info(
-                "   starrocks-br backup incremental --group my_daily_incremental --config config.yaml"
+
+            logger.info("Initializing ops schema...")
+            schema.initialize_ops_schema(
+                database, ops_database=ops_database, table_inventory_entries=table_inventory_entries
             )
+            logger.info("")
+
+            if table_inventory_entries:
+                logger.success(
+                    f"Table inventory bootstrapped from config with {len(table_inventory_entries)} entries"
+                )
+                logger.info("")
+                logger.info("Next steps:")
+                logger.info("1. Run your first backup:")
+                logger.info(
+                    f"   starrocks-br backup incremental --group <your_group_name> --config {config}"
+                )
+            else:
+                logger.info("Next steps:")
+                logger.info("1. Insert your table inventory records:")
+                logger.info(f"   INSERT INTO {ops_database}.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 exceptions.ConfigFileNotFoundError as e:
         error_handler.handle_config_file_not_found_error(e)
@@ -182,13 +219,17 @@ def backup_incremental(config, baseline_backup, group, name):
             tls_config=cfg.get("tls"),
         )
 
+        ops_database = config_module.get_ops_database(cfg)
+
         with database:
-            was_created = schema.ensure_ops_schema(database)
+            was_created = schema.ensure_ops_schema(database, ops_database=ops_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!")
+                logger.warning(
+                    "Remember to populate the table_inventory table with your backup groups!"
+                )
                 sys.exit(1)  # Exit if schema was just created, requires user action
 
             healthy, message = health.check_cluster_health(database)
@@ -207,6 +248,7 @@ def backup_incremental(config, baseline_backup, group, name):
                 backup_type="incremental",
                 database_name=cfg["database"],
                 custom_name=name,
+                ops_database=ops_database,
             )
 
             logger.success(f"Generated label: {label}")
@@ -225,7 +267,11 @@ def backup_incremental(config, baseline_backup, group, name):
                     )
 
             partitions = planner.find_recent_partitions(
-                database, cfg["database"], baseline_backup_label=baseline_backup, group_name=group
+                database,
+                cfg["database"],
+                baseline_backup_label=baseline_backup,
+                group_name=group,
+                ops_database=ops_database,
             )
 
             if not partitions:
@@ -238,9 +284,11 @@ def backup_incremental(config, baseline_backup, group, name):
                 partitions, cfg["repository"], label, cfg["database"]
             )
 
-            concurrency.reserve_job_slot(database, scope="backup", label=label)
+            concurrency.reserve_job_slot(
+                database, scope="backup", label=label, ops_database=ops_database
+            )
 
-            planner.record_backup_partitions(database, label, partitions)
+            planner.record_backup_partitions(database, label, partitions, ops_database=ops_database)
 
             logger.success("Job slot reserved")
             logger.info(f"Starting incremental backup for group '{group}'...")
@@ -251,6 +299,7 @@ def backup_incremental(config, baseline_backup, group, name):
                 backup_type="incremental",
                 scope="backup",
                 database=cfg["database"],
+                ops_database=ops_database,
             )
 
             if result["success"]:
@@ -274,7 +323,7 @@ def backup_incremental(config, baseline_backup, group, name):
                 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.tip("Enable run_status concurrency checks to prevent this.")
                 logger.error(f"{result['error_message']}")
                 sys.exit(1)
 
@@ -334,13 +383,17 @@ def backup_full(config, group, name):
             tls_config=cfg.get("tls"),
         )
 
+        ops_database = config_module.get_ops_database(cfg)
+
         with database:
-            was_created = schema.ensure_ops_schema(database)
+            was_created = schema.ensure_ops_schema(database, ops_database=ops_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!")
+                logger.warning(
+                    "Remember to populate the table_inventory table with your backup groups!"
+                )
                 sys.exit(1)  # Exit if schema was just created, requires user action
 
             healthy, message = health.check_cluster_health(database)
@@ -355,13 +408,25 @@ def backup_full(config, group, name):
             logger.success(f"Repository '{cfg['repository']}' verified")
 
             label = labels.determine_backup_label(
-                db=database, backup_type="full", database_name=cfg["database"], custom_name=name
+                db=database,
+                backup_type="full",
+                database_name=cfg["database"],
+                custom_name=name,
+                ops_database=ops_database,
             )
 
             logger.success(f"Generated label: {label}")
 
+            tables = planner.find_tables_by_group(database, group, ops_database)
+            planner.validate_tables_exist(database, cfg["database"], tables, group)
+
             backup_command = planner.build_full_backup_command(
-                database, group, cfg["repository"], label, cfg["database"]
+                database,
+                group,
+                cfg["repository"],
+                label,
+                cfg["database"],
+                ops_database=ops_database,
             )
 
             if not backup_command:
@@ -375,9 +440,13 @@ def backup_full(config, group, name):
                 database, cfg["database"], tables
             )
 
-            concurrency.reserve_job_slot(database, scope="backup", label=label)
+            concurrency.reserve_job_slot(
+                database, scope="backup", label=label, ops_database=ops_database
+            )
 
-            planner.record_backup_partitions(database, label, all_partitions)
+            planner.record_backup_partitions(
+                database, label, all_partitions, ops_database=ops_database
+            )
 
             logger.success("Job slot reserved")
             logger.info(f"Starting full backup for group '{group}'...")
@@ -388,6 +457,7 @@ def backup_full(config, group, name):
                 backup_type="full",
                 scope="backup",
                 database=cfg["database"],
+                ops_database=ops_database,
             )
 
             if result["success"]:
@@ -405,10 +475,13 @@ def backup_full(config, group, name):
                 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.tip("Enable run_status concurrency checks to prevent this.")
                 logger.error(f"{result['error_message']}")
                 sys.exit(1)
 
+    except exceptions.InvalidTablesInInventoryError as e:
+        error_handler.handle_invalid_tables_in_inventory_error(e, config)
+        sys.exit(1)
     except exceptions.ConcurrencyConflictError as e:
         error_handler.handle_concurrency_conflict_error(e, config)
         sys.exit(1)
@@ -485,13 +558,17 @@ def restore_command(config, target_label, group, table, rename_suffix, yes):
             tls_config=cfg.get("tls"),
         )
 
+        ops_database = config_module.get_ops_database(cfg)
+
         with database:
-            was_created = schema.ensure_ops_schema(database)
+            was_created = schema.ensure_ops_schema(database, ops_database=ops_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!")
+                logger.warning(
+                    "Remember to populate the table_inventory table with your backup groups!"
+                )
                 sys.exit(1)  # Exit if schema was just created, requires user action
 
             healthy, message = health.check_cluster_health(database)
@@ -507,7 +584,9 @@ def restore_command(config, target_label, group, table, rename_suffix, yes):
 
             logger.info(f"Finding restore sequence for target backup: {target_label}")
 
-            restore_pair = restore.find_restore_pair(database, target_label)
+            restore_pair = restore.find_restore_pair(
+                database, target_label, ops_database=ops_database
+            )
             logger.success(f"Found restore sequence: {' -> '.join(restore_pair)}")
 
             logger.info("Determining tables to restore from backup manifest...")
@@ -518,6 +597,7 @@ def restore_command(config, target_label, group, table, rename_suffix, yes):
                 group=group,
                 table=table,
                 database=cfg["database"] if table else None,
+                ops_database=ops_database,
             )
 
             if not tables_to_restore:
@@ -535,6 +615,7 @@ def restore_command(config, target_label, group, table, rename_suffix, yes):
                 tables_to_restore,
                 rename_suffix,
                 skip_confirmation=yes,
+                ops_database=ops_database,
             )
 
             if result["success"]:

starrocks_br/concurrency.py

@@ -1,47 +1,67 @@
+# Copyright 2025 deep-bi
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Literal
 
 from . import exceptions, logger, utils
 
 
-def reserve_job_slot(db, scope: str, label: str) -> None:
-    """Reserve a job slot in ops.run_status to prevent overlapping jobs.
+def reserve_job_slot(db, scope: str, label: str, ops_database: str = "ops") -> None:
+    """Reserve a job slot in the run_status table 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)
+    active_jobs = _get_active_jobs_for_scope(db, scope, ops_database)
 
     if not active_jobs:
-        _insert_new_job(db, scope, label)
+        _insert_new_job(db, scope, label, ops_database)
         return
 
-    _handle_active_job_conflicts(db, scope, active_jobs)
+    _handle_active_job_conflicts(db, scope, active_jobs, ops_database)
 
-    _insert_new_job(db, scope, label)
+    _insert_new_job(db, scope, label, ops_database)
 
 
-def _get_active_jobs_for_scope(db, scope: str) -> list[tuple[str, str, str]]:
+def _get_active_jobs_for_scope(
+    db, scope: str, ops_database: str = "ops"
+) -> 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'")
+    rows = db.query(
+        f"SELECT scope, label, state FROM {ops_database}.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:
+def _handle_active_job_conflicts(
+    db, scope: str, active_jobs: list[tuple[str, str, str]], ops_database: str = "ops"
+) -> 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)
+        if _can_heal_stale_job(active_scope, active_label, db, ops_database):
+            _cleanup_stale_job(db, active_scope, active_label, ops_database)
             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:
+def _can_heal_stale_job(scope: str, label: str, db, ops_database: str = "ops") -> 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)
+    return _is_backup_job_stale(db, label, ops_database)
 
 
 def _raise_concurrency_conflict(scope: str, active_jobs: list[tuple[str, str, str]]) -> None:
@@ -49,22 +69,22 @@ def _raise_concurrency_conflict(scope: str, active_jobs: list[tuple[str, str, st
     raise exceptions.ConcurrencyConflictError(scope, active_jobs)
 
 
-def _insert_new_job(db, scope: str, label: str) -> None:
+def _insert_new_job(db, scope: str, label: str, ops_database: str = "ops") -> None:
     """Insert a new active job record."""
     sql = f"""
-        INSERT INTO ops.run_status (scope, label, state, started_at)
+        INSERT INTO {ops_database}.run_status (scope, label, state, started_at)
         VALUES ({utils.quote_value(scope)}, {utils.quote_value(label)}, 'ACTIVE', NOW())
     """
     db.execute(sql)
 
 
-def _is_backup_job_stale(db, label: str) -> bool:
+def _is_backup_job_stale(db, label: str, ops_database: str = "ops") -> 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)
+        user_databases = _get_user_databases(db, ops_database)
 
         for database_name in user_databases:
             job_status = _check_backup_job_in_database(db, database_name, label)
@@ -84,9 +104,9 @@ def _is_backup_job_stale(db, label: str) -> bool:
         return False
 
 
-def _get_user_databases(db) -> list[str]:
+def _get_user_databases(db, ops_database: str = "ops") -> list[str]:
     """Get list of user databases (excluding system databases)."""
-    system_databases = {"information_schema", "mysql", "sys", "ops"}
+    system_databases = {"information_schema", "mysql", "sys", ops_database}
 
     databases = db.query("SHOW DATABASES")
     return [
@@ -145,10 +165,10 @@ def _extract_backup_info(result) -> tuple[str, str]:
     return snapshot_name, state
 
 
-def _cleanup_stale_job(db, scope: str, label: str) -> None:
+def _cleanup_stale_job(db, scope: str, label: str, ops_database: str = "ops") -> None:
     """Clean up a stale job by updating its state to CANCELLED."""
     sql = f"""
-        UPDATE ops.run_status
+        UPDATE {ops_database}.run_status
         SET state='CANCELLED', finished_at=NOW()
         WHERE scope={utils.quote_value(scope)} AND label={utils.quote_value(label)} AND state='ACTIVE'
     """
@@ -156,14 +176,18 @@ def _cleanup_stale_job(db, scope: str, label: str) -> None:
 
 
 def complete_job_slot(
-    db, scope: str, label: str, final_state: Literal["FINISHED", "FAILED", "CANCELLED"]
+    db,
+    scope: str,
+    label: str,
+    final_state: Literal["FINISHED", "FAILED", "CANCELLED"],
+    ops_database: str = "ops",
 ) -> None:
     """Complete job slot and persist final state.
 
     Simple approach: update the same row by scope/label.
     """
     sql = f"""
-        UPDATE ops.run_status
+        UPDATE {ops_database}.run_status
         SET state={utils.quote_value(final_state)}, finished_at=NOW()
         WHERE scope={utils.quote_value(scope)} AND label={utils.quote_value(label)}
     """

starrocks_br/config.py

@@ -1,3 +1,17 @@
+# Copyright 2025 deep-bi
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any
 
 import yaml
@@ -43,6 +57,34 @@ def validate_config(config: dict[str, Any]) -> None:
             raise exceptions.ConfigValidationError(f"Missing required config field: {field}")
 
     _validate_tls_section(config.get("tls"))
+    _validate_table_inventory_section(config.get("table_inventory"))
+
+
+def get_ops_database(config: dict[str, Any]) -> str:
+    """Get the ops database name from config, defaulting to 'ops'."""
+    return config.get("ops_database", "ops")
+
+
+def get_table_inventory_entries(config: dict[str, Any]) -> list[tuple[str, str, str]]:
+    """Extract table inventory entries from config.
+
+    Args:
+        config: Configuration dictionary
+
+    Returns:
+        List of tuples (group, database, table)
+    """
+    table_inventory = config.get("table_inventory")
+    if not table_inventory:
+        return []
+
+    entries = []
+    for group_entry in table_inventory:
+        group = group_entry["group"]
+        for table_entry in group_entry["tables"]:
+            entries.append((group, table_entry["database"], table_entry["table"]))
+
+    return entries
 
 
 def _validate_tls_section(tls_config) -> None:
@@ -74,3 +116,50 @@ def _validate_tls_section(tls_config) -> None:
             raise exceptions.ConfigValidationError(
                 "TLS configuration field 'tls_versions' must be a list of strings if provided"
             )
+
+
+def _validate_table_inventory_section(table_inventory) -> None:
+    if table_inventory is None:
+        return
+
+    if not isinstance(table_inventory, list):
+        raise exceptions.ConfigValidationError("'table_inventory' must be a list")
+
+    for entry in table_inventory:
+        if not isinstance(entry, dict):
+            raise exceptions.ConfigValidationError(
+                "Each entry in 'table_inventory' must be a dictionary"
+            )
+
+        if "group" not in entry:
+            raise exceptions.ConfigValidationError(
+                "Each entry in 'table_inventory' must have a 'group' field"
+            )
+
+        if "tables" not in entry:
+            raise exceptions.ConfigValidationError(
+                "Each entry in 'table_inventory' must have a 'tables' field"
+            )
+
+        if not isinstance(entry["group"], str):
+            raise exceptions.ConfigValidationError("'group' field must be a string")
+
+        tables = entry["tables"]
+        if not isinstance(tables, list):
+            raise exceptions.ConfigValidationError("'tables' field must be a list")
+
+        for table_entry in tables:
+            if not isinstance(table_entry, dict):
+                raise exceptions.ConfigValidationError("Each table entry must be a dictionary")
+
+            if "database" not in table_entry or "table" not in table_entry:
+                raise exceptions.ConfigValidationError(
+                    "Each table entry must have 'database' and 'table' fields"
+                )
+
+            if not isinstance(table_entry["database"], str) or not isinstance(
+                table_entry["table"], str
+            ):
+                raise exceptions.ConfigValidationError(
+                    "'database' and 'table' fields must be strings"
+                )

starrocks_br/db.py

@@ -1,3 +1,17 @@
+# Copyright 2025 deep-bi
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any
 
 import mysql.connector