diff --git a/weaviate/backup/backup.py b/weaviate/backup/backup.py index acfb91db8..e079d0532 100644 --- a/weaviate/backup/backup.py +++ b/weaviate/backup/backup.py @@ -4,7 +4,7 @@ from enum import Enum from time import sleep -from typing import Optional, Union, List, Tuple, Dict, Any, cast +from typing import Literal, Optional, Union, List, Tuple, Dict, Any, cast from pydantic import BaseModel, Field @@ -288,6 +288,8 @@ async def restore( backend: BackupStorage, include_collections: Union[List[str], str, None] = None, exclude_collections: Union[List[str], str, None] = None, + roles_restore: Optional[Literal["noRestore", "all"]] = None, + user_restore: Optional[Literal["noRestore", "all"]] = None, wait_for_completion: bool = False, config: Optional[BackupConfigRestore] = None, backup_location: Optional[BackupLocationType] = None, @@ -345,16 +347,14 @@ async def restore( "exclude": exclude_collections, } + configPayload = {} + if config is not None: if self._connection._weaviate_version.is_lower_than(1, 25, 0): raise WeaviateUnsupportedFeatureError( "BackupConfigRestore", str(self._connection._weaviate_version), "1.25.0" ) - if not isinstance(config, BackupConfigRestore): - raise WeaviateInvalidInputError( - f"Expected 'config' to be of type 'BackupConfigRestore', but got {type(config)}." - ) - payload["config"] = config._to_dict() + configPayload = config._to_dict() if backup_location is not None: if self._connection._weaviate_version.is_lower_than(1, 27, 2): @@ -366,7 +366,16 @@ async def restore( if "config" not in payload: payload["config"] = {} - payload["config"].update(backup_location._to_dict()) + configPayload.update(backup_location._to_dict()) + + if roles_restore is not None: + configPayload["rolesOptions"] = roles_restore + + if user_restore is not None: + configPayload["usersOptions"] = user_restore + + if len(configPayload) > 0: + payload["config"] = configPayload path = f"/backups/{backend.value}/{backup_id}/restore" response = await self._connection.post( diff --git a/weaviate/backup/sync.pyi b/weaviate/backup/sync.pyi index 8858aaf0b..5bcfbef76 100644 --- a/weaviate/backup/sync.pyi +++ b/weaviate/backup/sync.pyi @@ -1,4 +1,4 @@ -from typing import Optional, Union, List +from typing import Literal, Optional, Union, List from weaviate.backup.backup import ( BackupStorage, @@ -47,6 +47,8 @@ class _Backup: wait_for_completion: bool = False, config: Optional[BackupConfigRestore] = None, backup_location: Optional[BackupLocationType] = None, + roles_restore: Optional[Literal["noRestore", "all"]] = None, + user_restore: Optional[Literal["noRestore", "all"]] = None, ) -> BackupReturn: ... def get_restore_status( self, diff --git a/weaviate/collections/backups/backups.py b/weaviate/collections/backups/backups.py index 1c03b0ef4..8903b271d 100644 --- a/weaviate/collections/backups/backups.py +++ b/weaviate/collections/backups/backups.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Literal, Optional from weaviate.backup.backup import ( BackupConfigCreate, @@ -71,6 +71,8 @@ async def restore( self, backup_id: str, backend: BackupStorage, + roles_restore: Optional[Literal["noRestore", "all"]] = None, + user_restore: Optional[Literal["noRestore", "all"]] = None, wait_for_completion: bool = False, config: Optional[BackupConfigRestore] = None, backup_location: Optional[BackupLocationType] = None, @@ -103,7 +105,15 @@ async def restore( If the backup failed. """ restore = await self._backup.restore( - backup_id, backend, [self._name], None, wait_for_completion, config, backup_location + backup_id, + backend, + [self._name], + None, + roles_restore, + user_restore, + wait_for_completion, + config, + backup_location, ) return BackupStatusReturn( error=restore.error, status=restore.status, path=restore.path, id=backup_id diff --git a/weaviate/collections/backups/sync.pyi b/weaviate/collections/backups/sync.pyi index 3844946b2..06d6fe6b7 100644 --- a/weaviate/collections/backups/sync.pyi +++ b/weaviate/collections/backups/sync.pyi @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Literal, Optional from weaviate.backup.backup import ( BackupConfigCreate, BackupConfigRestore, @@ -19,6 +19,8 @@ class _CollectionBackup(_CollectionBackupBase): self, backup_id: str, backend: BackupStorage, + roles_restore: Optional[Literal["noRestore", "all"]] = None, + user_restore: Optional[Literal["noRestore", "all"]] = None, wait_for_completion: bool = False, config: Optional[BackupConfigRestore] = None, ) -> BackupStatusReturn: ...