diff --git a/ch_backup/backup/layout.py b/ch_backup/backup/layout.py index 42ed1bb9..f4753c4b 100644 --- a/ch_backup/backup/layout.py +++ b/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, diff --git a/ch_backup/backup/metadata/backup_metadata.py b/ch_backup/backup/metadata/backup_metadata.py index ced9c6b0..79b50427 100644 --- a/ch_backup/backup/metadata/backup_metadata.py +++ b/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]: diff --git a/ch_backup/backup/sources.py b/ch_backup/backup/sources.py index 2bac50c3..5fb50f98 100644 --- a/ch_backup/backup/sources.py +++ b/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,6 +33,7 @@ def for_backup( schema: bool = False, udf: bool = False, named_collections: bool = False, + workload_entities: bool = False, schema_only: bool = False, ) -> "BackupSources": """ @@ -38,10 +41,10 @@ def 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: diff --git a/ch_backup/ch_backup.py b/ch_backup/ch_backup.py index b481c9d9..a9f83c38 100644 --- a/ch_backup/ch_backup.py +++ b/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) diff --git a/ch_backup/clickhouse/control.py b/ch_backup/clickhouse/control.py index 88d3ca39..e28a79d5 100644 --- a/ch_backup/clickhouse/control.py +++ b/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: diff --git a/ch_backup/config.py b/ch_backup/config.py index 248803d2..88318bf5 100644 --- a/ch_backup/config.py +++ b/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, diff --git a/ch_backup/logic/workload_entities.py b/ch_backup/logic/workload_entities.py new file mode 100644 index 00000000..ad9c4135 --- /dev/null +++ b/ch_backup/logic/workload_entities.py @@ -0,0 +1,281 @@ +""" +Clickhouse backup logic for workload entities (WORKLOADs and RESOURCEs) +""" + +import os +import posixpath +from dataclasses import dataclass, field +from enum import Enum +from typing import List + +from ch_backup import logging +from ch_backup.backup_context import BackupContext +from ch_backup.clickhouse.config import ClickhouseConfig +from ch_backup.clickhouse.encryption import ClickHouseEncryption +from ch_backup.logic.backup_manager import BackupManager +from ch_backup.util import ( + chown_dir_contents, + copy_directory_content, + ensure_owned_directory, + escape_metadata_file_name, + temp_directory, +) +from ch_backup.zookeeper.zookeeper import ZookeeperCTL + + +class WorkloadEntitiesBackup(BackupManager): + """ + Workload entities (WORKLOAD and RESOURCE) backup class + """ + + def backup(self, context: BackupContext) -> None: + """ + Backup workload entities. + """ + if not context.ch_ctl.ch_version_ge("24.11"): + # CREATE WORKLOAD and CREATE RESOURCE SQL syntax added in 24.11 + # https://clickhouse.com/docs/en/operations/workload-scheduling + logging.info( + "Workload entities are not supported for version less than 24.11" + ) + return + + we_config = WorkloadEntitiesStorageConfig.from_ch_config( + context.ch_ctl_conf, context.ch_config + ) + + workload_entities = context.ch_ctl.get_workload_entities_query() + + if len(workload_entities) == 0: + return + + user = context.ch_ctl_conf["user"] + group = context.ch_ctl_conf["group"] + tmp_path = context.ch_ctl_conf["tmp_path"] + + ensure_owned_directory(tmp_path, user, group) + + with temp_directory( + tmp_path, + context.backup_meta.name, + ) as backup_tmp_path: + for entity_name in workload_entities: + context.backup_meta.add_workload_entity(entity_name) + + if we_config.is_local_storage(): + copy_directory_content(we_config.storage_path, backup_tmp_path) + elif we_config.is_storage_zookeeper(): + self._copy_directory_content_from_zookeeper( + context.zk_ctl, + we_config.storage_path, + backup_tmp_path, + ) + + if we_config.is_encrypted(): + decryptor = ClickHouseEncryption(context.ch_ctl) + decryptor.decrypt_directory_content( + backup_tmp_path, + we_config.encryption_key_hex, + ) + + chown_dir_contents(user, group, backup_tmp_path) + + for entity_name in workload_entities: + local_path = os.path.join( + backup_tmp_path, f"{escape_metadata_file_name(entity_name)}.sql" + ) + + context.backup_layout.upload_workload_entity_ddl_from_file( + local_path, + context.backup_meta.name, + entity_name, + ) + + def restore(self, context: BackupContext) -> None: + """ + Restore workload entities. + """ + if not context.ch_ctl.ch_version_ge("24.11"): + # CREATE WORKLOAD and CREATE RESOURCE SQL syntax added in 24.11 + # https://clickhouse.com/docs/en/operations/workload-scheduling + logging.info( + "Workload entities are not supported for version less than 24.11" + ) + return + + we_list = self.get_workload_entities_list(context) + if not we_list: + return + + logging.info("Restoring workload entities: {}", " ,".join(we_list)) + + we_on_clickhouse_list = context.ch_ctl.get_workload_entities_query() + + for entity_name in we_list: + logging.debug("Restoring workload entity {}", entity_name) + + statement = context.backup_layout.get_workload_entity_create_statement( + context.backup_meta, entity_name + ) + + if entity_name in we_on_clickhouse_list: + we_on_clickhouse_statement = ( + context.backup_layout.get_local_workload_entity_create_statement(entity_name) + ) + if we_on_clickhouse_statement != statement: + # The entity already on ClickHouse is the one being dropped, + # so derive its type from its own create statement when + # available, falling back to the backup statement. + self._drop_workload_entity( + context, entity_name, we_on_clickhouse_statement or statement + ) + context.ch_ctl.restore_workload_entity(statement) + + if entity_name not in we_on_clickhouse_list: + context.ch_ctl.restore_workload_entity(statement) + + logging.debug("Workload entity {} restored", entity_name) + + logging.info("All workload entities restored") + + @staticmethod + def get_workload_entities_list(context: BackupContext) -> List[str]: + """ + Get workload entities list + """ + return context.backup_meta.get_workload_entities() + + @staticmethod + def _drop_workload_entity( + context: BackupContext, entity_name: str, create_statement: str + ) -> None: + """ + Drop a workload entity, choosing WORKLOAD or RESOURCE based on its + create statement (`CREATE WORKLOAD ...` / `CREATE RESOURCE ...`). + + Both patterns are handled explicitly so that an unknown or malformed + statement fails loudly instead of being silently misclassified as a + RESOURCE and issuing the wrong DROP. + """ + normalized_statement = create_statement.lstrip().upper() + + if normalized_statement.startswith("CREATE WORKLOAD"): + context.ch_ctl.drop_workload(entity_name) + elif normalized_statement.startswith("CREATE RESOURCE"): + context.ch_ctl.drop_resource(entity_name) + else: + logging.error( + "Unsupported workload entity create statement for {}: {}", + entity_name, + create_statement, + ) + raise ValueError( + f"Unsupported workload entity create statement for '{entity_name}': " + f"{create_statement!r}" + ) + + def _copy_directory_content_from_zookeeper( + self, + zk_ctl: ZookeeperCTL, + from_path_dir: str, + to_path_dir: str, + ) -> None: + """ + Copy all files from zookeeper directory to destination. + """ + if posixpath.isabs(from_path_dir): + from_path_dir = from_path_dir[1:] + + with zk_ctl.zk_client as client: + target_dir = posixpath.normpath( + posixpath.join(zk_ctl.zk_root_path, from_path_dir) + ) + children_names = client.get_children(path=target_dir) + + for child_name in children_names: + subpath_from = posixpath.join(target_dir, child_name) + child_data, _ = client.get(subpath_from) + subpath_to = os.path.join(to_path_dir, child_name) + if not os.path.exists(subpath_to): + with open(subpath_to, "xb") as f: + f.write(child_data) + + +@dataclass +class WorkloadEntitiesStorageConfig: + """ + Class representing workload entities storage config. + """ + + class StorageType(Enum): + """ + Workload entities storage type. + """ + + LOCAL = "local" + LOCAL_ENCRYPTED = "local_encrypted" + ZOOKEEPER = "zookeeper" + ZOOKEEPER_ENCRYPTED = "zookeeper_encrypted" + + storage_type: StorageType = field(default=StorageType.LOCAL) + storage_path: str = field(default="/") + encryption_key_hex: str = field(default="") + + def is_local_storage(self) -> bool: + """ + Determines if config using local filesystem for storage. + """ + return self.storage_type in ( + self.StorageType.LOCAL, + self.StorageType.LOCAL_ENCRYPTED, + ) + + def is_storage_zookeeper(self) -> bool: + """ + Determines if config using zookeeper for storage. + """ + return self.storage_type in ( + self.StorageType.ZOOKEEPER, + self.StorageType.ZOOKEEPER_ENCRYPTED, + ) + + def is_encrypted(self) -> bool: + """ + Determines if config using encryption. + """ + return self.storage_type in ( + self.StorageType.LOCAL_ENCRYPTED, + self.StorageType.ZOOKEEPER_ENCRYPTED, + ) + + @classmethod + def from_ch_config( + cls, ch_backup_config: dict, ch_config: ClickhouseConfig + ) -> "WorkloadEntitiesStorageConfig": + """ + Create WorkloadEntitiesStorageConfig from ClickhouseConfig. + """ + we_config = ch_config.config.get("workload_entity_storage") + if not we_config: + storage_path = ch_backup_config.get("workload_path") + + assert storage_path, "workload_path missing from ch-backup config" + + return WorkloadEntitiesStorageConfig(storage_path=storage_path) + + storage_type_from_config = we_config.get("type") + storage_path_from_config = we_config.get("path") + encryption_key_hex_from_config = we_config.get("key_hex") + + assert ( + storage_path_from_config + ), "path missing from workload_entity_storage config" + + storage_type = WorkloadEntitiesStorageConfig.StorageType.LOCAL + storage_path = storage_path_from_config + encryption_key_hex = encryption_key_hex_from_config + + if storage_type_from_config: + storage_type = cls.StorageType(storage_type_from_config) + + return cls(storage_type, storage_path, encryption_key_hex) diff --git a/tests/integration/features/workload_entities.feature b/tests/integration/features/workload_entities.feature new file mode 100644 index 00000000..cbaeb31c --- /dev/null +++ b/tests/integration/features/workload_entities.feature @@ -0,0 +1,120 @@ +Feature: Workload entities (WORKLOADs and RESOURCEs) support + + Background: + Given default configuration + And a working s3 + And a working zookeeper on zookeeper01 + And a working clickhouse on clickhouse01 + And a working clickhouse on clickhouse02 + + @require_version_24.11 + Scenario: Check RESOURCE restore + Given we have executed queries on clickhouse01 + """ + CREATE RESOURCE test_resource (WRITE DISK disk_s3, READ DISK disk_s3); + """ + When we create clickhouse01 clickhouse backup + When we restore clickhouse backup #0 to clickhouse02 + When we execute query on clickhouse02 + """ + SELECT name FROM system.resources WHERE name = 'test_resource' LIMIT 1; + """ + Then we get response + """ + test_resource + """ + + @require_version_24.11 + Scenario: Check WORKLOAD restore + Given we have executed queries on clickhouse01 + """ + CREATE RESOURCE test_io_write (WRITE DISK disk_s3); + CREATE RESOURCE test_io_read (READ DISK disk_s3); + CREATE WORKLOAD test_workload SETTINGS max_requests = 100; + """ + When we create clickhouse01 clickhouse backup + When we restore clickhouse backup #0 to clickhouse02 + When we execute query on clickhouse02 + """ + SELECT name FROM system.workloads WHERE name = 'test_workload' LIMIT 1; + """ + Then we get response + """ + test_workload + """ + + @require_version_24.11 + Scenario: Check workload entity restore with same name + Given we have executed queries on clickhouse01 + """ + CREATE RESOURCE test_resource (WRITE DISK disk_s3); + """ + Given we have executed queries on clickhouse02 + """ + CREATE RESOURCE test_resource (READ DISK disk_s3); + """ + When we create clickhouse01 clickhouse backup + When we restore clickhouse backup #0 to clickhouse02 + When we execute query on clickhouse02 + """ + SELECT create_query FROM system.resources WHERE name = 'test_resource' LIMIT 1; + """ + Then we get response contains + """ + WRITE DISK disk_s3 + """ + + @require_version_24.11 + Scenario: Check workload entities restore-schema + Given we have executed queries on clickhouse01 + """ + CREATE RESOURCE test_resource (WRITE DISK disk_s3); + CREATE WORKLOAD test_workload SETTINGS max_requests = 100; + """ + When we create clickhouse01 clickhouse backup + """ + schema_only: true + """ + When we restore clickhouse backup #0 to clickhouse02 + """ + schema_only: true + """ + When we execute query on clickhouse02 + """ + SELECT count() FROM ( + SELECT name FROM system.resources WHERE name = 'test_resource' + UNION ALL + SELECT name FROM system.workloads WHERE name = 'test_workload' + ); + """ + Then we get response + """ + 2 + """ + + @require_version_24.11 + Scenario: Check workload entities restore-schema with same name + Given we have executed queries on clickhouse01 + """ + CREATE RESOURCE test_resource (WRITE DISK disk_s3); + """ + Given we have executed queries on clickhouse02 + """ + CREATE RESOURCE test_resource (READ DISK disk_s3); + """ + When we create clickhouse01 clickhouse backup + """ + schema_only: true + """ + When we restore clickhouse backup #0 to clickhouse02 + """ + schema_only: true + """ + When we execute query on clickhouse02 + """ + SELECT create_query FROM system.resources WHERE name = 'test_resource' LIMIT 1; + """ + Then we get response contains + """ + WRITE DISK disk_s3 + """