Add workload scheduling backup#303
Open
techkuz wants to merge 3 commits into
Open
Conversation
Contributor
Reviewer's GuideAdds backup and restore support for ClickHouse workload scheduling entities (WORKLOADs and RESOURCEs), including metadata wiring, storage abstraction (local/ZooKeeper, optional encryption), ClickHouse control operations, CLI source toggles, and integration tests for various restore scenarios. Sequence diagram for workload entities restore flowsequenceDiagram
participant CHBackup as CHBackup
participant WorkloadEntitiesBackup as WorkloadEntitiesBackup
participant BackupMetadata as BackupMetadata
participant BackupLayout as BackupLayout
participant ClickHouseControl as ClickHouseControl
CHBackup->>WorkloadEntitiesBackup: restore(context)
WorkloadEntitiesBackup->>ClickHouseControl: ch_version_ge("24.11")
ClickHouseControl-->>WorkloadEntitiesBackup: bool
alt [version < 24.11]
WorkloadEntitiesBackup-->>CHBackup: return
else [version >= 24.11]
WorkloadEntitiesBackup->>WorkloadEntitiesBackup: get_workload_entities_list(context)
WorkloadEntitiesBackup->>BackupMetadata: get_workload_entities()
BackupMetadata-->>WorkloadEntitiesBackup: we_list
WorkloadEntitiesBackup->>ClickHouseControl: get_workload_entities_query()
ClickHouseControl-->>WorkloadEntitiesBackup: we_on_clickhouse_list
loop for each entity_name in we_list
WorkloadEntitiesBackup->>BackupLayout: get_workload_entity_create_statement(backup_meta, entity_name)
BackupLayout-->>WorkloadEntitiesBackup: statement
alt [entity_name in we_on_clickhouse_list]
WorkloadEntitiesBackup->>BackupLayout: get_local_workload_entity_create_statement(entity_name)
BackupLayout-->>WorkloadEntitiesBackup: we_on_clickhouse_statement
alt [we_on_clickhouse_statement != statement]
WorkloadEntitiesBackup->>WorkloadEntitiesBackup: _drop_workload_entity(context, entity_name, we_on_clickhouse_statement or statement)
alt [create_statement starts with CREATE WORKLOAD]
WorkloadEntitiesBackup->>ClickHouseControl: drop_workload(entity_name)
else [otherwise]
WorkloadEntitiesBackup->>ClickHouseControl: drop_resource(entity_name)
end
WorkloadEntitiesBackup->>ClickHouseControl: restore_workload_entity(statement)
end
else [entity_name not in we_on_clickhouse_list]
WorkloadEntitiesBackup->>ClickHouseControl: restore_workload_entity(statement)
end
end
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
techkuz
changed the title
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- In
WorkloadEntitiesStorageConfig.from_ch_config, whenworkload_entity_storageis present butpathis missing,storage_pathcan becomeNonewithout any validation; consider asserting or raising a clear error ifpath(and when required,key_hex) are absent to avoid hard-to-debug runtime issues. - In
ClickhouseCTL.drop_workload_entity, you currently fall back to dropping a RESOURCE on anyExceptionfrom the WORKLOAD drop; it would be safer to restrict the fallback to expected errors (e.g. entity-not-found) rather than swallowing all failures, so that real errors (permissions, syntax, connectivity) are not masked.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `WorkloadEntitiesStorageConfig.from_ch_config`, when `workload_entity_storage` is present but `path` is missing, `storage_path` can become `None` without any validation; consider asserting or raising a clear error if `path` (and when required, `key_hex`) are absent to avoid hard-to-debug runtime issues.
- In `ClickhouseCTL.drop_workload_entity`, you currently fall back to dropping a RESOURCE on any `Exception` from the WORKLOAD drop; it would be safer to restrict the fallback to expected errors (e.g. entity-not-found) rather than swallowing all failures, so that real errors (permissions, syntax, connectivity) are not masked.
## Individual Comments
### Comment 1
<location path="ch_backup/backup/layout.py" line_range="347-354" />
<code_context>
+ 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
+
</code_context>
<issue_to_address>
**suggestion:** Error message mentions 'async upload' although the call is synchronous, which can mislead debugging.
Since `upload_file` is used synchronously, this message is misleading for anyone reading logs or debugging. Please update it to describe the actual failure (for example, `Failed to upload workload entity DDL to {remote_path}`) without implying async behavior.
```suggestion
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 upload workload entity DDL to {remote_path}"
raise StorageError(msg) from e
```
</issue_to_address>
### Comment 2
<location path="ch_backup/clickhouse/control.py" line_range="1006-1008" />
<code_context>
+ 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(
</code_context>
<issue_to_address>
**issue (bug_risk):** Catching a broad Exception when dropping workload entities may hide real errors.
Because any `DROP WORKLOAD` failure is treated as “must be a RESOURCE”, unrelated errors (permissions, syntax, connection issues, etc.) get silently misclassified and retried. If ClickHouse provides a specific code or message for “WORKLOAD not found”, please check for that explicitly before falling back to `DROP RESOURCE`, and re-raise all other errors.
</issue_to_address>
### Comment 3
<location path="ch_backup/logic/workload_entities.py" line_range="233-234" />
<code_context>
+ 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")
+
+ storage_type = WorkloadEntitiesStorageConfig.StorageType.LOCAL
</code_context>
<issue_to_address>
**issue:** Storage path from `workload_entity_storage` config is not validated and may be missing.
If `workload_entity_storage` exists but `path` is missing or empty, `storage_path_from_config` becomes `None` and is passed into `WorkloadEntitiesStorageConfig`. Subsequent operations that expect a valid directory (e.g., `copy_directory_content`, ZooKeeper reads) will then fail opaquely. Consider validating `storage_path_from_config` here and either raising a clear error or falling back to `workload_path` when it is not set.
</issue_to_address>
### Comment 4
<location path="ch_backup/logic/workload_entities.py" line_range="121-127" />
<code_context>
+ )
+
+ if entity_name in we_on_clickhouse_list:
+ we_on_clickhouse_statement = (
+ context.backup_layout.get_local_we_create_statement(entity_name)
+ )
+ if we_on_clickhouse_statement != statement:
+ context.ch_ctl.drop_workload_entity(entity_name)
</code_context>
<issue_to_address>
**suggestion:** Handling of a missing local workload entity statement relies on implicit `None != statement` comparison.
If `get_local_we_create_statement` returns `None` (e.g., missing file), this condition treats it as a mismatch and always drops/recreates the entity. If that behavior is desired, consider handling the `None` case explicitly (e.g., log a specific message, then drop/restore) so the intent and expectations around missing local files are clearer than relying on `None != statement`.
```suggestion
if entity_name in we_on_clickhouse_list:
we_on_clickhouse_statement = (
context.backup_layout.get_local_we_create_statement(entity_name)
)
if we_on_clickhouse_statement is None:
logging.warning(
"Local workload entity definition for %s is missing; "
"dropping and restoring %s from backup",
entity_name,
entity_name,
)
context.ch_ctl.drop_workload_entity(entity_name)
context.ch_ctl.restore_workload_entity(statement)
elif we_on_clickhouse_statement != statement:
logging.info(
"Local workload entity definition for %s differs from backup; "
"recreating entity from backup definition",
entity_name,
)
context.ch_ctl.drop_workload_entity(entity_name)
context.ch_ctl.restore_workload_entity(statement)
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+121
to
+127
| if entity_name in we_on_clickhouse_list: | ||
| we_on_clickhouse_statement = ( | ||
| context.backup_layout.get_local_we_create_statement(entity_name) | ||
| ) | ||
| if we_on_clickhouse_statement != statement: | ||
| context.ch_ctl.drop_workload_entity(entity_name) | ||
| context.ch_ctl.restore_workload_entity(statement) |
Contributor
There was a problem hiding this comment.
suggestion: Handling of a missing local workload entity statement relies on implicit None != statement comparison.
If get_local_we_create_statement returns None (e.g., missing file), this condition treats it as a mismatch and always drops/recreates the entity. If that behavior is desired, consider handling the None case explicitly (e.g., log a specific message, then drop/restore) so the intent and expectations around missing local files are clearer than relying on None != statement.
Suggested change
| if entity_name in we_on_clickhouse_list: | |
| we_on_clickhouse_statement = ( | |
| context.backup_layout.get_local_we_create_statement(entity_name) | |
| ) | |
| if we_on_clickhouse_statement != statement: | |
| context.ch_ctl.drop_workload_entity(entity_name) | |
| context.ch_ctl.restore_workload_entity(statement) | |
| if entity_name in we_on_clickhouse_list: | |
| we_on_clickhouse_statement = ( | |
| context.backup_layout.get_local_we_create_statement(entity_name) | |
| ) | |
| if we_on_clickhouse_statement is None: | |
| logging.warning( | |
| "Local workload entity definition for %s is missing; " | |
| "dropping and restoring %s from backup", | |
| entity_name, | |
| entity_name, | |
| ) | |
| context.ch_ctl.drop_workload_entity(entity_name) | |
| context.ch_ctl.restore_workload_entity(statement) | |
| elif we_on_clickhouse_statement != statement: | |
| logging.info( | |
| "Local workload entity definition for %s differs from backup; " | |
| "recreating entity from backup definition", | |
| entity_name, | |
| ) | |
| context.ch_ctl.drop_workload_entity(entity_name) | |
| context.ch_ctl.restore_workload_entity(statement) |
aalexfvk
requested changes
May 21, 2026
| remote_path = _named_collections_data_path(backup_meta.path, filename) | ||
| return self._storage_loader.download_data(remote_path, encryption=True) | ||
|
|
||
| def upload_workload_entity_ddl_from_file( |
Contributor
There was a problem hiding this comment.
Let's move this function above next to the other upload_ functions
| msg = f"Failed to create async upload of {remote_path}" | ||
| raise StorageError(msg) from e | ||
|
|
||
| def get_local_we_create_statement(self, entity_name: str) -> Optional[str]: |
Contributor
There was a problem hiding this comment.
we is misleading. Lets use workload_entity
| """ | ||
| try: | ||
| self._ch_client.query(DROP_WORKLOAD_SQL.format(entity_name=escape(entity_name))) | ||
| except Exception: |
Contributor
There was a problem hiding this comment.
Exception is too broad here.
Consider keeping entities (workload and resources) separate, or keeping the type together with each entity.
| ) | ||
| if we_on_clickhouse_statement != statement: | ||
| context.ch_ctl.drop_workload_entity(entity_name) | ||
| context.ch_ctl.restore_workload_entity(statement) |
Contributor
There was a problem hiding this comment.
А workload can reference some resource or another workload through the parent field, forming a hierarchy.
If you restore objects in an arbitrary order, a command like CREATE WORKLOAD child PARENT parent will fail if the parent has not been created yet.
Named collections do not have dependencies on each other, so this problem does not occur with them.
In short resources must be created before workloads, and workloads must be created in topological order (with parents before their children).
Contributor
|
Pls, resolve the Sourcery comments (either agree and fix or explain and deny as appropriate). |
techkuz
changed the title
Contributor
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- In
WorkloadEntitiesStorageConfig.from_ch_config, whenworkload_entity_storageis present butpathis omitted,storage_pathbecomesNonewithout validation; consider asserting or raising a clear error ifstorage_path_from_configis missing to avoid obscure failures at backup time. - The
_drop_workload_entityhelper inWorkloadEntitiesBackuptreats any non-CREATE WORKLOADstatement as aRESOURCE; it would be safer to explicitly check forCREATE RESOURCEand either raise or log an error for unexpected statements so misconfigured entities are not silently misclassified. - Some of the new logging calls (e.g. in
WorkloadEntitiesBackup.restoreandbackup/layout.py) use{}placeholders with theloggingmodule; if the project does not wrap logging to support this style, you may want to switch to%sformatting to avoid runtime formatting issues.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `WorkloadEntitiesStorageConfig.from_ch_config`, when `workload_entity_storage` is present but `path` is omitted, `storage_path` becomes `None` without validation; consider asserting or raising a clear error if `storage_path_from_config` is missing to avoid obscure failures at backup time.
- The `_drop_workload_entity` helper in `WorkloadEntitiesBackup` treats any non-`CREATE WORKLOAD` statement as a `RESOURCE`; it would be safer to explicitly check for `CREATE RESOURCE` and either raise or log an error for unexpected statements so misconfigured entities are not silently misclassified.
- Some of the new logging calls (e.g. in `WorkloadEntitiesBackup.restore` and `backup/layout.py`) use `{}` placeholders with the `logging` module; if the project does not wrap logging to support this style, you may want to switch to `%s` formatting to avoid runtime formatting issues.
## Individual Comments
### Comment 1
<location path="ch_backup/logic/workload_entities.py" line_range="148-159" />
<code_context>
+ 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 ...`).
+ """
+ if create_statement.lstrip().upper().startswith("CREATE WORKLOAD"):
+ context.ch_ctl.drop_workload(entity_name)
+ else:
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Treating every non-WORKLOAD entity as RESOURCE is brittle if new workload-related entity types appear.
Because anything not starting with `CREATE WORKLOAD` is treated as a RESOURCE, any new workload-like entity type (or malformed statement) would incorrectly call `drop_resource` without visibility. Consider explicitly handling both `CREATE WORKLOAD` and `CREATE RESOURCE`, and logging or raising if neither matches, so unknown or misconfigured statements fail loudly rather than issuing a wrong DROP.
```suggestion
@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 ...`).
For safety, explicitly handle both patterns and fail loudly on unknown
or malformed statements.
"""
normalized_stmt = create_statement.lstrip().upper()
if normalized_stmt.startswith("CREATE WORKLOAD"):
context.ch_ctl.drop_workload(entity_name)
elif normalized_stmt.startswith("CREATE RESOURCE"):
context.ch_ctl.drop_resource(entity_name)
else:
logging.error(
"Unsupported workload entity create statement for '%s': %r",
entity_name,
create_statement,
)
raise ValueError(
f"Unsupported workload entity create statement for '{entity_name}': "
f"{create_statement!r}"
)
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Author
|
@aalexfvk sorcery fixed |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
https://clickhouse.com/docs/operations/workload-scheduling
Summary by Sourcery
Add backup and restore support for ClickHouse workload scheduling entities (WORKLOADs and RESOURCEs), including storage configuration and metadata integration.
New Features:
Enhancements:
Tests: