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

starrocks_br/repository.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 __future__ import annotations
 
 
@@ -16,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

starrocks_br/restore.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 datetime
 import time
 
@@ -187,9 +201,22 @@ def execute_restore(
     max_polls: int = MAX_POLLS,
     poll_interval: float = 1.0,
     scope: str = "restore",
+    ops_database: str = "ops",
 ) -> dict:
     """Execute a complete restore workflow: submit command and monitor progress.
 
+    Args:
+        db: Database connection
+        restore_command: Restore SQL command to execute
+        backup_label: Label of the backup being restored
+        restore_type: Type of restore operation
+        repository: Repository name
+        database: Database name
+        max_polls: Maximum polling attempts
+        poll_interval: Seconds between polls
+        scope: Job scope (for concurrency control)
+        ops_database: Name of ops database (default: "ops")
+
     Returns dictionary with keys: success, final_status, error_message
     """
     cluster_tz = db.timezone
@@ -226,13 +253,18 @@ def execute_restore(
                     "finished_at": finished_at,
                     "error_message": None if success else final_status["state"],
                 },
+                ops_database=ops_database,
             )
         except Exception as e:
             logger.error(f"Failed to log restore history: {str(e)}")
 
         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 as e:
             logger.error(f"Failed to complete job slot: {str(e)}")
@@ -250,7 +282,7 @@ def execute_restore(
         return {"success": False, "final_status": None, "error_message": str(e)}
 
 
-def find_restore_pair(db, target_label: str) -> list[str]:
+def find_restore_pair(db, target_label: str, ops_database: str = "ops") -> list[str]:
     """Find the correct sequence of backups needed for restore.
 
     Args:
@@ -266,7 +298,7 @@ def find_restore_pair(db, target_label: str) -> list[str]:
     """
     query = f"""
     SELECT label, backup_type, finished_at
-    FROM ops.backup_history
+    FROM {ops_database}.backup_history
     WHERE label = {utils.quote_value(target_label)}
     AND status = 'FINISHED'
     """
@@ -285,7 +317,7 @@ def find_restore_pair(db, target_label: str) -> list[str]:
 
         full_backup_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_name}_%")}
@@ -312,6 +344,7 @@ def get_tables_from_backup(
     group: str | None = None,
     table: str | None = None,
     database: str | None = None,
+    ops_database: str = "ops",
 ) -> list[str]:
     """Get list of tables to restore from backup manifest.
 
@@ -340,7 +373,7 @@ def get_tables_from_backup(
 
     query = f"""
     SELECT DISTINCT database_name, table_name
-    FROM ops.backup_partitions
+    FROM {ops_database}.backup_partitions
     WHERE label = {utils.quote_value(label)}
     ORDER BY database_name, table_name
     """
@@ -363,7 +396,7 @@ def get_tables_from_backup(
     if group:
         group_query = f"""
         SELECT database_name, table_name
-        FROM ops.table_inventory
+        FROM {ops_database}.table_inventory
         WHERE inventory_group = {utils.quote_value(group)}
         """
 
@@ -390,6 +423,35 @@ def get_tables_from_backup(
     return tables
 
 
+def get_partitions_from_backup(
+    db, label: str, table: str, ops_database: str = "ops"
+) -> list[str]:
+    """Get list of partitions for a specific table from backup manifest.
+
+    Args:
+        db: Database connection
+        label: Backup label
+        table: Table name in format 'database.table'
+        ops_database: Operations database name
+
+    Returns:
+        List of partition names for the table in this backup
+    """
+    database_name, table_name = table.split(".", 1)
+
+    query = f"""
+    SELECT partition_name
+    FROM {ops_database}.backup_partitions
+    WHERE label = {utils.quote_value(label)}
+      AND database_name = {utils.quote_value(database_name)}
+      AND table_name = {utils.quote_value(table_name)}
+    ORDER BY partition_name
+    """
+
+    rows = db.query(query)
+    return [row[0] for row in rows]
+
+
 def execute_restore_flow(
     db,
     repo_name: str,
@@ -397,6 +459,7 @@ def execute_restore_flow(
     tables_to_restore: list[str],
     rename_suffix: str = "_restored",
     skip_confirmation: bool = False,
+    ops_database: str = "ops",
 ) -> dict:
     """Execute the complete restore flow with safety measures.
 
@@ -407,6 +470,7 @@ def execute_restore_flow(
         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
+        ops_database: Name of ops database (default: "ops")
 
     Returns:
         Dictionary with success status and details
@@ -438,59 +502,122 @@ def execute_restore_flow(
         database_name = tables_to_restore[0].split(".")[0]
 
         base_label = restore_pair[0]
-        logger.info("")
-        logger.info(f"Step 1: Restoring base backup '{base_label}'...")
-
-        base_timestamp = get_snapshot_timestamp(db, repo_name, base_label)
-
-        base_restore_command = _build_restore_command_with_rename(
-            base_label, repo_name, tables_to_restore, rename_suffix, database_name, base_timestamp
-        )
-
-        base_result = execute_restore(
-            db, base_restore_command, base_label, "full", repo_name, database_name, scope="restore"
-        )
-
-        if not base_result["success"]:
-            return {
-                "success": False,
-                "error_message": f"Base restore failed: {base_result['error_message']}",
-            }
 
-        logger.success("Base restore completed successfully")
+        tables_in_base = get_tables_from_backup(db, base_label, ops_database=ops_database)
+        tables_to_restore_from_base = [t for t in tables_to_restore if t in tables_in_base]
 
-        if len(restore_pair) > 1:
-            incremental_label = restore_pair[1]
+        if tables_to_restore_from_base:
             logger.info("")
-            logger.info(f"Step 2: Applying incremental backup '{incremental_label}'...")
+            logger.info(f"Step 1: Restoring base backup '{base_label}'...")
 
-            incremental_timestamp = get_snapshot_timestamp(db, repo_name, incremental_label)
+            base_timestamp = get_snapshot_timestamp(db, repo_name, base_label)
 
-            incremental_restore_command = _build_restore_command_without_rename(
-                incremental_label,
+            base_restore_command = _build_restore_command_with_rename(
+                base_label,
                 repo_name,
-                tables_to_restore,
+                tables_to_restore_from_base,
+                rename_suffix,
                 database_name,
-                incremental_timestamp,
+                base_timestamp,
             )
 
-            incremental_result = execute_restore(
+            base_result = execute_restore(
                 db,
-                incremental_restore_command,
-                incremental_label,
-                "incremental",
+                base_restore_command,
+                base_label,
+                "full",
                 repo_name,
                 database_name,
                 scope="restore",
+                ops_database=ops_database,
             )
 
-            if not incremental_result["success"]:
+            if not base_result["success"]:
                 return {
                     "success": False,
-                    "error_message": f"Incremental restore failed: {incremental_result['error_message']}",
+                    "error_message": f"Base restore failed: {base_result['error_message']}",
                 }
 
-            logger.success("Incremental restore completed successfully")
+            logger.success("Base restore completed successfully")
+        else:
+            logger.info("")
+            logger.info(
+                f"Step 1: Skipping base backup '{base_label}' (no requested tables in this backup)"
+            )
+
+        if len(restore_pair) > 1:
+            incremental_label = restore_pair[1]
+
+            tables_in_incremental = get_tables_from_backup(
+                db, incremental_label, ops_database=ops_database
+            )
+            tables_to_restore_from_incremental = [
+                t for t in tables_to_restore if t in tables_in_incremental
+            ]
+
+            if not tables_to_restore_from_incremental:
+                logger.info("")
+                logger.info(
+                    f"Step 2: Skipping incremental backup '{incremental_label}' (no requested tables in this backup)"
+                )
+            else:
+                logger.info("")
+                logger.info(f"Step 2: Applying incremental backup '{incremental_label}'...")
+
+                incremental_timestamp = get_snapshot_timestamp(db, repo_name, incremental_label)
+
+                for table in tables_to_restore_from_incremental:
+                    partitions = get_partitions_from_backup(
+                        db, incremental_label, table, ops_database=ops_database
+                    )
+
+                    if not partitions:
+                        logger.warning(f"No partitions found for {table} in {incremental_label}, skipping")
+                        continue
+
+                    table_was_in_base = table in tables_to_restore_from_base
+
+                    if table_was_in_base:
+                        _, table_name = table.split(".", 1)
+                        target_table_name = f"{table_name}{rename_suffix}"
+                        incremental_restore_command = _build_partition_restore_command(
+                            incremental_label,
+                            repo_name,
+                            f"{database_name}.{target_table_name}",
+                            partitions,
+                            database_name,
+                            incremental_timestamp,
+                            rename_suffix=None,
+                        )
+                    else:
+                        incremental_restore_command = _build_partition_restore_command(
+                            incremental_label,
+                            repo_name,
+                            table,
+                            partitions,
+                            database_name,
+                            incremental_timestamp,
+                            rename_suffix=rename_suffix,
+                        )
+
+                    incremental_result = execute_restore(
+                        db,
+                        incremental_restore_command,
+                        incremental_label,
+                        "incremental",
+                        repo_name,
+                        database_name,
+                        scope="restore",
+                        ops_database=ops_database,
+                    )
+
+                    if not incremental_result["success"]:
+                        return {
+                            "success": False,
+                            "error_message": f"Incremental restore failed for {table}: {incremental_result['error_message']}",
+                        }
+
+                logger.success("Incremental restore completed successfully")
 
         logger.info("")
         logger.info("Step 3: Performing atomic rename...")
@@ -557,6 +684,50 @@ def _build_restore_command_without_rename(
     PROPERTIES ("backup_timestamp" = "{backup_timestamp}")"""
 
 
+def _build_partition_restore_command(
+    backup_label: str,
+    repo_name: str,
+    table: str,
+    partitions: list[str],
+    database: str,
+    backup_timestamp: str,
+    rename_suffix: str | None = None,
+) -> str:
+    """Build partition-level restore command with optional AS clause.
+
+    Args:
+        backup_label: Backup snapshot label
+        repo_name: Repository name
+        table: Table name in format 'database.table'
+        partitions: List of partition names to restore
+        database: Database name
+        backup_timestamp: Backup timestamp
+        rename_suffix: Optional suffix for AS clause (e.g., '_restored')
+
+    Returns:
+        SQL RESTORE command string
+    """
+    _, table_name = table.split(".", 1)
+
+    # Build partition list
+    partition_list = ", ".join([utils.quote_identifier(p) for p in partitions])
+
+    # Build table clause
+    if rename_suffix:
+        # Table only in incremental: use AS clause
+        temp_table_name = f"{table_name}{rename_suffix}"
+        table_clause = f"TABLE {utils.quote_identifier(table_name)} PARTITION ({partition_list}) AS {utils.quote_identifier(temp_table_name)}"
+    else:
+        # Table in base: target the _restored table directly (no AS)
+        table_clause = f"TABLE {utils.quote_identifier(table_name)} PARTITION ({partition_list})"
+
+    return f"""RESTORE SNAPSHOT {utils.quote_identifier(backup_label)}
+    FROM {utils.quote_identifier(repo_name)}
+    DATABASE {utils.quote_identifier(database)}
+    ON ({table_clause})
+    PROPERTIES ("backup_timestamp" = "{backup_timestamp}")"""
+
+
 def _generate_timestamped_backup_name(table_name: str) -> str:
     """Generate a timestamped backup table name.
 

starrocks_br/schema.py

@@ -1,71 +1,131 @@
+# 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 . import logger
 
 
-def initialize_ops_schema(db) -> None:
+def initialize_ops_schema(
+    db, ops_database: str = "ops", table_inventory_entries: list[tuple[str, str, str]] | None = None
+) -> None:
     """Initialize the ops database and all required control tables.
 
-    Creates empty ops tables. Does NOT populate with sample data.
-    Users must manually insert their table inventory records.
+    Creates empty ops tables. Optionally populates table_inventory from entries.
+
+    Args:
+        db: Database connection
+        ops_database: Name of the ops database (defaults to "ops")
+        table_inventory_entries: Optional list of (group, database, table) tuples to bootstrap
     """
 
-    logger.info("Creating ops database...")
-    db.execute("CREATE DATABASE IF NOT EXISTS ops")
-    logger.success("ops database created")
+    logger.info(f"Creating {ops_database} database...")
+    db.execute(f"CREATE DATABASE IF NOT EXISTS {ops_database}")
+    logger.success(f"{ops_database} database created")
+
+    logger.info(f"Creating {ops_database}.table_inventory...")
+    db.execute(get_table_inventory_schema(ops_database=ops_database))
+    logger.success(f"{ops_database}.table_inventory created")
 
-    logger.info("Creating ops.table_inventory...")
-    db.execute(get_table_inventory_schema())
-    logger.success("ops.table_inventory created")
+    logger.info(f"Creating {ops_database}.backup_history...")
+    db.execute(get_backup_history_schema(ops_database=ops_database))
+    logger.success(f"{ops_database}.backup_history created")
 
-    logger.info("Creating ops.backup_history...")
-    db.execute(get_backup_history_schema())
-    logger.success("ops.backup_history created")
+    logger.info(f"Creating {ops_database}.restore_history...")
+    db.execute(get_restore_history_schema(ops_database=ops_database))
+    logger.success(f"{ops_database}.restore_history created")
 
-    logger.info("Creating ops.restore_history...")
-    db.execute(get_restore_history_schema())
-    logger.success("ops.restore_history created")
+    logger.info(f"Creating {ops_database}.run_status...")
+    db.execute(get_run_status_schema(ops_database=ops_database))
+    logger.success(f"{ops_database}.run_status created")
 
-    logger.info("Creating ops.run_status...")
-    db.execute(get_run_status_schema())
-    logger.success("ops.run_status created")
+    logger.info(f"Creating {ops_database}.backup_partitions...")
+    db.execute(get_backup_partitions_schema(ops_database=ops_database))
+    logger.success(f"{ops_database}.backup_partitions created")
+
+    if table_inventory_entries:
+        logger.info(f"Bootstrapping {ops_database}.table_inventory from configuration...")
+        bootstrap_table_inventory(db, table_inventory_entries, ops_database=ops_database)
+        logger.success(
+            f"{ops_database}.table_inventory bootstrapped with {len(table_inventory_entries)} entries"
+        )
 
-    logger.info("Creating ops.backup_partitions...")
-    db.execute(get_backup_partitions_schema())
-    logger.success("ops.backup_partitions created")
     logger.info("")
     logger.success("Schema initialized successfully!")
 
 
-def ensure_ops_schema(db) -> bool:
+def ensure_ops_schema(db, ops_database: str = "ops") -> bool:
     """Ensure ops schema exists, creating it if necessary.
 
     Returns True if schema was created, False if it already existed.
     This is called automatically before backup/restore operations.
     """
     try:
-        result = db.query("SHOW DATABASES LIKE 'ops'")
+        result = db.query(f"SHOW DATABASES LIKE '{ops_database}'")
 
         if not result:
-            initialize_ops_schema(db)
+            initialize_ops_schema(db, ops_database=ops_database)
             return True
 
-        db.execute("USE ops")
+        db.execute(f"USE {ops_database}")
         tables_result = db.query("SHOW TABLES")
 
         if not tables_result or len(tables_result) < 5:
-            initialize_ops_schema(db)
+            initialize_ops_schema(db, ops_database=ops_database)
             return True
 
         return False
 
     except Exception:
-        initialize_ops_schema(db)
+        initialize_ops_schema(db, ops_database=ops_database)
         return True
 
 
-def get_table_inventory_schema() -> str:
+def bootstrap_table_inventory(
+    db, entries: list[tuple[str, str, str]], ops_database: str = "ops"
+) -> None:
+    """Bootstrap table_inventory table with entries from configuration.
+
+    Args:
+        db: Database connection
+        entries: List of (group, database, table) tuples
+        ops_database: Name of the ops database (defaults to "ops")
+    """
+    if not entries:
+        return
+
+    unique_databases = {database for _, database, _ in entries}
+
+    for database_name in unique_databases:
+        result = db.query(f"SHOW DATABASES LIKE '{database_name}'")
+        if not result:
+            logger.warning(
+                f"Database '{database_name}' does not exist. "
+                f"Table inventory entries will be created, but backups will fail until the database is created."
+            )
+
+    for group, database, table in entries:
+        sql = f"""
+            INSERT INTO {ops_database}.table_inventory
+            (inventory_group, database_name, table_name)
+            VALUES ('{group}', '{database}', '{table}')
+        """
+        db.execute(sql)
+
+
+def get_table_inventory_schema(ops_database: str = "ops") -> str:
     """Get CREATE TABLE statement for table_inventory."""
-    return """
-    CREATE TABLE IF NOT EXISTS ops.table_inventory (
+    return f"""
+    CREATE TABLE IF NOT EXISTS {ops_database}.table_inventory (
         inventory_group STRING NOT NULL COMMENT "Group name for a set of tables",
         database_name STRING NOT NULL COMMENT "Database name",
         table_name STRING NOT NULL COMMENT "Table name, or '*' for all tables in database",
@@ -78,10 +138,10 @@ def get_table_inventory_schema() -> str:
     """
 
 
-def get_backup_history_schema() -> str:
+def get_backup_history_schema(ops_database: str = "ops") -> str:
     """Get CREATE TABLE statement for backup_history."""
-    return """
-    CREATE TABLE IF NOT EXISTS ops.backup_history (
+    return f"""
+    CREATE TABLE IF NOT EXISTS {ops_database}.backup_history (
         label STRING NOT NULL COMMENT "Unique backup snapshot label",
         backup_type STRING NOT NULL COMMENT "Type of backup: full or incremental",
         status STRING NOT NULL COMMENT "Final backup status: FINISHED, FAILED, CANCELLED, TIMEOUT",
@@ -96,10 +156,10 @@ def get_backup_history_schema() -> str:
     """
 
 
-def get_restore_history_schema() -> str:
+def get_restore_history_schema(ops_database: str = "ops") -> str:
     """Get CREATE TABLE statement for restore_history."""
-    return """
-    CREATE TABLE IF NOT EXISTS ops.restore_history (
+    return f"""
+    CREATE TABLE IF NOT EXISTS {ops_database}.restore_history (
         job_id STRING NOT NULL COMMENT "Unique restore job identifier",
         backup_label STRING NOT NULL COMMENT "Source backup snapshot label",
         restore_type STRING NOT NULL COMMENT "Type of restore: partition, table, or database",
@@ -116,10 +176,10 @@ def get_restore_history_schema() -> str:
     """
 
 
-def get_run_status_schema() -> str:
+def get_run_status_schema(ops_database: str = "ops") -> str:
     """Get CREATE TABLE statement for run_status."""
-    return """
-    CREATE TABLE IF NOT EXISTS ops.run_status (
+    return f"""
+    CREATE TABLE IF NOT EXISTS {ops_database}.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, FINISHED, FAILED, or CANCELLED",
@@ -131,12 +191,12 @@ def get_run_status_schema() -> str:
     """
 
 
-def get_backup_partitions_schema() -> str:
+def get_backup_partitions_schema(ops_database: str = "ops") -> str:
     """Get CREATE TABLE statement for backup_partitions."""
-    return """
-    CREATE TABLE IF NOT EXISTS ops.backup_partitions (
+    return f"""
+    CREATE TABLE IF NOT EXISTS {ops_database}.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.",
+        label STRING NOT NULL COMMENT "The backup label this partition belongs to. FK to {ops_database}.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.",

starrocks_br/timezone.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 datetime
 from zoneinfo import ZoneInfo
 

starrocks_br/utils.py

@@ -1,3 +1,18 @@
+# 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.
+
+
 def quote_identifier(identifier):
     """
     Quote a SQL identifier (database, table, or column name) with backticks.