Add workload scheduling backup

ch_backup/backup/layout.py

@@ -42,6 +42,7 @@ def __init__(self, config: Config) -> None:
         self._access_control_path = config["clickhouse"]["access_control_path"]
         self._metadata_path = config["clickhouse"]["metadata_path"]
         self._named_collections_path = config["clickhouse"]["named_collections_path"]
+        self._workload_path = config["clickhouse"]["workload_path"]
         enc_conf = config["encryption"]
         self._encryption_chunk_size = enc_conf["chunk_size"]
         self._encryption_metadata_size = get_encryption(
@@ -297,6 +298,24 @@ def upload_named_collections_ddl_from_file(
             msg = f"Failed to create async upload of {remote_path}"
             raise StorageError(msg) from e
 
+    def upload_workload_entity_ddl_from_file(
+        self, local_path: str, backup_name: str, entity_name: str
+    ) -> None:
+        """
+        Upload workload entity create statement file.
+        """
+        remote_path = _workload_entities_data_path(
+            self.get_backup_path(backup_name), entity_name
+        )
+        try:
+            logging.debug('Uploading workload entity create statement "{}"', entity_name)
+            self._storage_loader.upload_file(
+                local_path, remote_path=remote_path, encryption=True
+            )
+        except Exception as e:
+            msg = f"Failed to create async upload of {remote_path}"
+            raise StorageError(msg) from e
+
     def get_udf_create_statement(
         self, backup_meta: BackupMetadata, filename: str
     ) -> str:
@@ -324,7 +343,7 @@ def get_local_nc_create_statement(self, nc_name: str) -> Optional[str]:
                 str(e),
             )
             return None
-
+        
     def get_named_collection_create_statement(
         self, backup_meta: BackupMetadata, filename: str
     ) -> str:
@@ -334,6 +353,32 @@ def get_named_collection_create_statement(
         remote_path = _named_collections_data_path(backup_meta.path, filename)
         return self._storage_loader.download_data(remote_path, encryption=True)
 
+    def get_local_workload_entity_create_statement(self, entity_name: str) -> Optional[str]:
+        """
+        Read workload entity create statement from local file.
+        """
+        local_path = os.path.join(
+            self._workload_path, f"{escape_metadata_file_name(entity_name)}.sql"
+        )
+        try:
+            return Path(local_path).read_bytes().decode("utf-8")
+        except OSError as e:
+            logging.debug(
+                'Cannot load a create statement of the workload entity "{}": {}',
+                entity_name,
+                str(e),
+            )
+            return None
+
+    def get_workload_entity_create_statement(
+        self, backup_meta: BackupMetadata, filename: str
+    ) -> str:
+        """
+        Download workload entity create statement.
+        """
+        remote_path = _workload_entities_data_path(backup_meta.path, filename)
+        return self._storage_loader.download_data(remote_path, encryption=True)
+
     def get_backup_names(self) -> Sequence[str]:
         """
         Get names of existing backups.
@@ -880,6 +925,13 @@ def _named_collections_data_path(backup_path: str, nc_name: str) -> str:
     return os.path.join(backup_path, "named_collections", _quote(nc_name) + ".sql")
 
 
+def _workload_entities_data_path(backup_path: str, entity_name: str) -> str:
+    """
+    Return S3 path to workload entities.
+    """
+    return os.path.join(backup_path, "workload_entities", _quote(entity_name) + ".sql")
+
+
 def _part_path(
     backup_path: str,
     db_name: str,

ch_backup/backup/metadata/backup_metadata.py

@@ -71,6 +71,7 @@ def __init__(
         self._access_control = AccessControlMetadata()
         self._user_defined_functions: List[str] = []
         self._named_collections: List[str] = []
+        self._workload_entities: List[str] = []
 
     def __str__(self) -> str:
         return self.dump_json()
@@ -128,6 +129,7 @@ def dump(self, light: bool = False) -> dict:
             "access_controls": self._access_control.dump() if not light else {},
             "user_defined_functions": self._user_defined_functions if not light else [],
             "named_collections": self._named_collections if not light else [],
+            "workload_entities": self._workload_entities if not light else [],
             "cloud_storage": self.cloud_storage.dump(),
             "meta": {
                 "name": self.name,
@@ -208,6 +210,9 @@ def load(cls, data: dict) -> "BackupMetadata":
             backup._named_collections = data.get(
                 "named_collections", meta.get("named_collections", [])
             )
+            backup._workload_entities = data.get(
+                "workload_entities", meta.get("workload_entities", [])
+            )
 
             return backup
 
@@ -326,6 +331,19 @@ def get_named_collections(self) -> List[str]:
         """
         return self._named_collections
 
+    def add_workload_entity(self, entity_name: str) -> None:
+        """
+        Add workload entity in metadata.
+        """
+        assert entity_name not in self._workload_entities
+        self._workload_entities.append(entity_name)
+
+    def get_workload_entities(self) -> List[str]:
+        """
+        Get workload entities data.
+        """
+        return self._workload_entities
+
     def find_part(
         self, db_name: str, table_name: str, part_name: str
     ) -> Optional[PartMetadata]:

ch_backup/backup/sources.py

@@ -21,6 +21,8 @@ class BackupSources:
     udf: bool = True
     # Perform operation for named collections
     named_collections: bool = True
+    # Perform operation for workload entities (WORKLOADs and RESOURCEs)
+    workload_entities: bool = True
 
     @classmethod
     # pylint: disable=too-many-positional-arguments
@@ -31,17 +33,18 @@ def for_backup(
         schema: bool = False,
         udf: bool = False,
         named_collections: bool = False,
+        workload_entities: bool = False,
         schema_only: bool = False,
     ) -> "BackupSources":
         """
         Setting up sources ready for backup.
 
         @todo: `schema_only` is deprecated and will be replaced soon.
         """
-        if any([access, data, schema, udf, named_collections]):
+        if any([access, data, schema, udf, named_collections, workload_entities]):
             schema = data or schema
         else:
-            access, schema, udf, named_collections = True, True, True, True
+            access, schema, udf, named_collections, workload_entities = True, True, True, True, True
             data = not schema_only
 
         return cls(
@@ -50,6 +53,7 @@ def for_backup(
             schema=schema,
             udf=udf,
             named_collections=named_collections,
+            workload_entities=workload_entities,
         )
 
     @classmethod
@@ -61,6 +65,7 @@ def for_restore(
         schema: bool = False,
         udf: bool = False,
         named_collections: bool = False,
+        workload_entities: bool = False,
         schema_only: bool = False,
     ) -> "BackupSources":
         """
@@ -69,10 +74,10 @@ def for_restore(
         @todo: method will be merged with `for_backup` when they'll have similar logic
         @todo: `schema_only` is deprecated and will be replaced soon.
         """
-        if any([access, data, schema, udf, named_collections]):
+        if any([access, data, schema, udf, named_collections, workload_entities]):
             schema = data or schema
         else:
-            access, schema, udf, named_collections = False, True, True, True
+            access, schema, udf, named_collections, workload_entities = False, True, True, True, True
             data = not schema_only
 
         return cls(
@@ -81,6 +86,7 @@ def for_restore(
             schema=schema,
             udf=udf,
             named_collections=named_collections,
+            workload_entities=workload_entities,
         )
 
     def schemas_included(self) -> bool:

ch_backup/ch_backup.py

@@ -28,6 +28,7 @@
 from ch_backup.logic.access import AccessBackup
 from ch_backup.logic.database import DatabaseBackup
 from ch_backup.logic.named_collections import NamedCollectionsBackup
+from ch_backup.logic.workload_entities import WorkloadEntitiesBackup
 from ch_backup.logic.partial_restore import PartialRestoreFilter
 from ch_backup.logic.table import TableBackup
 from ch_backup.logic.udf import UDFBackup
@@ -61,6 +62,7 @@ def __init__(self, config: Config) -> None:
         self._table_backup_manager = TableBackup()
         self._udf_backup_manager = UDFBackup()
         self._nc_backup_manager = NamedCollectionsBackup()
+        self._we_backup_manager = WorkloadEntitiesBackup()
 
     @property
     def config(self) -> Config:
@@ -185,6 +187,8 @@ def backup(
                     self._udf_backup_manager.backup(self._context)
                 if sources.named_collections:
                     self._nc_backup_manager.backup(self._context)
+                if sources.workload_entities:
+                    self._we_backup_manager.backup(self._context)
                 if sources.schemas_included():
                     databases = self._database_backup_manager.backup(
                         self._context, databases
@@ -563,6 +567,10 @@ def _restore(
             # Restore named collections
             self._nc_backup_manager.restore(self._context)
 
+        if sources.workload_entities:
+            # Restore workload entities
+            self._we_backup_manager.restore(self._context)
+
         if sources.schemas_included():
             databases: Dict[str, Database] = {
                 db_name: self._context.backup_meta.get_database(db_name)

ch_backup/clickhouse/control.py

@@ -243,6 +243,18 @@
 """
 )
 
+DROP_WORKLOAD_SQL = strip_query(
+    """
+    DROP WORKLOAD `{entity_name}`
+"""
+)
+
+DROP_RESOURCE_SQL = strip_query(
+    """
+    DROP RESOURCE `{entity_name}`
+"""
+)
+
 TRUNCATE_TABLE_IF_EXISTS_SQL = strip_query(
     """
     TRUNCATE TABLE IF EXISTS `{db_name}`.`{table_name}`
@@ -446,6 +458,18 @@
 """
 )
 
+GET_WORKLOAD_ENTITIES_QUERY_SQL = strip_query(
+    """
+    SELECT name FROM (
+        SELECT name FROM system.workloads
+        UNION ALL
+        SELECT name FROM system.resources
+    )
+    ORDER BY name
+    FORMAT JSON
+"""
+)
+
 DECRYPT_AES_CTR_QUERY_SQL = strip_query(
     """
     SELECT decrypt('aes-{key_size}-ctr', unhex('{data_hex}'), unhex('{key_hex}'), unhex('{iv_hex}')) AS data
@@ -908,6 +932,12 @@ def restore_named_collection(self, nc_statement):
         """
         self._ch_client.query(nc_statement)
 
+    def restore_workload_entity(self, entity_statement):
+        """
+        Restore workload entity (WORKLOAD or RESOURCE).
+        """
+        self._ch_client.query(entity_statement)
+
     def create_table(self, table: Table) -> None:
         """
         Restore table.
@@ -968,6 +998,18 @@ def drop_named_collection(self, nc_name: str) -> None:
         """
         self._ch_client.query(DROP_NAMED_COLLECTION_SQL.format(nc_name=escape(nc_name)))
 
+    def drop_workload(self, entity_name: str) -> None:
+        """
+        Drop WORKLOAD entity.
+        """
+        self._ch_client.query(DROP_WORKLOAD_SQL.format(entity_name=escape(entity_name)))
+
+    def drop_resource(self, entity_name: str) -> None:
+        """
+        Drop RESOURCE entity.
+        """
+        self._ch_client.query(DROP_RESOURCE_SQL.format(entity_name=escape(entity_name)))
+
     def system_drop_replica(self, replica: str, zookeeper_path: str) -> None:
         """
         System drop replica query.
@@ -1205,6 +1247,13 @@ def get_named_collections_query(self) -> List[str]:
         resp = self._ch_client.query(GET_NAMED_COLLECTIONS_QUERY_SQL)
         return [row["name"] for row in resp.get("data", [])]
 
+    def get_workload_entities_query(self) -> List[str]:
+        """
+        Get workload entities (WORKLOADs and RESOURCEs) from system tables.
+        """
+        resp = self._ch_client.query(GET_WORKLOAD_ENTITIES_QUERY_SQL)
+        return [row["name"] for row in resp.get("data", [])]
+
     def decrypt_aes_ctr(
         self, data_hex: str, key_hex: str, key_size: int, iv_hex: str
     ) -> str:

ch_backup/config.py

@@ -25,6 +25,7 @@ def _as_seconds(t: str) -> int:
         "config_dir": "/etc/clickhouse-server/config.d/",
         "preprocessed_config_path": "/var/lib/clickhouse/preprocessed_configs/config.xml",
         "named_collections_path": "/var/lib/clickhouse/named_collections",
+        "workload_path": "/var/lib/clickhouse/workload",
         "host": socket.gethostname(),
         "protocol": "http",
         "port": None,