starrocks-br 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

starrocks_br/concurrency.py

@@ -1,5 +1,6 @@
-from typing import Literal, List, Tuple
-from . import logger
+from typing import Literal
+
+from . import logger, utils
 
 
 def reserve_job_slot(db, scope: str, label: str) -> None:
@@ -9,23 +10,23 @@ def reserve_job_slot(db, scope: str, label: str) -> None:
     However, we implement self-healing logic to automatically clean up stale locks.
     """
     active_jobs = _get_active_jobs_for_scope(db, scope)
-    
+
     if not active_jobs:
         _insert_new_job(db, scope, label)
         return
-    
+
     _handle_active_job_conflicts(db, scope, active_jobs)
-    
+
     _insert_new_job(db, scope, label)
 
 
-def _get_active_jobs_for_scope(db, scope: str) -> List[Tuple[str, str, str]]:
+def _get_active_jobs_for_scope(db, scope: str) -> 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'")
     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]]) -> 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):
@@ -37,17 +38,17 @@ def _handle_active_job_conflicts(db, scope: str, active_jobs: List[Tuple[str, st
 
 def _can_heal_stale_job(scope: str, label: str, db) -> bool:
     """Check if a stale job can be healed (only for backup jobs)."""
-    if scope != 'backup':
+    if scope != "backup":
         return False
-    
+
     return _is_backup_job_stale(db, label)
 
 
-def _raise_concurrency_conflict(scope: str, active_jobs: List[Tuple[str, str, str]]) -> None:
+def _raise_concurrency_conflict(scope: str, active_jobs: list[tuple[str, str, str]]) -> None:
     """Raise a concurrency conflict error with helpful message."""
     active_job_strings = [f"{job[0]}:{job[1]}" for job in active_jobs]
     active_labels = [job[1] for job in active_jobs]
-    
+
     raise RuntimeError(
         f"Concurrency conflict: Another '{scope}' job is already ACTIVE: {', '.join(active_job_strings)}. "
         f"Wait for it to complete or cancel it via: UPDATE ops.run_status SET state='CANCELLED' "
@@ -57,48 +58,48 @@ def _raise_concurrency_conflict(scope: str, active_jobs: List[Tuple[str, str, st
 
 def _insert_new_job(db, scope: str, label: str) -> None:
     """Insert a new active job record."""
-    sql = (
-        "INSERT INTO ops.run_status (scope, label, state, started_at) "
-        "VALUES ('%s','%s','ACTIVE', NOW())" % (scope, label)
-    )
+    sql = f"""
+        INSERT INTO ops.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:
     """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)
-        
+
         for database_name in user_databases:
             job_status = _check_backup_job_in_database(db, database_name, label)
-            
+
             if job_status is None:
                 continue
-                
+
             if job_status == "active":
                 return False
             elif job_status == "stale":
                 return True
-                
+
         return True
-        
+
     except Exception as e:
         logger.error(f"Error checking backup job status: {e}")
         return False
 
 
-def _get_user_databases(db) -> List[str]:
+def _get_user_databases(db) -> list[str]:
     """Get list of user databases (excluding system databases)."""
-    SYSTEM_DATABASES = {'information_schema', 'mysql', 'sys', 'ops'}
-    
+    system_databases = {"information_schema", "mysql", "sys", "ops"}
+
     databases = db.query("SHOW DATABASES")
     return [
-        _extract_database_name(db_row) 
-        for db_row in databases 
-        if _extract_database_name(db_row) not in SYSTEM_DATABASES
+        _extract_database_name(db_row)
+        for db_row in databases
+        if _extract_database_name(db_row) not in system_databases
     ]
 
 
@@ -106,40 +107,40 @@ def _extract_database_name(db_row) -> str:
     """Extract database name from database query result."""
     if isinstance(db_row, (list, tuple)):
         return db_row[0]
-    return db_row.get('Database', '')
+    return db_row.get("Database", "")
 
 
 def _check_backup_job_in_database(db, database_name: str, label: str) -> str:
     """Check if backup job exists in specific database and return its status.
-    
+
     Returns:
         'active' if job is still running
-        'stale' if job is in terminal state  
+        'stale' if job is in terminal state
         None if job not found in this database
     """
     try:
-        show_backup_query = f"SHOW BACKUP FROM {database_name}"
+        show_backup_query = f"SHOW BACKUP FROM {utils.quote_identifier(database_name)}"
         backup_rows = db.query(show_backup_query)
-        
+
         if not backup_rows:
             return None
-            
+
         result = backup_rows[0]
         snapshot_name, state = _extract_backup_info(result)
-        
+
         if snapshot_name != label:
             return None
-            
+
         if state in ["FINISHED", "CANCELLED", "FAILED"]:
             return "stale"
         else:
             return "active"
-            
+
     except Exception:
         return None
 
 
-def _extract_backup_info(result) -> Tuple[str, str]:
+def _extract_backup_info(result) -> tuple[str, str]:
     """Extract snapshot name and state from SHOW BACKUP result."""
     if isinstance(result, dict):
         snapshot_name = result.get("SnapshotName", "")
@@ -147,31 +148,30 @@ def _extract_backup_info(result) -> Tuple[str, str]:
     else:
         snapshot_name = result[1] if len(result) > 1 else ""
         state = result[3] if len(result) > 3 else "UNKNOWN"
-    
+
     return snapshot_name, state
 
 
 def _cleanup_stale_job(db, scope: str, label: str) -> None:
     """Clean up a stale job by updating its state to CANCELLED."""
-    sql = (
-        "UPDATE ops.run_status SET state='CANCELLED', finished_at=NOW() "
-        "WHERE scope='%s' AND label='%s' AND state='ACTIVE'" % (scope, label)
-    )
+    sql = f"""
+        UPDATE ops.run_status
+        SET state='CANCELLED', finished_at=NOW()
+        WHERE scope={utils.quote_value(scope)} AND label={utils.quote_value(label)} AND state='ACTIVE'
+    """
     db.execute(sql)
 
 
 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"]
 ) -> None:
     """Complete job slot and persist final state.
 
     Simple approach: update the same row by scope/label.
     """
-    sql = (
-        "UPDATE ops.run_status SET state='%s', finished_at=NOW() WHERE scope='%s' AND label='%s'"
-        % (final_state, scope, label)
-    )
+    sql = f"""
+        UPDATE ops.run_status
+        SET state={utils.quote_value(final_state)}, finished_at=NOW()
+        WHERE scope={utils.quote_value(scope)} AND label={utils.quote_value(label)}
+    """
     db.execute(sql)

starrocks_br/config.py

@@ -1,45 +1,46 @@
+from typing import Any
+
 import yaml
-from typing import Any, Dict, Optional
 
 
-def load_config(config_path: str) -> Dict[str, Any]:
+def load_config(config_path: str) -> dict[str, Any]:
     """Load and parse YAML configuration file.
-    
+
     Args:
         config_path: Path to the YAML config file
-        
+
     Returns:
         Dictionary containing configuration
-        
+
     Raises:
         FileNotFoundError: If config file doesn't exist
         yaml.YAMLError: If config file is not valid YAML
     """
-    with open(config_path, 'r') as f:
+    with open(config_path) as f:
         config = yaml.safe_load(f)
-    
+
     if not isinstance(config, dict):
         raise ValueError("Config must be a dictionary")
-    
+
     return config
 
 
-def validate_config(config: Dict[str, Any]) -> None:
+def validate_config(config: dict[str, Any]) -> None:
     """Validate that config contains required fields.
-    
+
     Args:
         config: Configuration dictionary
-        
+
     Raises:
         ValueError: If required fields are missing
     """
-    required_fields = ['host', 'port', 'user', 'database', 'repository']
-    
+    required_fields = ["host", "port", "user", "database", "repository"]
+
     for field in required_fields:
         if field not in config:
             raise ValueError(f"Missing required config field: {field}")
 
-    _validate_tls_section(config.get('tls'))
+    _validate_tls_section(config.get("tls"))
 
 
 def _validate_tls_section(tls_config) -> None:
@@ -49,16 +50,23 @@ def _validate_tls_section(tls_config) -> None:
     if not isinstance(tls_config, dict):
         raise ValueError("TLS configuration must be a dictionary")
 
-    enabled = bool(tls_config.get('enabled', False))
+    enabled = bool(tls_config.get("enabled", False))
 
-    if enabled and not tls_config.get('ca_cert'):
+    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")
+    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,10 +1,11 @@
+from typing import Any, Optional
+
 import mysql.connector
-from typing import Any, Dict, List, Optional
 
 
 class StarRocksDB:
     """Database connection wrapper for StarRocks."""
-    
+
     def __init__(
         self,
         host: str,
@@ -12,10 +13,10 @@ class StarRocksDB:
         user: str,
         password: str,
         database: str,
-        tls_config: Optional[Dict[str, Any]] = None,
+        tls_config: Optional[dict[str, Any]] = None,
     ):
         """Initialize database connection.
-        
+
         Args:
             host: Database host
             port: Database port
@@ -31,35 +32,35 @@ class StarRocksDB:
         self._connection = None
         self.tls_config = tls_config or {}
         self._timezone: Optional[str] = None
-    
+
     def connect(self) -> None:
         """Establish database connection."""
-        conn_args: Dict[str, Any] = {
-            '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),
+        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'])
+            tls_versions = self.tls_config.get("tls_versions", ["TLSv1.2", "TLSv1.3"])
             if tls_versions:
-                ssl_args['tls_versions'] = 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):
+            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}. "
@@ -67,42 +68,42 @@ class StarRocksDB:
                     f"Alternatively, set 'enabled: false' in the tls section of your config file."
                 ) from e
             raise
-    
+
     def close(self) -> None:
         """Close database connection."""
         if self._connection:
             self._connection.close()
             self._connection = None
-    
+
     def execute(self, sql: str) -> None:
         """Execute a SQL statement that doesn't return results.
-        
+
         Args:
             sql: SQL statement to execute
         """
         if not self._connection:
             self.connect()
-        
+
         cursor = self._connection.cursor()
         try:
             cursor.execute(sql)
             self._connection.commit()
         finally:
             cursor.close()
-    
-    def query(self, sql: str, params: tuple = None) -> List[tuple]:
+
+    def query(self, sql: str, params: tuple = None) -> list[tuple]:
         """Execute a SQL query and return results.
-        
+
         Args:
             sql: SQL query to execute
             params: Optional tuple of parameters for parameterized queries
-            
+
         Returns:
             List of tuples containing query results
         """
         if not self._connection:
             self.connect()
-        
+
         cursor = self._connection.cursor()
         try:
             if params:
@@ -112,24 +113,24 @@ class StarRocksDB:
             return cursor.fetchall()
         finally:
             cursor.close()
-    
+
     def __enter__(self):
         """Context manager entry."""
         self.connect()
         return self
-    
+
     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.
@@ -138,7 +139,7 @@ class StarRocksDB:
             try:
                 query = "SHOW VARIABLES LIKE 'time_zone'"
                 rows = self.query(query)
-                
+
                 if not rows:
                     self._timezone = "UTC"
                 else:
@@ -149,6 +150,5 @@ class StarRocksDB:
                         self._timezone = row[1] if len(row) > 1 else "UTC"
             except Exception:
                 self._timezone = "UTC"
-        
-        return self._timezone
 
+        return self._timezone