diff --git a/tests/Makefile b/tests/Makefile new file mode 100644 index 0000000000..720eb396e6 --- /dev/null +++ b/tests/Makefile @@ -0,0 +1,95 @@ +# tests/Makefile — lint helpers for Python files under tests/. +# +# All targets operate on the *PR-changed* set: files under tests/**/*.py that +# differ between $(BASE_REF) and HEAD (three-dot diff = since merge-base), +# matching the GitHub Actions "Python Lint (tests/)" workflow. +# +# Usage (run from this tests/ directory): +# make lint # ruff check (CI equivalent) +# make format-check # ruff format --check (CI equivalent) +# make ci # lint + format-check +# make lint-fix # ruff check --fix +# make format # ruff format +# +# Variables: +# BASE_REF diff base ref. Auto-detected from the current branch's open PR +# via `gh pr view`. If no PR exists, you MUST set it explicitly, +# e.g.: make ci BASE_REF=origin/master +# RUFF_VERSION ruff version (default: 0.15.11, pinned to .github/workflows/python-lint.yaml) +# +# Requires `uvx` (ships with `uv`). Ruff is fetched & cached on first run. +# Auto-detection requires `gh` authenticated (`gh auth status`). + +RUFF_VERSION ?= 0.15.11 +RUFF := uvx ruff@$(RUFF_VERSION) +REPO_ROOT := $(shell git rev-parse --show-toplevel) + +# Auto-detect the base ref from the current branch's open PR. +# Steps: +# 1. Read the PR URL (always lives under the base repo) and base branch via gh. +# Extract base repo as "/" from the URL path. +# 2. Find which local remote points to that base repo (must NOT assume "origin" +# — for forks, "upstream" is typically the base; for direct clones, "origin"). +# 3. Emit "/". +# Empty when no PR exists, no remote matches, or gh is unauthenticated — +# user must then set BASE_REF explicitly. +BASE_REF ?= $(shell \ + pr_url=$$(gh pr view --json url -q .url 2>/dev/null); \ + base_branch=$$(gh pr view --json baseRefName -q .baseRefName 2>/dev/null); \ + if [ -z "$$pr_url" ] || [ -z "$$base_branch" ]; then exit 0; fi; \ + base_repo=$$(echo "$$pr_url" | sed -E 's|^https?://[^/]+/([^/]+/[^/]+)/.*|\1|'); \ + if [ -z "$$base_repo" ]; then exit 0; fi; \ + remote=$$(git -C $(REPO_ROOT) remote -v | awk -v p="$$base_repo" 'index($$2, p) > 0 && $$3=="(fetch)" {print $$1; exit}'); \ + if [ -z "$$remote" ]; then exit 0; fi; \ + echo "$$remote/$$base_branch") + +CHANGED_FILES_CMD = git -C $(REPO_ROOT) diff --name-only --diff-filter=ACMR $(BASE_REF)...HEAD -- 'tests/**/*.py' + +.PHONY: help lint format-check ci lint-fix format _require-base-ref + +help: + @echo "tests/ Makefile — lint PR-changed *.py using ruff $(RUFF_VERSION)" + @echo + @echo "Targets:" + @echo " make lint ruff check (CI equivalent)" + @echo " make format-check ruff format --check (CI equivalent)" + @echo " make ci lint + format-check" + @echo " make lint-fix ruff check --fix" + @echo " make format ruff format" + @echo + @echo "BASE_REF: auto-detected from the current branch's open PR via 'gh pr view'." + @echo " When no PR exists, you must set it explicitly:" + @echo " make ci BASE_REF=origin/master" + @if [ -n "$(BASE_REF)" ]; then echo; echo "Currently detected: BASE_REF=$(BASE_REF)"; \ + else echo; echo "Currently: BASE_REF is unset (no open PR detected)."; fi + +_require-base-ref: + @if [ -z "$(BASE_REF)" ]; then \ + echo "Error: BASE_REF is not set."; \ + echo " No open PR detected for the current branch via 'gh pr view'."; \ + echo " Set it explicitly, e.g.:"; \ + echo " make $(or $(MAKECMDGOALS),ci) BASE_REF=origin/master"; \ + exit 1; \ + fi + +lint: _require-base-ref + @files=$$($(CHANGED_FILES_CMD)); \ + if [ -z "$$files" ]; then echo "No changed Python files under tests/ (base: $(BASE_REF))."; \ + else cd $(REPO_ROOT) && $(RUFF) check $$files; fi + +format-check: _require-base-ref + @files=$$($(CHANGED_FILES_CMD)); \ + if [ -z "$$files" ]; then echo "No changed Python files under tests/ (base: $(BASE_REF))."; \ + else cd $(REPO_ROOT) && $(RUFF) format --check --diff $$files; fi + +ci: lint format-check + +lint-fix: _require-base-ref + @files=$$($(CHANGED_FILES_CMD)); \ + if [ -z "$$files" ]; then echo "No changed Python files under tests/ (base: $(BASE_REF))."; \ + else cd $(REPO_ROOT) && $(RUFF) check --fix $$files; fi + +format: _require-base-ref + @files=$$($(CHANGED_FILES_CMD)); \ + if [ -z "$$files" ]; then echo "No changed Python files under tests/ (base: $(BASE_REF))."; \ + else cd $(REPO_ROOT) && $(RUFF) format $$files; fi diff --git a/tests/README.md b/tests/README.md index b5fd0cf25b..e2524b9f05 100644 --- a/tests/README.md +++ b/tests/README.md @@ -108,3 +108,38 @@ $ uv run ruff format --check . # format check only (CI-friendly) Rules enabled: `E`, `F`, `W`, `I`, `UP`. Target Python version: `3.10`. +#### Lint only PR-changed files (Python Lint CI parity) + +The GitHub Actions workflow `.github/workflows/python-lint.yaml` (job +**"Python Lint (tests/) / Ruff (changed files only)"**) runs `ruff check` +and `ruff format --check` against the *changed* `tests/**/*.py` set on +every PR. Running `uv run ruff check .` over the whole tree is too coarse +because the historical contents of many files predate the lint config and +will fail unrelated rules. + +To reproduce the CI step locally, use the `Makefile` shipped in this directory: + +```shell +$ cd tests/ +$ make ci # ruff check + format --check on PR-changed *.py (CI equivalent) +$ make lint-fix # ruff check --fix on PR-changed *.py +$ make format # ruff format on PR-changed *.py +$ make help # show all targets and the detected BASE_REF +``` + +`BASE_REF` is auto-detected from the current branch's open PR via `gh pr view`: +the PR URL is parsed to discover the base `/` and matched against +your local git remotes, producing e.g. `upstream/master`, `upstream/2.x`, or +`origin/main` for direct clones. + +If no open PR exists for the branch, you **must** set `BASE_REF` explicitly +(no guessing — a wrong base diffs against unrelated commits): + +```shell +$ make ci BASE_REF=upstream/master +``` + +Requires `uv` (provides `uvx`) and `gh` authenticated against GitHub +(`gh auth status`). The ruff version is pinned in the `Makefile` via +`RUFF_VERSION` to match the workflow. + diff --git a/tests/README_CN.md b/tests/README_CN.md index d896b7c65f..342056ce71 100644 --- a/tests/README_CN.md +++ b/tests/README_CN.md @@ -108,3 +108,37 @@ $ uv run ruff format --check . # 只检查不修改 (CI 友好) 启用的规则：`E`、`F`、`W`、`I`、`UP`；目标 Python 版本：`3.10`。 +#### 仅检查 PR 修改的文件 (与 Python Lint CI 一致) + +GitHub Actions workflow `.github/workflows/python-lint.yaml` +（job **"Python Lint (tests/) / Ruff (changed files only)"**） +只对 PR 修改的 `tests/**/*.py` 文件跑 `ruff check` 和 `ruff format --check`。 +直接对整个 `tests/` 目录执行 `uv run ruff check .` 太粗——很多历史文件早于 +当前 lint 配置，会因为无关规则失败。 + +要在本地复现 CI 这一步，使用本目录下的 `Makefile`： + +```shell +$ cd tests/ +$ make ci # 对 PR 修改文件跑 ruff check + format --check（与 CI 等价） +$ make lint-fix # 对 PR 修改文件跑 ruff check --fix +$ make format # 对 PR 修改文件跑 ruff format +$ make help # 显示所有 target 以及当前自动检测到的 BASE_REF +``` + +`BASE_REF` 通过 `gh pr view` 从当前分支的开放 PR 中自动检测： +解析 PR URL 得到 base `/`，再匹配本地 git remotes， +得到形如 `upstream/master`、`upstream/2.x` 或（直接 clone 主仓库时的） +`origin/main`。 + +如果当前分支没有开放 PR，**必须**显式设置 `BASE_REF` +（不再做猜测——base 选错会 diff 出无关 commit）： + +```shell +$ make ci BASE_REF=upstream/master +``` + +需要本地已安装 `uv`（提供 `uvx`）和已认证的 `gh` +（`gh auth status`）。Makefile 通过 `RUFF_VERSION` 锁定与 +workflow 一致的 ruff 版本。 + diff --git a/tests/python_client/base/client_v2_base.py b/tests/python_client/base/client_v2_base.py index d4336c7073..4402e3ed8d 100644 --- a/tests/python_client/base/client_v2_base.py +++ b/tests/python_client/base/client_v2_base.py @@ -1,22 +1,22 @@ import sys import time -from typing import Optional -from pymilvus import MilvusClient, DataType + +from pymilvus import DataType, MilvusClient sys.path.append("..") -from check.func_check import ResponseChecker -from utils.api_request import api_request -from utils.wrapper import trace -from utils.util_log import test_log as log -from common import common_func as cf, common_type as ct from base.client_base import Base +from check.func_check import ResponseChecker +from common import common_func as cf +from common import common_type as ct +from utils.api_request import api_request +from utils.util_log import test_log as log +from utils.wrapper import trace TIMEOUT = 120 INDEX_NAME = "" class TestMilvusClientV2Base(Base): - # milvus_client = None active_trace = False @@ -31,28 +31,52 @@ class TestMilvusClientV2Base(Base): self.async_milvus_client_wrap.init_async_client(**kwargs) def _client(self, active_trace=False, **kwargs): - """ return MilvusClient instance if connected successfully, otherwise return None""" + """return MilvusClient instance if connected successfully, otherwise return None""" if self.skip_connection: return None if cf.param_info.param_uri: uri = cf.param_info.param_uri else: uri = "http://" + cf.param_info.param_host + ":" + str(cf.param_info.param_port) - res, is_succ = self.init_milvus_client(uri=uri, token=cf.param_info.param_token, active_trace=active_trace, **kwargs) + res, is_succ = self.init_milvus_client( + uri=uri, token=cf.param_info.param_token, active_trace=active_trace, **kwargs + ) if is_succ: # self.milvus_client = res - log.info(f"server version: {res.get_server_version()}") + log.info(f"server version: {res.get_server_version(detail=True)}") return res - def init_milvus_client(self, uri, user="", password="", db_name="", token="", timeout=None, - check_task=None, check_items=None, active_trace=False, **kwargs): + def init_milvus_client( + self, + uri, + user="", + password="", + db_name="", + token="", + timeout=None, + check_task=None, + check_items=None, + active_trace=False, + **kwargs, + ): self.active_trace = active_trace func_name = sys._getframe().f_code.co_name res, is_succ = api_request([MilvusClient, uri, user, password, db_name, token, timeout], **kwargs) # self.milvus_client = res if is_succ else None - check_result = ResponseChecker(res, func_name, check_task, check_items, is_succ, - uri=uri, user=user, password=password, db_name=db_name, token=token, - timeout=timeout, **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + is_succ, + uri=uri, + user=user, + password=password, + db_name=db_name, + token=token, + timeout=timeout, + **kwargs, + ).run() return res, check_result @trace() @@ -63,78 +87,105 @@ class TestMilvusClientV2Base(Base): return res, check_result @trace() - def create_schema(self, client, timeout=None, check_task=None, - check_items=None, **kwargs): + def create_schema(self, client, timeout=None, check_task=None, check_items=None, **kwargs): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_schema], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result + @trace() - def create_struct_field_schema(self, client, check_task=None, - check_items=None, **kwargs): + def create_struct_field_schema(self, client, check_task=None, check_items=None, **kwargs): func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_struct_field_schema], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() - return res, check_result - + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() + return res, check_result @trace() def add_field(self, schema, field_name, datatype, check_task=None, check_items=None, **kwargs): # Set default parameters for specific field types - if datatype == DataType.VARCHAR and 'max_length' not in kwargs: - kwargs['max_length'] = ct.default_length + if datatype == DataType.VARCHAR and "max_length" not in kwargs: + kwargs["max_length"] = ct.default_length elif datatype == DataType.ARRAY: - if 'element_type' not in kwargs: - kwargs['element_type'] = DataType.INT64 - if 'max_capacity' not in kwargs: - kwargs['max_capacity'] = ct.default_max_capacity + if "element_type" not in kwargs: + kwargs["element_type"] = DataType.INT64 + if "max_capacity" not in kwargs: + kwargs["max_capacity"] = ct.default_max_capacity func_name = sys._getframe().f_code.co_name res, check = api_request([schema.add_field, field_name, datatype], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result - @trace() - def create_collection(self, client, collection_name, dimension=None, primary_field_name='id', - id_type='int', vector_field_name='vector', metric_type='COSINE', - auto_id=False, schema=None, index_params=None, timeout=None, force_teardown=True, - check_task=None, check_items=None, **kwargs): + def create_collection( + self, + client, + collection_name, + dimension=None, + primary_field_name="id", + id_type="int", + vector_field_name="vector", + metric_type="COSINE", + auto_id=False, + schema=None, + index_params=None, + timeout=None, + force_teardown=True, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout consistency_level = kwargs.get("consistency_level", "Strong") kwargs.update({"consistency_level": consistency_level}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.create_collection, collection_name, dimension, primary_field_name, - id_type, vector_field_name, metric_type, auto_id, timeout, schema, - index_params], **kwargs) + res, check = api_request( + [ + client.create_collection, + collection_name, + dimension, + primary_field_name, + id_type, + vector_field_name, + metric_type, + auto_id, + timeout, + schema, + index_params, + ], + **kwargs, + ) # Register for cleanup BEFORE assertion check, so collection is cleaned up even if check fails # (e.g., when check_task=err_res but API actually succeeds due to server bug) if check and force_teardown: self.tear_down_collection_names.append(collection_name) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, dimension=dimension, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + dimension=dimension, + **kwargs, + ).run() return res, check_result - def has_collection(self, client, collection_name, timeout=None, check_task=None, - check_items=None, **kwargs): + def has_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.has_collection, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -144,8 +195,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.insert, collection_name, data], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() @@ -155,9 +205,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.upsert, collection_name, data], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, data=data, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, data=data, **kwargs + ).run() return res, check_result @trace() @@ -167,58 +217,132 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.get_collection_stats, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() - def search(self, client, collection_name, data=None, limit=10, filter=None, output_fields=None, search_params=None, - timeout=None, check_task=None, check_items=None, **kwargs): + def search( + self, + client, + collection_name, + data=None, + limit=10, + filter=None, + output_fields=None, + search_params=None, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout # kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.search, collection_name, data, filter, limit, - output_fields, search_params], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, data=data, limit=limit, filter=filter, - output_fields=output_fields, search_params=search_params, - **kwargs).run() + res, check = api_request( + [client.search, collection_name, data, filter, limit, output_fields, search_params], **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + data=data, + limit=limit, + filter=filter, + output_fields=output_fields, + search_params=search_params, + **kwargs, + ).run() return res, check_result - @trace() - def search_iterator(self, client, collection_name, data, batch_size, limit=-1, filter=None, output_fields=None, - search_params=None, timeout=None, check_task=None, check_items=None, **kwargs): + def search_iterator( + self, + client, + collection_name, + data, + batch_size, + limit=-1, + filter=None, + output_fields=None, + search_params=None, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.search_iterator, collection_name, data, batch_size, filter, limit, - output_fields, search_params], **kwargs) - if any(k in kwargs for k in ['use_rbac_mul_db', 'use_mul_db']): - self.using_database(client, kwargs.get('another_db')) - if kwargs.get('use_alias', False) is True: + res, check = api_request( + [client.search_iterator, collection_name, data, batch_size, filter, limit, output_fields, search_params], + **kwargs, + ) + if any(k in kwargs for k in ["use_rbac_mul_db", "use_mul_db"]): + self.using_database(client, kwargs.get("another_db")) + if kwargs.get("use_alias", False) is True: alias = collection_name - self.alter_alias(client, kwargs.get('another_collection'), alias) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, data=data, batch_size=batch_size, filter=filter, - limit=limit, output_fields=output_fields, search_params=search_params, - **kwargs).run() + self.alter_alias(client, kwargs.get("another_collection"), alias) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + data=data, + batch_size=batch_size, + filter=filter, + limit=limit, + output_fields=output_fields, + search_params=search_params, + **kwargs, + ).run() return res, check_result @trace() - def hybrid_search(self, client, collection_name, reqs, ranker, limit=10, - output_fields=None, timeout=None, partition_names=None, - check_task=None, check_items=None, **kwargs): + def hybrid_search( + self, + client, + collection_name, + reqs, + ranker, + limit=10, + output_fields=None, + timeout=None, + partition_names=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout # kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.hybrid_search, collection_name, reqs, ranker, limit, - output_fields, timeout, partition_names], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, reqs=reqs, ranker=ranker, limit=limit, - output_fields=output_fields, timeout=timeout, partition_names=partition_names, **kwargs).run() + res, check = api_request( + [client.hybrid_search, collection_name, reqs, ranker, limit, output_fields, timeout, partition_names], + **kwargs, + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + reqs=reqs, + ranker=ranker, + limit=limit, + output_fields=output_fields, + timeout=timeout, + partition_names=partition_names, + **kwargs, + ).run() return res, check_result @trace() @@ -228,9 +352,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.query, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -240,23 +364,39 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.query_iterator, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() - def get(self, client, collection_name, ids, output_fields=None, - timeout=None, check_task=None, check_items=None, **kwargs): + def get( + self, + client, + collection_name, + ids, + output_fields=None, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.get, collection_name, ids, output_fields], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, ids=ids, - output_fields=output_fields, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + ids=ids, + output_fields=output_fields, + **kwargs, + ).run() return res, check_result # No client.num_entities method @@ -273,15 +413,25 @@ class TestMilvusClientV2Base(Base): # return res, check_result @trace() - def delete(self, client, collection_name, ids=None, timeout=None, filter=None, partition_name=None, - check_task=None, check_items=None, **kwargs): + def delete( + self, + client, + collection_name, + ids=None, + timeout=None, + filter=None, + partition_name=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout func_name = sys._getframe().f_code.co_name res, check = api_request([client.delete, collection_name, ids, timeout, filter, partition_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -291,9 +441,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.flush, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -303,9 +453,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_collection, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -315,17 +465,16 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_collections], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() def drop_collection(self, client, collection_name, check_task=None, check_items=None, **kwargs): func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_collection, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() if check_result is True and collection_name in self.tear_down_collection_names: self.tear_down_collection_names.remove(collection_name) return res, check_result @@ -334,30 +483,29 @@ class TestMilvusClientV2Base(Base): def list_indexes(self, client, collection_name, field_name=None, check_task=None, check_items=None, **kwargs): func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_indexes, collection_name, field_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() def describe_replica(self, client, collection_name, check_task=None, check_items=None, **kwargs): func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_replica, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() def get_load_state(self, client, collection_name, check_task=None, check_items=None, **kwargs): func_name = sys._getframe().f_code.co_name res, check = api_request([client.get_load_state, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result - @trace() def load_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs): timeout = TIMEOUT if timeout is None else timeout @@ -365,9 +513,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.load_collection, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -377,9 +525,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.refresh_load, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -389,9 +537,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.release_collection, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -401,49 +549,65 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.truncate_collection, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() - def list_persistent_segments(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs): + def list_persistent_segments( + self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_persistent_segments, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() - def load_partitions(self, client, collection_name, partition_names, timeout=None, check_task=None, check_items=None, **kwargs): + def load_partitions( + self, client, collection_name, partition_names, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.load_partitions, collection_name, partition_names], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - partition_names=partition_names, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + partition_names=partition_names, + **kwargs, + ).run() return res, check_result @trace() - def release_partitions(self, client, collection_name, partition_names, timeout=None, check_task=None, check_items=None, **kwargs): + def release_partitions( + self, client, collection_name, partition_names, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.release_partitions, collection_name, partition_names], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - partition_names=partition_names, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + partition_names=partition_names, + **kwargs, + ).run() return res, check_result @trace() @@ -453,36 +617,42 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.rename_collection, old_name, new_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - old_name=old_name, - new_name=new_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, old_name=old_name, new_name=new_name, **kwargs + ).run() self.tear_down_collection_names.append(new_name) return res, check_result @trace() - def create_database(self, client, db_name, properties: Optional[dict] = None, check_task=None, check_items=None, **kwargs): + def create_database( + self, client, db_name, properties: dict | None = None, check_task=None, check_items=None, **kwargs + ): func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_database, db_name, properties], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - db_name=db_name, properties=properties, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, db_name=db_name, properties=properties, **kwargs + ).run() return res, check_result @trace() - def create_partition(self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs): + def create_partition( + self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_partition, collection_name, partition_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - partition_name=partition_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + partition_name=partition_name, + **kwargs, + ).run() return res, check_result @trace() @@ -492,52 +662,72 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_partitions, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() - def drop_partition(self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs): + def drop_partition( + self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_partition, collection_name, partition_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - partition_name=partition_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + partition_name=partition_name, + **kwargs, + ).run() return res, check_result @trace() - def has_partition(self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs): + def has_partition( + self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.has_partition, collection_name, partition_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - partition_name=partition_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + partition_name=partition_name, + **kwargs, + ).run() return res, check_result @trace() - def get_partition_stats(self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs): + def get_partition_stats( + self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.get_partition_stats, collection_name, partition_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - partition_name=partition_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + partition_name=partition_name, + **kwargs, + ).run() return res, check_result @trace() @@ -545,51 +735,70 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.prepare_index_params], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def create_index(self, client, collection_name, index_params, timeout=None, check_task=None, check_items=None, **kwargs): + def create_index( + self, client, collection_name, index_params, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_index, collection_name, index_params], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - index_params=index_params, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + index_params=index_params, + **kwargs, + ).run() return res, check_result @trace() - def drop_index(self, client, collection_name, index_name, timeout=None, check_task=None, check_items=None, **kwargs): + def drop_index( + self, client, collection_name, index_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_index, collection_name, index_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - index_name=index_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + index_name=index_name, + **kwargs, + ).run() return res, check_result @trace() - def describe_index(self, client, collection_name, index_name, timeout=None, check_task=None, check_items=None, **kwargs): + def describe_index( + self, client, collection_name, index_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_index, collection_name, index_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - index_name=index_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + index_name=index_name, + **kwargs, + ).run() return res, check_result def wait_for_index_ready(self, client, collection_name, index_name, timeout=None, **kwargs): @@ -619,11 +828,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_alias, collection_name, alias], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - alias=alias, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, alias=alias, **kwargs + ).run() return res, check_result @trace() @@ -633,10 +840,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_alias, alias], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - alias=alias, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, alias=alias, **kwargs).run() return res, check_result @trace() @@ -646,11 +850,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.alter_alias, collection_name, alias], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - collection_name=collection_name, - alias=alias, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, alias=alias, **kwargs + ).run() return res, check_result @trace() @@ -660,10 +862,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_alias, alias], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, - alias=alias, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, alias=alias, **kwargs).run() return res, check_result @trace() @@ -673,9 +872,9 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_aliases, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() @@ -703,9 +902,9 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_user, user_name, password], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, user_name=user_name, - password=password, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, user_name=user_name, password=password, **kwargs + ).run() return res, check_result @trace() @@ -714,22 +913,42 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_user, user_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, user_name=user_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, user_name=user_name, **kwargs + ).run() return res, check_result @trace() - def update_password(self, client, user_name, old_password, new_password, reset_connection=False, - timeout=None, check_task=None, check_items=None, **kwargs): + def update_password( + self, + client, + user_name, + old_password, + new_password, + reset_connection=False, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.update_password, user_name, old_password, new_password, - reset_connection], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, user_name=user_name, old_password=old_password, - new_password=new_password, reset_connection=reset_connection, - **kwargs).run() + res, check = api_request( + [client.update_password, user_name, old_password, new_password, reset_connection], **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + user_name=user_name, + old_password=old_password, + new_password=new_password, + reset_connection=reset_connection, + **kwargs, + ).run() return res, check_result @trace() @@ -738,8 +957,7 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_users], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() @@ -748,8 +966,9 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_user, user_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, user_name=user_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, user_name=user_name, **kwargs + ).run() return res, check_result @trace() @@ -758,8 +977,9 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_role, role_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, role_name=role_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, role_name=role_name, **kwargs + ).run() return res, check_result @trace() @@ -768,8 +988,9 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_role, role_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, role_name=role_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, role_name=role_name, **kwargs + ).run() return res, check_result @trace() @@ -778,8 +999,9 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_role, role_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, role_name=role_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, role_name=role_name, **kwargs + ).run() return res, check_result @trace() @@ -788,8 +1010,7 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_roles], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, - check_items, check, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() @@ -798,8 +1019,9 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.grant_role, user_name, role_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - user_name=user_name, role_name=role_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, user_name=user_name, role_name=role_name, **kwargs + ).run() return res, check_result @trace() @@ -808,56 +1030,107 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.revoke_role, user_name, role_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - user_name=user_name, role_name=role_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, user_name=user_name, role_name=role_name, **kwargs + ).run() return res, check_result @trace() - def grant_privilege(self, client, role_name, object_type, privilege, object_name, db_name="", - timeout=None, check_task=None, check_items=None, **kwargs): + def grant_privilege( + self, + client, + role_name, + object_type, + privilege, + object_name, + db_name="", + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.grant_privilege, role_name, object_type, privilege, - object_name, db_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - role_name=role_name, object_type=object_type, privilege=privilege, - object_name=object_name, db_name=db_name, **kwargs).run() + res, check = api_request( + [client.grant_privilege, role_name, object_type, privilege, object_name, db_name], **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + role_name=role_name, + object_type=object_type, + privilege=privilege, + object_name=object_name, + db_name=db_name, + **kwargs, + ).run() return res, check_result @trace() - def revoke_privilege(self, client, role_name, object_type, privilege, object_name, db_name="", - timeout=None, check_task=None, check_items=None, **kwargs): + def revoke_privilege( + self, + client, + role_name, + object_type, + privilege, + object_name, + db_name="", + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.revoke_privilege, role_name, object_type, privilege, - object_name, db_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - role_name=role_name, object_type=object_type, privilege=privilege, - object_name=object_name, db_name=db_name, **kwargs).run() + res, check = api_request( + [client.revoke_privilege, role_name, object_type, privilege, object_name, db_name], **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + role_name=role_name, + object_type=object_type, + privilege=privilege, + object_name=object_name, + db_name=db_name, + **kwargs, + ).run() return res, check_result @trace() - def create_privilege_group(self, client, privilege_group: str, timeout=None, check_task=None, check_items=None, **kwargs): + def create_privilege_group( + self, client, privilege_group: str, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_privilege_group, privilege_group], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - privilege_group=privilege_group, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, privilege_group=privilege_group, **kwargs + ).run() return res, check_result @trace() - def drop_privilege_group(self, client, privilege_group: str, timeout=None, check_task=None, check_items=None, **kwargs): + def drop_privilege_group( + self, client, privilege_group: str, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_privilege_group, privilege_group], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - privilege_group=privilege_group, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, privilege_group=privilege_group, **kwargs + ).run() return res, check_result @trace() @@ -871,119 +1144,191 @@ class TestMilvusClientV2Base(Base): return res, check_result @trace() - def add_privileges_to_group(self, client, privilege_group: str, privileges: list, timeout=None, - check_task=None, check_items=None, **kwargs): + def add_privileges_to_group( + self, client, privilege_group: str, privileges: list, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.add_privileges_to_group, privilege_group, privileges], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - privilege_group=privilege_group, privileges=privileges, **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + privilege_group=privilege_group, + privileges=privileges, + **kwargs, + ).run() return res, check_result @trace() - def remove_privileges_from_group(self, client, privilege_group: str, privileges: list, timeout=None, - check_task=None, check_items=None, **kwargs): + def remove_privileges_from_group( + self, client, privilege_group: str, privileges: list, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.remove_privileges_from_group, privilege_group, privileges], - **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - privilege_group=privilege_group, privileges=privileges, **kwargs).run() + res, check = api_request([client.remove_privileges_from_group, privilege_group, privileges], **kwargs) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + privilege_group=privilege_group, + privileges=privileges, + **kwargs, + ).run() return res, check_result @trace() - def grant_privilege_v2(self, client, role_name: str, privilege: str, collection_name: str, db_name=None, timeout=None, - check_task=None, check_items=None, **kwargs): + def grant_privilege_v2( + self, + client, + role_name: str, + privilege: str, + collection_name: str, + db_name=None, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.grant_privilege_v2, role_name, privilege, collection_name, db_name], - **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - role_name=role_name, privilege=privilege, - collection_name=collection_name, db_name=db_name, **kwargs).run() + res, check = api_request([client.grant_privilege_v2, role_name, privilege, collection_name, db_name], **kwargs) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + role_name=role_name, + privilege=privilege, + collection_name=collection_name, + db_name=db_name, + **kwargs, + ).run() return res, check_result @trace() - def revoke_privilege_v2(self, client, role_name: str, privilege: str, collection_name: str, db_name=None, - timeout=None, check_task=None, check_items=None, **kwargs): + def revoke_privilege_v2( + self, + client, + role_name: str, + privilege: str, + collection_name: str, + db_name=None, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.revoke_privilege_v2, role_name, privilege, collection_name, db_name], - **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - role_name=role_name, privilege=privilege, - collection_name=collection_name, db_name=db_name, **kwargs).run() + res, check = api_request([client.revoke_privilege_v2, role_name, privilege, collection_name, db_name], **kwargs) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + role_name=role_name, + privilege=privilege, + collection_name=collection_name, + db_name=db_name, + **kwargs, + ).run() return res, check_result @trace() - def alter_index_properties(self, client, collection_name, index_name, properties, timeout=None, - check_task=None, check_items=None, **kwargs): + def alter_index_properties( + self, client, collection_name, index_name, properties, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.alter_index_properties, collection_name, index_name, properties], - **kwargs) + res, check = api_request([client.alter_index_properties, collection_name, index_name, properties], **kwargs) check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def drop_index_properties(self, client, collection_name, index_name, property_keys, timeout=None, - check_task=None, check_items=None, **kwargs): + def drop_index_properties( + self, + client, + collection_name, + index_name, + property_keys, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.drop_index_properties, collection_name, index_name, property_keys], - **kwargs) + res, check = api_request([client.drop_index_properties, collection_name, index_name, property_keys], **kwargs) check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def alter_collection_properties(self, client, collection_name, properties, timeout=None, - check_task=None, check_items=None, **kwargs): + def alter_collection_properties( + self, client, collection_name, properties, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.alter_collection_properties, collection_name, properties], - **kwargs) + res, check = api_request([client.alter_collection_properties, collection_name, properties], **kwargs) check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def drop_collection_properties(self, client, collection_name, property_keys, timeout=None, - check_task=None, check_items=None, **kwargs): + def drop_collection_properties( + self, client, collection_name, property_keys, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout func_name = sys._getframe().f_code.co_name - res, check = api_request([client.drop_collection_properties, collection_name, property_keys, timeout], - **kwargs) + res, check = api_request([client.drop_collection_properties, collection_name, property_keys, timeout], **kwargs) check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def alter_collection_field(self, client, collection_name, field_name, field_params, timeout=None, - check_task=None, check_items=None, **kwargs): + def alter_collection_field( + self, + client, + collection_name, + field_name, + field_params, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout func_name = sys._getframe().f_code.co_name - res, check = api_request([client.alter_collection_field, collection_name, field_name, field_params, timeout], - **kwargs) + res, check = api_request( + [client.alter_collection_field, collection_name, field_name, field_params, timeout], **kwargs + ) check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def alter_database_properties(self, client, db_name, properties, timeout=None, - check_task=None, check_items=None, **kwargs): + def alter_database_properties( + self, client, db_name, properties, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) @@ -993,8 +1338,9 @@ class TestMilvusClientV2Base(Base): return res, check_result @trace() - def drop_database_properties(self, client, db_name, property_keys, timeout=None, - check_task=None, check_items=None, **kwargs): + def drop_database_properties( + self, client, db_name, property_keys, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) @@ -1039,18 +1385,29 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.run_analyzer, text, analyzer_params], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, text=text, - analyzer_params=analyzer_params, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, text=text, analyzer_params=analyzer_params, **kwargs + ).run() return res, check_result - def compact(self, client, collection_name, is_clustering=False, timeout=None, check_task=None, check_items=None, **kwargs): + def compact( + self, client, collection_name, is_clustering=False, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.compact, collection_name, is_clustering], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, is_clustering=is_clustering, **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + collection_name=collection_name, + is_clustering=is_clustering, + **kwargs, + ).run() return res, check_result @trace() @@ -1070,8 +1427,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.create_resource_group, name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - name=name, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, name=name, **kwargs).run() return res, check_result @trace() @@ -1081,8 +1437,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.update_resource_groups, configs], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - configs=configs, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, configs=configs, **kwargs).run() return res, check_result @trace() @@ -1092,8 +1447,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_resource_group, name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - name=name, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, name=name, **kwargs).run() return res, check_result @trace() @@ -1103,8 +1457,7 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_resource_group, name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - name=name, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, name=name, **kwargs).run() return res, check_result @trace() @@ -1114,33 +1467,62 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_resource_groups], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result @trace() - def transfer_replica(self, client, source_group, target_group, collection_name, num_replicas, - timeout=None, check_task=None, check_items=None, **kwargs): + def transfer_replica( + self, + client, + source_group, + target_group, + collection_name, + num_replicas, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.transfer_replica, source_group, target_group, collection_name, num_replicas], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - source_group=source_group, target_group=target_group, - collection_name=collection_name, num_replicas=num_replicas, - **kwargs).run() + res, check = api_request( + [client.transfer_replica, source_group, target_group, collection_name, num_replicas], **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + source_group=source_group, + target_group=target_group, + collection_name=collection_name, + num_replicas=num_replicas, + **kwargs, + ).run() return res, check_result @trace() - def add_collection_field(self, client, collection_name, field_name, data_type, desc="", timeout=None, check_task=None, check_items=None, **kwargs): + def add_collection_field( + self, + client, + collection_name, + field_name, + data_type, + desc="", + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.add_collection_field, collection_name, field_name, data_type, desc], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run() return res, check_result def wait_schema_version_consistent(self, client, collection_name, timeout=30, poll_interval=0.01): @@ -1195,9 +1577,20 @@ class TestMilvusClientV2Base(Base): f"{collection_name}: consistent={last_consistent}, total={last_total}" ) - def add_collection_field_wait_schema_version_consistency(self, client, collection_name, field_name, data_type, - desc="", timeout=None, check_task=None, check_items=None, - wait_timeout=30, poll_interval=0.01, **kwargs): + def add_collection_field_wait_schema_version_consistency( + self, + client, + collection_name, + field_name, + data_type, + desc="", + timeout=None, + check_task=None, + check_items=None, + wait_timeout=30, + poll_interval=0.01, + **kwargs, + ): """ Wrapper around add_collection_field that first waits for the schema-version consistency gate to pass. Use this in E2E tests that issue successive @@ -1207,16 +1600,32 @@ class TestMilvusClientV2Base(Base): See wait_schema_version_consistent for the polling logic and pass conditions. All add_collection_field arguments are forwarded unchanged. """ - self.wait_schema_version_consistent(client, collection_name, - timeout=wait_timeout, poll_interval=poll_interval) - return self.add_collection_field(client, collection_name, field_name, data_type, desc=desc, - timeout=timeout, check_task=check_task, check_items=check_items, - **kwargs) + self.wait_schema_version_consistent(client, collection_name, timeout=wait_timeout, poll_interval=poll_interval) + return self.add_collection_field( + client, + collection_name, + field_name, + data_type, + desc=desc, + timeout=timeout, + check_task=check_task, + check_items=check_items, + **kwargs, + ) # ====================== Snapshot ====================== @trace() - def create_snapshot(self, client, collection_name, snapshot_name, description="", - timeout=None, check_task=None, check_items=None, **kwargs): + def create_snapshot( + self, + client, + collection_name, + snapshot_name, + description="", + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): """Create a snapshot for a collection. Note: wrapper keeps ``collection_name`` before ``snapshot_name`` for test @@ -1227,24 +1636,40 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.create_snapshot, snapshot_name, collection_name], - description=description, **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - snapshot_name=snapshot_name, collection_name=collection_name, - **kwargs).run() + res, check = api_request( + [client.create_snapshot, snapshot_name, collection_name], description=description, **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + snapshot_name=snapshot_name, + collection_name=collection_name, + **kwargs, + ).run() return res, check_result @trace() - def drop_snapshot(self, client, snapshot_name, collection_name, - timeout=None, check_task=None, check_items=None, **kwargs): + def drop_snapshot( + self, client, snapshot_name, collection_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.drop_snapshot, snapshot_name, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - snapshot_name=snapshot_name, collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + snapshot_name=snapshot_name, + collection_name=collection_name, + **kwargs, + ).run() return res, check_result @trace() @@ -1254,27 +1679,44 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_snapshots, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, **kwargs).run() + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() return res, check_result @trace() - def describe_snapshot(self, client, snapshot_name, collection_name, - timeout=None, check_task=None, check_items=None, **kwargs): + def describe_snapshot( + self, client, snapshot_name, collection_name, timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.describe_snapshot, snapshot_name, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - snapshot_name=snapshot_name, collection_name=collection_name, - **kwargs).run() + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + snapshot_name=snapshot_name, + collection_name=collection_name, + **kwargs, + ).run() return res, check_result @trace() - def restore_snapshot(self, client, snapshot_name, target_collection_name, - source_collection_name="", - timeout=None, check_task=None, check_items=None, **kwargs): + def restore_snapshot( + self, + client, + snapshot_name, + target_collection_name, + source_collection_name="", + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): """Restore a snapshot into a new collection. SDK positional order is ``(snapshot_name, source_collection_name, @@ -1284,13 +1726,20 @@ class TestMilvusClientV2Base(Base): kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name - res, check = api_request([client.restore_snapshot, snapshot_name, - source_collection_name, target_collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - snapshot_name=snapshot_name, - target_collection_name=target_collection_name, - source_collection_name=source_collection_name, - **kwargs).run() + res, check = api_request( + [client.restore_snapshot, snapshot_name, source_collection_name, target_collection_name], **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + snapshot_name=snapshot_name, + target_collection_name=target_collection_name, + source_collection_name=source_collection_name, + **kwargs, + ).run() return res, check_result @trace() @@ -1300,17 +1749,67 @@ class TestMilvusClientV2Base(Base): func_name = sys._getframe().f_code.co_name res, check = api_request([client.get_restore_snapshot_state, job_id], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - job_id=job_id, **kwargs).run() + check_result = ResponseChecker(res, func_name, check_task, check_items, check, job_id=job_id, **kwargs).run() return res, check_result @trace() - def list_restore_snapshot_jobs(self, client, collection_name="", timeout=None, check_task=None, check_items=None, **kwargs): + def list_restore_snapshot_jobs( + self, client, collection_name="", timeout=None, check_task=None, check_items=None, **kwargs + ): timeout = TIMEOUT if timeout is None else timeout kwargs.update({"timeout": timeout}) func_name = sys._getframe().f_code.co_name res, check = api_request([client.list_restore_snapshot_jobs, collection_name], **kwargs) - check_result = ResponseChecker(res, func_name, check_task, check_items, check, - collection_name=collection_name, **kwargs).run() - return res, check_result \ No newline at end of file + check_result = ResponseChecker( + res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs + ).run() + return res, check_result + + @trace() + def pin_snapshot_data( + self, + client, + snapshot_name, + collection_name, + ttl_seconds=0, + timeout=None, + check_task=None, + check_items=None, + **kwargs, + ): + """Pin snapshot data to prevent GC / compaction. Returns pin_id (int). + + SDK signature: ``pin_snapshot_data(snapshot_name, collection_name, + db_name="", ttl_seconds=0, ...)``. + """ + timeout = TIMEOUT if timeout is None else timeout + kwargs.update({"timeout": timeout}) + + func_name = sys._getframe().f_code.co_name + res, check = api_request( + [client.pin_snapshot_data, snapshot_name, collection_name], ttl_seconds=ttl_seconds, **kwargs + ) + check_result = ResponseChecker( + res, + func_name, + check_task, + check_items, + check, + snapshot_name=snapshot_name, + collection_name=collection_name, + ttl_seconds=ttl_seconds, + **kwargs, + ).run() + return res, check_result + + @trace() + def unpin_snapshot_data(self, client, pin_id, timeout=None, check_task=None, check_items=None, **kwargs): + """Release a pin previously created by ``pin_snapshot_data``.""" + timeout = TIMEOUT if timeout is None else timeout + kwargs.update({"timeout": timeout}) + + func_name = sys._getframe().f_code.co_name + res, check = api_request([client.unpin_snapshot_data, pin_id], **kwargs) + check_result = ResponseChecker(res, func_name, check_task, check_items, check, pin_id=pin_id, **kwargs).run() + return res, check_result diff --git a/tests/python_client/milvus_client/test_milvus_client_snapshot.py b/tests/python_client/milvus_client/test_milvus_client_snapshot.py index 2eb6c55678..1be7498d17 100644 --- a/tests/python_client/milvus_client/test_milvus_client_snapshot.py +++ b/tests/python_client/milvus_client/test_milvus_client_snapshot.py @@ -1,16 +1,15 @@ -import time import threading -import pytest -import numpy as np +import time +import numpy as np +import pytest from base.client_v2_base import TestMilvusClientV2Base -from utils.util_log import test_log as log from common import common_func as cf from common import common_type as ct from common.common_type import CaseLabel, CheckTasks -from pymilvus import DataType from ml_dtypes import bfloat16 - +from pymilvus import DataType +from utils.util_log import test_log as log prefix = "snapshot" default_dim = 128 @@ -27,6 +26,8 @@ def wait_for_restore_complete(client, job_id, timeout=60): raise Exception(f"Restore snapshot failed: {state.reason}") time.sleep(1) raise TimeoutError(f"Restore snapshot job {job_id} did not complete within {timeout}s") + + default_nb = 3000 default_nq = 2 default_limit = 10 @@ -54,18 +55,20 @@ class TestMilvusClientSnapshotDefault(TestMilvusClientV2Base): # 1. Create collection and insert data self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - default_float_field_name: i * 1.0, - default_string_field_name: str(i) - } for i in range(default_nb)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + default_float_field_name: i * 1.0, + default_string_field_name: str(i), + } + for i in range(default_nb) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # 2. Create snapshot - self.create_snapshot(client, collection_name, snapshot_name, - description="Test snapshot for L0") + self.create_snapshot(client, collection_name, snapshot_name, description="Test snapshot for L0") # 3. List snapshots snapshots, _ = self.list_snapshots(client, collection_name=collection_name) @@ -99,12 +102,15 @@ class TestMilvusClientSnapshotDefault(TestMilvusClientV2Base): # 1. Create collection and insert data self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - default_float_field_name: i * 1.0, - default_string_field_name: str(i) - } for i in range(default_nb)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + default_float_field_name: i * 1.0, + default_string_field_name: str(i), + } + for i in range(default_nb) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -112,8 +118,9 @@ class TestMilvusClientSnapshotDefault(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # 3. Restore snapshot to new collection - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) assert job_id > 0, "restore_snapshot should return a valid job_id" # 4. Wait for restore to complete @@ -121,18 +128,15 @@ class TestMilvusClientSnapshotDefault(TestMilvusClientV2Base): # 5. Verify restored collection data count self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, filter="", - output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="", output_fields=["count(*)"]) restored_count = res[0]["count(*)"] - assert restored_count == default_nb, \ - f"Restored collection should have {default_nb} rows, got {restored_count}" + assert restored_count == default_nb, f"Restored collection should have {default_nb} rows, got {restored_count}" # 6. Cleanup self.drop_snapshot(client, snapshot_name, collection_name) self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotCreateInvalid(TestMilvusClientV2Base): """Test create_snapshot with invalid parameters - L1""" @@ -150,8 +154,7 @@ class TestMilvusClientSnapshotCreateInvalid(TestMilvusClientV2Base): # SDK validates snapshot_name and raises ParamError error = {ct.err_code: 1, ct.err_msg: "snapshot_name must be a non-empty string"} - self.create_snapshot(client, collection_name, snapshot_name, - check_task=CheckTasks.err_res, check_items=error) + self.create_snapshot(client, collection_name, snapshot_name, check_task=CheckTasks.err_res, check_items=error) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_create_whitespace_name(self): @@ -168,8 +171,7 @@ class TestMilvusClientSnapshotCreateInvalid(TestMilvusClientV2Base): # Server validates snapshot name and rejects whitespace-only names error = {ct.err_code: 1100, ct.err_msg: "snapshot name should be not empty"} - self.create_snapshot(client, collection_name, " ", - check_task=CheckTasks.err_res, check_items=error) + self.create_snapshot(client, collection_name, " ", check_task=CheckTasks.err_res, check_items=error) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_create_collection_not_exist(self): @@ -183,8 +185,9 @@ class TestMilvusClientSnapshotCreateInvalid(TestMilvusClientV2Base): non_existent_collection = cf.gen_unique_str("non_existent") error = {ct.err_code: 100, ct.err_msg: "collection not found"} - self.create_snapshot(client, non_existent_collection, snapshot_name, - check_task=CheckTasks.err_res, check_items=error) + self.create_snapshot( + client, non_existent_collection, snapshot_name, check_task=CheckTasks.err_res, check_items=error + ) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_create_duplicate_name(self): @@ -202,8 +205,7 @@ class TestMilvusClientSnapshotCreateInvalid(TestMilvusClientV2Base): # Try to create another snapshot with same name error = {ct.err_code: 1, ct.err_msg: "already exists"} - self.create_snapshot(client, collection_name, snapshot_name, - check_task=CheckTasks.err_res, check_items=error) + self.create_snapshot(client, collection_name, snapshot_name, check_task=CheckTasks.err_res, check_items=error) # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -229,6 +231,45 @@ class TestMilvusClientSnapshotCreateInvalid(TestMilvusClientV2Base): # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) + @pytest.mark.tags(CaseLabel.L1) + def test_snapshot_create_same_name_different_collections(self): + """ + target: test that snapshot name uniqueness is per-collection, not global + method: create the same snapshot_name under two different collections + expected: both succeed; describe returns the owning collection for each + note: server enforces uniqueness keyed by (collection_id, snapshot_name), + see internal/datacoord/services.go:2093-2099 + """ + client = self._client() + col_a = cf.gen_collection_name_by_testcase_name() + "_a" + col_b = cf.gen_collection_name_by_testcase_name() + "_b" + shared_snapshot = cf.gen_unique_str(prefix + "_shared") + + self.create_collection(client, col_a, default_dim) + self.create_collection(client, col_b, default_dim) + + # Both snapshot creations should succeed under different collections + self.create_snapshot(client, col_a, shared_snapshot) + self.create_snapshot(client, col_b, shared_snapshot) + + # Each collection's list returns its own snapshot + snaps_a, _ = self.list_snapshots(client, collection_name=col_a) + snaps_b, _ = self.list_snapshots(client, collection_name=col_b) + assert shared_snapshot in snaps_a + assert shared_snapshot in snaps_b + + # describe returns the owning collection id/name, not the other one + info_a, _ = self.describe_snapshot(client, shared_snapshot, col_a) + info_b, _ = self.describe_snapshot(client, shared_snapshot, col_b) + assert info_a.collection_name == col_a + assert info_b.collection_name == col_b + + # Cleanup + self.drop_snapshot(client, shared_snapshot, col_a) + self.drop_snapshot(client, shared_snapshot, col_b) + self.drop_collection(client, col_a) + self.drop_collection(client, col_b) + class TestMilvusClientSnapshotDropInvalid(TestMilvusClientV2Base): """Test drop_snapshot with invalid parameters - L1""" @@ -245,8 +286,7 @@ class TestMilvusClientSnapshotDropInvalid(TestMilvusClientV2Base): # SDK validates snapshot_name and raises ParamError before checking collection_name error = {ct.err_code: 1, ct.err_msg: "snapshot_name must be a non-empty string"} - self.drop_snapshot(client, snapshot_name, "", - check_task=CheckTasks.err_res, check_items=error) + self.drop_snapshot(client, snapshot_name, "", check_task=CheckTasks.err_res, check_items=error) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_drop_whitespace_name(self): @@ -263,8 +303,7 @@ class TestMilvusClientSnapshotDropInvalid(TestMilvusClientV2Base): # Server validates snapshot name and rejects whitespace-only names error = {ct.err_code: 1100, ct.err_msg: "snapshot name should be not empty"} - self.drop_snapshot(client, " ", collection_name, - check_task=CheckTasks.err_res, check_items=error) + self.drop_snapshot(client, " ", collection_name, check_task=CheckTasks.err_res, check_items=error) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_drop_not_exist(self): @@ -281,6 +320,69 @@ class TestMilvusClientSnapshotDropInvalid(TestMilvusClientV2Base): # Should not raise exception (idempotent) self.drop_snapshot(client, snapshot_name, collection_name) + @pytest.mark.tags(CaseLabel.L1) + def test_snapshot_drop_during_restore(self): + """ + target: test drop snapshot while restore job is still in progress + method: create snapshot -> start restore -> immediately drop snapshot + expected: drop should fail with error about active restore operations + verified: https://github.com/milvus-io/milvus/issues/47578 + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_collection_name = cf.gen_unique_str(prefix + "_restored") + + # 1. Create collection and insert data + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + # 2. Create snapshot + self.create_snapshot(client, collection_name, snapshot_name) + + # 3. Start restore (creates CopySegment jobs referencing the snapshot) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) + + # 4. Wait until restore is actively in progress (ref count registered) + # before attempting drop, to avoid timing-dependent flakiness + start = time.time() + while time.time() - start < 30: + state = client.get_restore_snapshot_state(job_id) + if state.state not in ("RestoreSnapshotPending",): + break + time.sleep(0.5) + log.info(f"Restore state before drop attempt: {state.state}") + + # 5. Attempt to drop the snapshot while restore is in progress. + # PR #48143 introduced pin-based protection: restore jobs pin the snapshot + # and Drop fails with "active pins exist, unpin before dropping". + error = {ct.err_code: 2601, ct.err_msg: "active pins exist"} + self.drop_snapshot(client, snapshot_name, collection_name, check_task=CheckTasks.err_res, check_items=error) + + # 6. Wait for restore to complete + wait_for_restore_complete(client, job_id) + + # 7. After restore completes, drop should succeed (ref count is 0) + self.drop_snapshot(client, snapshot_name, collection_name) + + # 8. Verify snapshot is actually dropped + snapshots, _ = self.list_snapshots(client, collection_name=collection_name) + assert snapshot_name not in snapshots + + # Cleanup + self.drop_collection(client, restored_collection_name) + class TestMilvusClientSnapshotListDescribe(TestMilvusClientV2Base): """Test list_snapshots and describe_snapshot - L1""" @@ -366,8 +468,7 @@ class TestMilvusClientSnapshotListDescribe(TestMilvusClientV2Base): snapshot_name = cf.gen_unique_str("non_existent") error = {ct.err_code: 1, ct.err_msg: "not found"} - self.describe_snapshot(client, snapshot_name, collection_name, - check_task=CheckTasks.err_res, check_items=error) + self.describe_snapshot(client, snapshot_name, collection_name, check_task=CheckTasks.err_res, check_items=error) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_describe_with_description(self): @@ -390,6 +491,106 @@ class TestMilvusClientSnapshotListDescribe(TestMilvusClientV2Base): # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_list_by_db_name_from_other_context(self): + """ + target: test list_snapshots honors the db_name kwarg from a different active-db context + method: create collection + snapshot in a non-default db, then switch client to + default db and call list_snapshots(collection_name=, db_name=) + expected: snapshot is returned; same call with db_name="default" returns empty/missing + note: server requires collection_name for ListSnapshots + (internal/proxy/task_snapshot.go:448-450) + """ + client = self._client() + target_db = cf.gen_unique_str("test_db_list") + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + + self.create_database(client, target_db) + try: + client.using_database(target_db) + self.create_collection(client, collection_name, default_dim) + self.create_snapshot(client, collection_name, snapshot_name) + + # switch active db back to default but query via db_name kwarg + client.using_database("default") + snapshots, _ = self.list_snapshots(client, collection_name=collection_name, db_name=target_db) + assert snapshot_name in snapshots, ( + f"{snapshot_name} missing when listing via db_name={target_db}, got {snapshots}" + ) + + # cleanup while still in target db context + client.using_database(target_db) + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, collection_name) + finally: + client.using_database("default") + try: + client.drop_database(target_db) + except Exception as e: + log.warning(f"Failed to drop test database '{target_db}': {e}") + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_list_restore_jobs_by_db_name(self): + """ + target: test list_restore_snapshot_jobs(db_name=X) filters by database + method: create snapshot + trigger restore entirely in a non-default db + (explicit source_db_name / target_db_name) -> list jobs via db_name + expected: returned jobs include the one created in target db; default db sees none + """ + client = self._client() + target_db = cf.gen_unique_str("test_db_jobs") + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_name = cf.gen_unique_str(prefix + "_restored") + + self.create_database(client, target_db) + try: + client.using_database(target_db) + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # Explicitly pin both source and target to target_db so the job is + # recorded under target_db's DbId (default empty target_db_name + # resolves to "default" on the server side) + job_id = client.restore_snapshot( + snapshot_name, collection_name, restored_name, source_db_name=target_db, target_db_name=target_db + ) + wait_for_restore_complete(client, job_id) + + # list jobs via explicit db_name kwarg from default db context + client.using_database("default") + jobs_target, _ = self.list_restore_snapshot_jobs(client, collection_name="", db_name=target_db) + job_ids = [j.job_id for j in jobs_target] + assert job_id in job_ids, f"Job {job_id} should appear when listing via db_name={target_db}, got {job_ids}" + + # jobs in default db must not include this job + jobs_default, _ = self.list_restore_snapshot_jobs(client, collection_name="", db_name="default") + default_job_ids = [j.job_id for j in jobs_default] + assert job_id not in default_job_ids, f"Job {job_id} leaked into default db listing: {default_job_ids}" + + # Cleanup: both collections are in target_db + client.using_database(target_db) + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, collection_name) + self.drop_collection(client, restored_name) + finally: + client.using_database("default") + try: + client.drop_database(target_db) + except Exception as e: + log.warning(f"Failed to drop test database '{target_db}': {e}") + class TestMilvusClientSnapshotRestoreInvalid(TestMilvusClientV2Base): """Test restore_snapshot with invalid parameters - L1""" @@ -408,9 +609,14 @@ class TestMilvusClientSnapshotRestoreInvalid(TestMilvusClientV2Base): target_collection_name = cf.gen_unique_str(prefix + "_target") error = {ct.err_code: 1, ct.err_msg: "not found"} - self.restore_snapshot(client, snapshot_name, target_collection_name, - source_collection_name=collection_name, - check_task=CheckTasks.err_res, check_items=error) + self.restore_snapshot( + client, + snapshot_name, + target_collection_name, + source_collection_name=collection_name, + check_task=CheckTasks.err_res, + check_items=error, + ) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_restore_collection_exist(self): @@ -432,9 +638,14 @@ class TestMilvusClientSnapshotRestoreInvalid(TestMilvusClientV2Base): self.create_collection(client, target_collection_name, default_dim) error = {ct.err_code: 65535, ct.err_msg: "duplicate collection"} - self.restore_snapshot(client, snapshot_name, target_collection_name, - source_collection_name=collection_name, - check_task=CheckTasks.err_res, check_items=error) + self.restore_snapshot( + client, + snapshot_name, + target_collection_name, + source_collection_name=collection_name, + check_task=CheckTasks.err_res, + check_items=error, + ) # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -454,8 +665,7 @@ class TestMilvusClientSnapshotRestoreState(TestMilvusClientV2Base): invalid_job_id = 999999999 error = {ct.err_code: 1, ct.err_msg: "not found"} - self.get_restore_snapshot_state(client, invalid_job_id, - check_task=CheckTasks.err_res, check_items=error) + self.get_restore_snapshot_state(client, invalid_job_id, check_task=CheckTasks.err_res, check_items=error) @pytest.mark.tags(CaseLabel.L1) def test_snapshot_list_restore_jobs(self): @@ -488,23 +698,26 @@ class TestMilvusClientSnapshotDataTypes(TestMilvusClientV2Base): # Create collection with Int64 PK self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["id"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["id"]) assert len(res) == 100 ids = sorted([r["id"] for r in res]) assert ids == list(range(100)) @@ -533,27 +746,29 @@ class TestMilvusClientSnapshotDataTypes(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "pk": f"key_{i}", - "vector": list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + "pk": f"key_{i}", + "vector": list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="pk like 'key_%'", output_fields=["pk"]) + res, _ = self.query(client, restored_collection_name, filter="pk like 'key_%'", output_fields=["pk"]) assert len(res) == 100 # Cleanup @@ -582,28 +797,30 @@ class TestMilvusClientSnapshotDataTypes(TestMilvusClientV2Base): index_params.add_index("float_vector", metric_type="COSINE") index_params.add_index("binary_vector", metric_type="HAMMING") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "float_vector": list(rng.random((1, default_dim))[0]), - "binary_vector": bytes(rng.integers(0, 256, size=16, dtype=np.uint8)), - } for i in range(100)] + rows = [ + { + "id": i, + "float_vector": list(rng.random((1, default_dim))[0]), + "binary_vector": bytes(rng.integers(0, 256, size=16, dtype=np.uint8)), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data count self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 100 # Cleanup @@ -631,28 +848,30 @@ class TestMilvusClientSnapshotDataTypes(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random((1, default_dim))[0]), - "metadata": {"key": f"value_{i}", "number": i, "nested": {"a": i}}, - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random((1, default_dim))[0]), + "metadata": {"key": f"value_{i}", "number": i, "nested": {"a": i}}, + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify JSON data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id == 0", output_fields=["metadata"]) + res, _ = self.query(client, restored_collection_name, filter="id == 0", output_fields=["metadata"]) assert res[0]["metadata"]["key"] == "value_0" # Cleanup @@ -679,29 +898,33 @@ class TestMilvusClientSnapshotDataTypes(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random((1, default_dim))[0]), - "dynamic_field_1": f"dynamic_{i}", - "dynamic_field_2": i * 10, - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random((1, default_dim))[0]), + "dynamic_field_1": f"dynamic_{i}", + "dynamic_field_2": i * 10, + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify dynamic field data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id == 0", output_fields=["dynamic_field_1", "dynamic_field_2"]) + res, _ = self.query( + client, restored_collection_name, filter="id == 0", output_fields=["dynamic_field_1", "dynamic_field_2"] + ) assert res[0]["dynamic_field_1"] == "dynamic_0" assert res[0]["dynamic_field_2"] == 0 @@ -710,7 +933,6 @@ class TestMilvusClientSnapshotDataTypes(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): """Test snapshot with partitions - L2""" @@ -737,17 +959,21 @@ class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): # Insert data into each partition rng = np.random.default_rng(seed=19530) for p_name in partition_names: - rows = [{ - default_primary_key_field_name: i + partition_names.index(p_name) * 100, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i + partition_names.index(p_name) * 100, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows, partition_name=p_name) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify partitions are preserved @@ -758,9 +984,9 @@ class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): # Verify data in each partition self.load_collection(client, restored_collection_name) for p_name in partition_names: - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", partition_names=[p_name], - output_fields=["count(*)"]) + res, _ = self.query( + client, restored_collection_name, filter="id >= 0", partition_names=[p_name], output_fields=["count(*)"] + ) assert res[0]["count(*)"] == 100, f"Partition {p_name} should have 100 rows" # Cleanup @@ -785,10 +1011,13 @@ class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): self.create_partition(client, collection_name, partition_name) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows, partition_name=partition_name) self.flush(client, collection_name) @@ -800,8 +1029,9 @@ class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): self.drop_partition(client, collection_name, partition_name) # Restore snapshot - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify partition is restored @@ -810,9 +1040,13 @@ class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): # Verify data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", partition_names=[partition_name], - output_fields=["count(*)"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id >= 0", + partition_names=[partition_name], + output_fields=["count(*)"], + ) assert res[0]["count(*)"] == 100 # Cleanup @@ -820,7 +1054,6 @@ class TestMilvusClientSnapshotPartition(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): """Test snapshot with data operations - L2""" @@ -839,10 +1072,13 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): # Create and insert data self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -855,14 +1091,14 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Restore - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify only 50 rows remain self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 50 # Cleanup @@ -884,10 +1120,13 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): # Create and insert initial data self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -895,28 +1134,30 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Insert more data after snapshot - more_rows = [{ - default_primary_key_field_name: i + 100, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(50)] + more_rows = [ + { + default_primary_key_field_name: i + 100, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(50) + ] self.insert(client, collection_name, more_rows) self.flush(client, collection_name) # Verify source collection has 150 rows self.load_collection(client, collection_name) - res, _ = self.query(client, collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 150 # Restore snapshot - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Restored collection should only have 100 rows (point-in-time) self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 100 # Cleanup @@ -955,27 +1196,32 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): rng = np.random.default_rng(seed=19530) # First batch: insert and flush (this data should be in snapshot) - flushed_rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(100)] + flushed_rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, flushed_rows) self.flush(client, collection_name) log.info("Inserted and flushed 100 rows") # Second batch: insert WITHOUT flush (growing segment, data in buffer) - unflushed_rows = [{ - default_primary_key_field_name: i + 100, - default_vector_field_name: list(rng.random((1, default_dim))[0]), - } for i in range(50)] + unflushed_rows = [ + { + default_primary_key_field_name: i + 100, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for i in range(50) + ] self.insert(client, collection_name, unflushed_rows) # Intentionally NOT calling flush - data stays in growing segment buffer log.info("Inserted 50 rows WITHOUT flush (growing segment)") # Verify source collection can query all 150 rows (growing + flushed) self.load_collection(client, collection_name) - res, _ = self.query(client, collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, collection_name, filter="id >= 0", output_fields=["count(*)"]) source_count = res[0]["count(*)"] log.info(f"Source collection total rows (flushed + growing): {source_count}") assert source_count == 150, f"Source should have 150 rows, got {source_count}" @@ -985,14 +1231,14 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): log.info("Created snapshot (without triggering flush)") # Restore snapshot to new collection - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify restored collection data count self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) restored_count = res[0]["count(*)"] log.info(f"Restored collection rows: {restored_count}") @@ -1000,16 +1246,17 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): # Growing segment data (50 rows) should NOT be captured # NOTE: This assertion documents the current behavior - snapshot does NOT include # growing segment data. If this test fails, it means the behavior has changed. - assert restored_count == 100, \ - f"Expected 100 rows (only flushed data), got {restored_count}. " \ + assert restored_count == 100, ( + f"Expected 100 rows (only flushed data), got {restored_count}. " f"Growing segment data should NOT be included in snapshot." + ) # Also verify the specific IDs: only 0-99 should exist, not 100-149 - res, _ = self.query(client, restored_collection_name, - filter="id >= 100", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 100", output_fields=["count(*)"]) growing_data_count = res[0]["count(*)"] - assert growing_data_count == 0, \ + assert growing_data_count == 0, ( f"Growing segment data (id >= 100) should NOT be in snapshot, found {growing_data_count}" + ) log.info("Verified: Snapshot does NOT include growing segment data") @@ -1018,7 +1265,6 @@ class TestMilvusClientSnapshotDataOperations(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotIndex(TestMilvusClientV2Base): """Test snapshot with various index types - L2""" @@ -1040,25 +1286,28 @@ class TestMilvusClientSnapshotIndex(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="COSINE", - index_type="HNSW", - params={"M": 16, "efConstruction": 200}) + index_params.add_index( + "vector", metric_type="COSINE", index_type="HNSW", params={"M": 16, "efConstruction": 200} + ) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random((1, default_dim))[0]), - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random((1, default_dim))[0]), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify index is preserved @@ -1068,8 +1317,7 @@ class TestMilvusClientSnapshotIndex(TestMilvusClientV2Base): # Verify search works self.load_collection(client, restored_collection_name) search_vectors = [list(rng.random((1, default_dim))[0])] - res, _ = self.search(client, restored_collection_name, search_vectors, - limit=10, output_fields=["id"]) + res, _ = self.search(client, restored_collection_name, search_vectors, limit=10, output_fields=["id"]) assert len(res[0]) == 10 # Cleanup @@ -1077,7 +1325,6 @@ class TestMilvusClientSnapshotIndex(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotDataIntegrity(TestMilvusClientV2Base): """Test snapshot data integrity - verify actual data content, not just counts""" @@ -1097,23 +1344,26 @@ class TestMilvusClientSnapshotDataIntegrity(TestMilvusClientV2Base): self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=12345) # Fixed seed for reproducibility original_vectors = [list(rng.random(default_dim)) for _ in range(100)] - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: original_vectors[i], - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: original_vectors[i], + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Query all vectors from restored collection self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["id", "vector"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["id", "vector"]) # Verify each vector is identical for row in res: @@ -1121,8 +1371,7 @@ class TestMilvusClientSnapshotDataIntegrity(TestMilvusClientV2Base): restored_vec = row["vector"] # Compare with tolerance for floating point for j in range(default_dim): - assert abs(original_vec[j] - restored_vec[j]) < 1e-6, \ - f"Vector mismatch at id={row['id']}, dim={j}" + assert abs(original_vec[j] - restored_vec[j]) < 1e-6, f"Vector mismatch at id={row['id']}, dim={j}" # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -1143,36 +1392,41 @@ class TestMilvusClientSnapshotDataIntegrity(TestMilvusClientV2Base): # Create collection and insert data self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(1000)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(1000) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) self.load_collection(client, collection_name) # Search on original collection query_vectors = [list(rng.random(default_dim)) for _ in range(10)] - original_results, _ = self.search(client, collection_name, query_vectors, - limit=10, output_fields=["id"]) + original_results, _ = self.search(client, collection_name, query_vectors, limit=10, output_fields=["id"]) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Search on restored collection with same queries self.load_collection(client, restored_collection_name) - restored_results, _ = self.search(client, restored_collection_name, query_vectors, - limit=10, output_fields=["id"]) + restored_results, _ = self.search( + client, restored_collection_name, query_vectors, limit=10, output_fields=["id"] + ) # Compare search results for i in range(len(query_vectors)): original_ids = [r["id"] for r in original_results[i]] restored_ids = [r["id"] for r in restored_results[i]] - assert original_ids == restored_ids, \ + assert original_ids == restored_ids, ( f"Search results mismatch for query {i}: original={original_ids}, restored={restored_ids}" + ) # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -1202,33 +1456,39 @@ class TestMilvusClientSnapshotDataIntegrity(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data with various values rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "int_field": i * 10, - "float_field": i * 0.5, - "bool_field": i % 2 == 0, - "varchar_field": f"string_value_{i}", - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "int_field": i * 10, + "float_field": i * 0.5, + "bool_field": i % 2 == 0, + "varchar_field": f"string_value_{i}", + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Query and verify all scalar values self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", - output_fields=["id", "int_field", "float_field", "bool_field", "varchar_field"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id >= 0", + output_fields=["id", "int_field", "float_field", "bool_field", "varchar_field"], + ) for row in res: i = row["id"] @@ -1242,7 +1502,6 @@ class TestMilvusClientSnapshotDataIntegrity(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): """Test snapshot boundary conditions and edge cases""" @@ -1330,16 +1589,20 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): # Create collection with more data to slow down restore self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(5000)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(5000) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) # Track progress progress_values = [] @@ -1360,8 +1623,9 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): assert 100 in progress_values, "Progress should reach 100 when completed" # Verify progress was monotonically increasing (or at least non-decreasing) for i in range(1, len(progress_values)): - assert progress_values[i] >= progress_values[i-1], \ - f"Progress should not decrease: {progress_values[i-1]} -> {progress_values[i]}" + assert progress_values[i] >= progress_values[i - 1], ( + f"Progress should not decrease: {progress_values[i - 1]} -> {progress_values[i]}" + ) # Verify start_time and time_cost are set final_state, _ = self.get_restore_snapshot_state(client, job_id) @@ -1391,10 +1655,13 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): # Create 3 snapshots at different data states for batch in range(3): # Insert 100 rows - rows = [{ - default_primary_key_field_name: i + batch * 100, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i + batch * 100, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -1407,16 +1674,18 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): # Restore each snapshot and verify correct count for i, snapshot_name in enumerate(snapshots): restored_name = cf.gen_unique_str(prefix + "_restored") - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) self.load_collection(client, restored_name) res, _ = self.query(client, restored_name, filter="id >= 0", output_fields=["count(*)"]) actual_count = res[0]["count(*)"] - assert actual_count == expected_counts[i], \ + assert actual_count == expected_counts[i], ( f"Snapshot {i} should have {expected_counts[i]} rows, got {actual_count}" + ) self.drop_collection(client, restored_name) @@ -1438,10 +1707,13 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): # Create collection with data self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(500)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -1456,8 +1728,9 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): for i in range(num_restores): restored_name = cf.gen_unique_str(prefix + f"_restored_{i}") restored_names.append(restored_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_name, source_collection_name=collection_name + ) restore_jobs.append(job_id) # Wait for all to complete @@ -1475,7 +1748,6 @@ class TestMilvusClientSnapshotBoundary(TestMilvusClientV2Base): self.drop_snapshot(client, snapshot_name, collection_name) - class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): """Test snapshot negative scenarios and error handling""" @@ -1494,10 +1766,13 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): # Create collection and snapshot self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) self.create_snapshot(client, collection_name, snapshot_name) @@ -1507,9 +1782,14 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): # Try to restore - should fail error = {ct.err_code: 1, ct.err_msg: "not found"} - self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name, - check_task=CheckTasks.err_res, check_items=error) + self.restore_snapshot( + client, + snapshot_name, + restored_collection_name, + source_collection_name=collection_name, + check_task=CheckTasks.err_res, + check_items=error, + ) @pytest.mark.tags(CaseLabel.L2) def test_snapshot_cascade_delete_on_drop_collection(self): @@ -1528,10 +1808,13 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): # Create collection and snapshot self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) self.create_snapshot(client, collection_name, snapshot_name) @@ -1548,8 +1831,7 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): # Snapshots should not exist for the new collection snapshots, _ = self.list_snapshots(client, collection_name=collection_name) - assert len(snapshots) == 0, \ - "Snapshots should be cascade deleted when collection is dropped" + assert len(snapshots) == 0, "Snapshots should be cascade deleted when collection is dropped" @pytest.mark.tags(CaseLabel.L2) def test_snapshot_schema_consistency_autoID(self): @@ -1571,8 +1853,7 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data (no id needed since auto_id=True) rng = np.random.default_rng(seed=19530) @@ -1582,15 +1863,16 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify schema of restored collection desc = client.describe_collection(restored_collection_name) # Check auto_id is preserved pk_field = [f for f in desc["fields"] if f.get("is_primary")][0] - assert pk_field.get("auto_id", False) == True, "auto_id should be preserved" + assert pk_field.get("auto_id", False), "auto_id should be preserved" # Verify can insert without id new_rows = [{"vector": list(rng.random(default_dim))} for _ in range(10)] @@ -1601,7 +1883,6 @@ class TestMilvusClientSnapshotNegative(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): """ L2 Test - Snapshot with all data types matrix testing @@ -1635,37 +1916,51 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data with all scalar types rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "int8_field": np.int8(i % 127), - "int16_field": np.int16(i * 10), - "int32_field": np.int32(i * 100), - "bool_field": i % 2 == 0, - "float_field": float(i * 0.5), - "double_field": float(i * 1.5), - "varchar_field": f"string_{i}", - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "int8_field": np.int8(i % 127), + "int16_field": np.int16(i * 10), + "int32_field": np.int32(i * 100), + "bool_field": i % 2 == 0, + "float_field": float(i * 0.5), + "double_field": float(i * 1.5), + "varchar_field": f"string_{i}", + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify all scalar data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", - output_fields=["id", "int8_field", "int16_field", "int32_field", - "bool_field", "float_field", "double_field", "varchar_field"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id >= 0", + output_fields=[ + "id", + "int8_field", + "int16_field", + "int32_field", + "bool_field", + "float_field", + "double_field", + "varchar_field", + ], + ) assert len(res) == 100 # Verify specific values @@ -1701,40 +1996,47 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field("int_array", DataType.ARRAY, element_type=DataType.INT64, max_capacity=50) schema.add_field("float_array", DataType.ARRAY, element_type=DataType.FLOAT, max_capacity=50) - schema.add_field("varchar_array", DataType.ARRAY, element_type=DataType.VARCHAR, - max_length=100, max_capacity=50) + schema.add_field( + "varchar_array", DataType.ARRAY, element_type=DataType.VARCHAR, max_length=100, max_capacity=50 + ) schema.add_field("bool_array", DataType.ARRAY, element_type=DataType.BOOL, max_capacity=50) index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data with array types rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "int_array": [i * j for j in range(10)], - "float_array": [float(i * j * 0.1) for j in range(10)], - "varchar_array": [f"str_{i}_{j}" for j in range(5)], - "bool_array": [j % 2 == 0 for j in range(5)], - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "int_array": [i * j for j in range(10)], + "float_array": [float(i * j * 0.1) for j in range(10)], + "varchar_array": [f"str_{i}_{j}" for j in range(5)], + "bool_array": [j % 2 == 0 for j in range(5)], + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify array data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id == 5", - output_fields=["id", "int_array", "float_array", "varchar_array", "bool_array"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id == 5", + output_fields=["id", "int_array", "float_array", "varchar_array", "bool_array"], + ) assert len(res) == 1 row = res[0] assert row["int_array"] == [5 * j for j in range(10)] @@ -1772,8 +2074,7 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): index_params.add_index("float16_vector", metric_type="L2") index_params.add_index("sparse_vector", metric_type="IP", index_type="SPARSE_INVERTED_INDEX") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Generate test data rng = np.random.default_rng(seed=19530) @@ -1785,32 +2086,35 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): # Sparse vector: {dim_index: value} sparse_vec = {j: float(rng.random()) for j in rng.choice(1000, size=10, replace=False)} - rows.append({ - "id": i, - "float_vector": float_vec, - "binary_vector": binary_vec, - "float16_vector": float16_vec, - "sparse_vector": sparse_vec, - }) + rows.append( + { + "id": i, + "float_vector": float_vec, + "binary_vector": binary_vec, + "float16_vector": float16_vec, + "sparse_vector": sparse_vec, + } + ) self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data count self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 100 # Verify search on float_vector works search_vectors = [list(rng.random(default_dim))] - search_res, _ = self.search(client, restored_collection_name, search_vectors, - anns_field="float_vector", limit=10, output_fields=["id"]) + search_res, _ = self.search( + client, restored_collection_name, search_vectors, anns_field="float_vector", limit=10, output_fields=["id"] + ) assert len(search_res[0]) == 10 # Cleanup @@ -1840,8 +2144,7 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data with some null values rng = np.random.default_rng(seed=19530) @@ -1860,26 +2163,33 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify nullable fields self.load_collection(client, restored_collection_name) # Check rows with null values - res, _ = self.query(client, restored_collection_name, - filter="id == 0", # i=0 should have nullable_int=None - output_fields=["nullable_int", "nullable_varchar", "nullable_float"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id == 0", # i=0 should have nullable_int=None + output_fields=["nullable_int", "nullable_varchar", "nullable_float"], + ) assert len(res) == 1 assert res[0]["nullable_int"] is None, "nullable_int should be None for id=0" assert res[0]["nullable_varchar"] is None, "nullable_varchar should be None for id=0" assert res[0]["nullable_float"] is None, "nullable_float should be None for id=0" # Check rows with non-null values - res, _ = self.query(client, restored_collection_name, - filter="id == 7", # i=7: nullable_int=70, nullable_varchar='str_7', nullable_float=3.5 - output_fields=["nullable_int", "nullable_varchar", "nullable_float"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id == 7", # i=7: nullable_int=70, nullable_varchar='str_7', nullable_float=3.5 + output_fields=["nullable_int", "nullable_varchar", "nullable_float"], + ) assert len(res) == 1 assert res[0]["nullable_int"] == 70 assert res[0]["nullable_varchar"] == "str_7" @@ -1905,23 +2215,18 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): # Generate comprehensive schema with all data types schema = cf.gen_all_datatype_collection_schema( - dim=default_dim, - enable_struct_array_field=True, - enable_dynamic_field=True, - nullable=True + dim=default_dim, enable_struct_array_field=True, enable_dynamic_field=True, nullable=True ) # Create indexes for vector fields index_params = client.prepare_index_params() index_params.add_index("float_vector", metric_type="COSINE") - index_params.add_index("text_sparse_emb", metric_type="BM25", - index_type="SPARSE_INVERTED_INDEX") + index_params.add_index("text_sparse_emb", metric_type="BM25", index_type="SPARSE_INVERTED_INDEX") index_params.add_index("minhash_emb", metric_type="HAMMING") # Struct array inner vector field also needs index index_params.add_index("array_struct[float_vector]", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Generate row data using schema-based data generator nb = 200 @@ -1931,43 +2236,43 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): # Verify source collection self.load_collection(client, collection_name) - res, _ = self.query(client, collection_name, filter="", - output_fields=["count(*)"]) + res, _ = self.query(client, collection_name, filter="", output_fields=["count(*)"]) assert res[0]["count(*)"] == nb # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify restored data count self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, filter="", - output_fields=["count(*)"]) - assert res[0]["count(*)"] == nb, \ - f"Expected {nb} rows, got {res[0]['count(*)']}" + res, _ = self.query(client, restored_collection_name, filter="", output_fields=["count(*)"]) + assert res[0]["count(*)"] == nb, f"Expected {nb} rows, got {res[0]['count(*)']}" # Verify text match works (BM25 function preserved) - res, _ = self.query(client, restored_collection_name, - filter='TEXT_MATCH(text, "the")', - output_fields=["id", "text"]) + res, _ = self.query( + client, restored_collection_name, filter='TEXT_MATCH(text, "the")', output_fields=["id", "text"] + ) log.info(f"Text match results: {len(res)} rows") # Verify float vector search works rng = np.random.default_rng(seed=19530) search_vectors = [list(rng.random(default_dim))] - search_res, _ = self.search(client, restored_collection_name, search_vectors, - anns_field="float_vector", limit=10, - output_fields=["id"]) + search_res, _ = self.search( + client, restored_collection_name, search_vectors, anns_field="float_vector", limit=10, output_fields=["id"] + ) assert len(search_res[0]) > 0, "Float vector search should return results" # Verify scalar fields are preserved (PK field name is "int64") - res, _ = self.query(client, restored_collection_name, - filter="int64 >= 0", - output_fields=["int64", "varchar", "json_field", - "array_int", "array_bool", "array_struct"], - limit=10) + res, _ = self.query( + client, + restored_collection_name, + filter="int64 >= 0", + output_fields=["int64", "varchar", "json_field", "array_int", "array_bool", "array_struct"], + limit=10, + ) assert len(res) > 0 # Check at least one row has non-null array fields has_array_data = False @@ -2009,35 +2314,42 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): index_params.add_index("float_vector", metric_type="COSINE") index_params.add_index("bfloat16_vector", metric_type="L2") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "float_vector": list(rng.random(default_dim)), - "bfloat16_vector": np.array(rng.random(default_dim), dtype=bfloat16), - } for i in range(100)] + rows = [ + { + "id": i, + "float_vector": list(rng.random(default_dim)), + "bfloat16_vector": np.array(rng.random(default_dim), dtype=bfloat16), + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data count self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 100 # Verify search on bfloat16 vector works search_vectors = [np.array(rng.random(default_dim), dtype=bfloat16)] - res, _ = self.search(client, restored_collection_name, search_vectors, - anns_field="bfloat16_vector", limit=10, - output_fields=["id"]) + res, _ = self.search( + client, + restored_collection_name, + search_vectors, + anns_field="bfloat16_vector", + limit=10, + output_fields=["id"], + ) assert len(res[0]) == 10 # Cleanup @@ -2066,43 +2378,56 @@ class TestMilvusClientSnapshotAllDataTypes(TestMilvusClientV2Base): schema.add_field("arr_float", DataType.ARRAY, element_type=DataType.FLOAT, max_capacity=20) schema.add_field("arr_double", DataType.ARRAY, element_type=DataType.DOUBLE, max_capacity=20) schema.add_field("arr_bool", DataType.ARRAY, element_type=DataType.BOOL, max_capacity=20) - schema.add_field("arr_varchar", DataType.ARRAY, element_type=DataType.VARCHAR, - max_length=100, max_capacity=20) + schema.add_field("arr_varchar", DataType.ARRAY, element_type=DataType.VARCHAR, max_length=100, max_capacity=20) index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "arr_int8": [int(np.int8(j)) for j in range(5)], - "arr_int16": [int(np.int16(i * 10 + j)) for j in range(5)], - "arr_int32": [i * 100 + j for j in range(5)], - "arr_int64": [i * 1000 + j for j in range(5)], - "arr_float": [float(i * 0.1 + j * 0.01) for j in range(5)], - "arr_double": [float(i * 1.1 + j * 0.11) for j in range(5)], - "arr_bool": [j % 2 == 0 for j in range(5)], - "arr_varchar": [f"s_{i}_{j}" for j in range(5)], - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "arr_int8": [int(np.int8(j)) for j in range(5)], + "arr_int16": [int(np.int16(i * 10 + j)) for j in range(5)], + "arr_int32": [i * 100 + j for j in range(5)], + "arr_int64": [i * 1000 + j for j in range(5)], + "arr_float": [float(i * 0.1 + j * 0.01) for j in range(5)], + "arr_double": [float(i * 1.1 + j * 0.11) for j in range(5)], + "arr_bool": [j % 2 == 0 for j in range(5)], + "arr_varchar": [f"s_{i}_{j}" for j in range(5)], + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id == 5", - output_fields=["arr_int8", "arr_int16", "arr_int32", "arr_int64", - "arr_float", "arr_double", "arr_bool", "arr_varchar"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id == 5", + output_fields=[ + "arr_int8", + "arr_int16", + "arr_int32", + "arr_int64", + "arr_float", + "arr_double", + "arr_bool", + "arr_varchar", + ], + ) assert len(res) == 1 row = res[0] assert row["arr_int32"] == [500, 501, 502, 503, 504] @@ -2138,12 +2463,9 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="L2", - index_type="IVF_FLAT", - params={"nlist": 128}) + index_params.add_index("vector", metric_type="L2", index_type="IVF_FLAT", params={"nlist": 128}) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(1000)] @@ -2152,15 +2474,22 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify index and search self.load_collection(client, restored_collection_name) search_vectors = [list(rng.random(default_dim))] - res, _ = self.search(client, restored_collection_name, search_vectors, - search_params={"nprobe": 16}, limit=10, output_fields=["id"]) + res, _ = self.search( + client, + restored_collection_name, + search_vectors, + search_params={"nprobe": 16}, + limit=10, + output_fields=["id"], + ) assert len(res[0]) == 10 # Cleanup @@ -2184,12 +2513,9 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="L2", - index_type="IVF_SQ8", - params={"nlist": 128}) + index_params.add_index("vector", metric_type="L2", index_type="IVF_SQ8", params={"nlist": 128}) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(1000)] @@ -2198,15 +2524,22 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify search works self.load_collection(client, restored_collection_name) search_vectors = [list(rng.random(default_dim))] - res, _ = self.search(client, restored_collection_name, search_vectors, - search_params={"nprobe": 16}, limit=10, output_fields=["id"]) + res, _ = self.search( + client, + restored_collection_name, + search_vectors, + search_params={"nprobe": 16}, + limit=10, + output_fields=["id"], + ) assert len(res[0]) == 10 # Cleanup @@ -2230,12 +2563,11 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="L2", - index_type="IVF_PQ", - params={"nlist": 128, "m": 16, "nbits": 8}) + index_params.add_index( + "vector", metric_type="L2", index_type="IVF_PQ", params={"nlist": 128, "m": 16, "nbits": 8} + ) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(1000)] @@ -2244,15 +2576,22 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify search works self.load_collection(client, restored_collection_name) search_vectors = [list(rng.random(default_dim))] - res, _ = self.search(client, restored_collection_name, search_vectors, - search_params={"nprobe": 16}, limit=10, output_fields=["id"]) + res, _ = self.search( + client, + restored_collection_name, + search_vectors, + search_params={"nprobe": 16}, + limit=10, + output_fields=["id"], + ) assert len(res[0]) == 10 # Cleanup @@ -2276,11 +2615,9 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="L2", - index_type="DISKANN") + index_params.add_index("vector", metric_type="L2", index_type="DISKANN") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(1000)] @@ -2289,15 +2626,22 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify search works self.load_collection(client, restored_collection_name) search_vectors = [list(rng.random(default_dim))] - res, _ = self.search(client, restored_collection_name, search_vectors, - search_params={"search_list": 100}, limit=10, output_fields=["id"]) + res, _ = self.search( + client, + restored_collection_name, + search_vectors, + search_params={"search_list": 100}, + limit=10, + output_fields=["id"], + ) assert len(res[0]) == 10 # Cleanup @@ -2321,12 +2665,9 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="L2", - index_type="SCANN", - params={"nlist": 128}) + index_params.add_index("vector", metric_type="L2", index_type="SCANN", params={"nlist": 128}) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(1000)] @@ -2335,15 +2676,22 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify search works self.load_collection(client, restored_collection_name) search_vectors = [list(rng.random(default_dim))] - res, _ = self.search(client, restored_collection_name, search_vectors, - search_params={"nprobe": 16}, limit=10, output_fields=["id"]) + res, _ = self.search( + client, + restored_collection_name, + search_vectors, + search_params={"nprobe": 16}, + limit=10, + output_fields=["id"], + ) assert len(res[0]) == 10 # Cleanup @@ -2373,33 +2721,34 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): index_params.add_index("category", index_type="STL_SORT") index_params.add_index("tag", index_type="INVERTED") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "category": i % 10, - "tag": f"tag_{i % 5}", - } for i in range(500)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "category": i % 10, + "tag": f"tag_{i % 5}", + } + for i in range(500) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify scalar index works with filter self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="category == 5", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="category == 5", output_fields=["count(*)"]) assert res[0]["count(*)"] == 50 # 500/10 = 50 rows with category=5 - res, _ = self.query(client, restored_collection_name, - filter="tag == 'tag_3'", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="tag == 'tag_3'", output_fields=["count(*)"]) assert res[0]["count(*)"] == 100 # 500/5 = 100 rows with tag='tag_3' # Cleanup @@ -2407,7 +2756,6 @@ class TestMilvusClientSnapshotAllIndexTypes(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): """ L2 Test - Snapshot with collection properties testing @@ -2428,16 +2776,14 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): description = "Test collection for snapshot with description preservation" - schema = client.create_schema(enable_dynamic_field=False, auto_id=False, - description=description) + schema = client.create_schema(enable_dynamic_field=False, auto_id=False, description=description) schema.add_field("id", DataType.INT64, is_primary=True) schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(100)] @@ -2446,14 +2792,16 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify description is preserved desc = client.describe_collection(restored_collection_name) - assert desc.get("description") == description, \ + assert desc.get("description") == description, ( f"Description should be preserved, got: {desc.get('description')}" + ) # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -2480,8 +2828,7 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params, num_shards=num_shards) + self.create_collection(client, collection_name, schema=schema, index_params=index_params, num_shards=num_shards) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(100)] @@ -2495,16 +2842,16 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify shard count is preserved desc = client.describe_collection(restored_collection_name) restored_shards = desc.get("num_shards") or desc.get("shards_num") log.info(f"Restored collection shards: {restored_shards}") - assert restored_shards == num_shards, \ - f"Shard count should be {num_shards}, got: {restored_shards}" + assert restored_shards == num_shards, f"Shard count should be {num_shards}, got: {restored_shards}" # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -2530,8 +2877,9 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): index_params.add_index("vector", metric_type="COSINE") # Create collection with Bounded consistency - self.create_collection(client, collection_name, schema=schema, - index_params=index_params, consistency_level="Bounded") + self.create_collection( + client, collection_name, schema=schema, index_params=index_params, consistency_level="Bounded" + ) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(100)] @@ -2545,8 +2893,9 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify consistency level is preserved @@ -2580,22 +2929,25 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params, num_partitions=16) + self.create_collection(client, collection_name, schema=schema, index_params=index_params, num_partitions=16) rng = np.random.default_rng(seed=19530) - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "category": i % 100, - } for i in range(500)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "category": i % 100, + } + for i in range(500) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify partition key is preserved in schema @@ -2603,13 +2955,11 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): fields = desc.get("fields", []) category_field = [f for f in fields if f.get("name") == "category"] assert len(category_field) == 1 - assert category_field[0].get("is_partition_key") == True, \ - "Partition key should be preserved" + assert category_field[0].get("is_partition_key"), "Partition key should be preserved" # Verify data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 500 # Cleanup @@ -2617,7 +2967,6 @@ class TestMilvusClientSnapshotCollectionProperties(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): """ L2 Test - Snapshot after various data operations @@ -2640,22 +2989,28 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): rng = np.random.default_rng(seed=19530) # Initial insert: ids 0-99 - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - default_float_field_name: float(i), - default_string_field_name: f"original_{i}", - } for i in range(100)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + default_float_field_name: float(i), + default_string_field_name: f"original_{i}", + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Upsert: update ids 50-99 and insert ids 100-149 - upsert_rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - default_float_field_name: float(i * 10), # Updated value - default_string_field_name: f"updated_{i}", - } for i in range(50, 150)] + upsert_rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + default_float_field_name: float(i * 10), # Updated value + default_string_field_name: f"updated_{i}", + } + for i in range(50, 150) + ] self.upsert(client, collection_name, upsert_rows) self.flush(client, collection_name) @@ -2663,36 +3018,30 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Restore - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify self.load_collection(client, restored_collection_name) # Total count should be 150 (0-149) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 150, f"Expected 150 rows, got {res[0]['count(*)']}" # Check original data (0-49) unchanged - res, _ = self.query(client, restored_collection_name, - filter="id == 25", - output_fields=["float", "varchar"]) + res, _ = self.query(client, restored_collection_name, filter="id == 25", output_fields=["float", "varchar"]) assert res[0]["float"] == 25.0 assert res[0]["varchar"] == "original_25" # Check updated data (50-99) - res, _ = self.query(client, restored_collection_name, - filter="id == 75", - output_fields=["float", "varchar"]) + res, _ = self.query(client, restored_collection_name, filter="id == 75", output_fields=["float", "varchar"]) assert res[0]["float"] == 750.0 # Updated value assert res[0]["varchar"] == "updated_75" # Check new data (100-149) - res, _ = self.query(client, restored_collection_name, - filter="id == 125", - output_fields=["float", "varchar"]) + res, _ = self.query(client, restored_collection_name, filter="id == 125", output_fields=["float", "varchar"]) assert res[0]["float"] == 1250.0 assert res[0]["varchar"] == "updated_125" @@ -2716,10 +3065,13 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): rng = np.random.default_rng(seed=19530) # Insert data - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(1000)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(1000) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -2739,19 +3091,18 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Restore - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data count (should be 700: 1000 - 300 deleted) self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 700, f"Expected 700 rows, got {res[0]['count(*)']}" # Verify deleted data is not present - res, _ = self.query(client, restored_collection_name, - filter="id < 300", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id < 300", output_fields=["count(*)"]) assert res[0]["count(*)"] == 0, "Deleted data should not be present" # Cleanup @@ -2778,12 +3129,11 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): schema.add_field("vector", DataType.FLOAT_VECTOR, dim=default_dim) index_params = client.prepare_index_params() - index_params.add_index("vector", metric_type="COSINE", - index_type="HNSW", - params={"M": 16, "efConstruction": 200}) + index_params.add_index( + "vector", metric_type="COSINE", index_type="HNSW", params={"M": 16, "efConstruction": 200} + ) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) rows = [{"id": i, "vector": list(rng.random(default_dim))} for i in range(500)] @@ -2800,9 +3150,7 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): # Create new IVF_FLAT index new_index_params = client.prepare_index_params() - new_index_params.add_index("vector", metric_type="L2", - index_type="IVF_FLAT", - params={"nlist": 128}) + new_index_params.add_index("vector", metric_type="L2", index_type="IVF_FLAT", params={"nlist": 128}) self.create_index(client, collection_name, new_index_params) log.info("Reindexed with IVF_FLAT") @@ -2811,21 +3159,22 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): log.info("Created snapshot with IVF_FLAT index") # Restore snapshot 1 (HNSW) - job_id_1, _ = self.restore_snapshot(client, snapshot_name_1, restored_collection_name_1, - source_collection_name=collection_name) + job_id_1, _ = self.restore_snapshot( + client, snapshot_name_1, restored_collection_name_1, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id_1) # Restore snapshot 2 (IVF_FLAT) - job_id_2, _ = self.restore_snapshot(client, snapshot_name_2, restored_collection_name_2, - source_collection_name=collection_name) + job_id_2, _ = self.restore_snapshot( + client, snapshot_name_2, restored_collection_name_2, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id_2) # Verify both collections can search for restored_name in [restored_collection_name_1, restored_collection_name_2]: self.load_collection(client, restored_name) search_vectors = [list(rng.random(default_dim))] - res, _ = self.search(client, restored_name, search_vectors, - limit=10, output_fields=["id"]) + res, _ = self.search(client, restored_name, search_vectors, limit=10, output_fields=["id"]) assert len(res[0]) == 10, f"Search should return 10 results for {restored_name}" # Cleanup @@ -2852,10 +3201,13 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): total_rows = 0 # Insert multiple batches to create multiple segments for batch in range(5): - rows = [{ - default_primary_key_field_name: i + batch * 200, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(200)] + rows = [ + { + default_primary_key_field_name: i + batch * 200, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(200) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) total_rows += 200 @@ -2865,20 +3217,18 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Restore - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify all data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) - assert res[0]["count(*)"] == total_rows, \ - f"Expected {total_rows} rows, got {res[0]['count(*)']}" + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == total_rows, f"Expected {total_rows} rows, got {res[0]['count(*)']}" # Verify data range - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["id"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["id"]) ids = sorted([r["id"] for r in res]) assert ids == list(range(total_rows)), "All IDs should be present" @@ -2902,11 +3252,14 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): rng = np.random.default_rng(seed=19530) # Step 1: Initial insert (0-199) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - default_string_field_name: f"original_{i}", - } for i in range(200)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + default_string_field_name: f"original_{i}", + } + for i in range(200) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) log.info("Initial insert: 200 rows (0-199)") @@ -2918,11 +3271,14 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): log.info("Deleted rows 0-49") # Step 3: Upsert (update 100-149, insert 200-249) - upsert_rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - default_string_field_name: f"upserted_{i}", - } for i in range(100, 250)] + upsert_rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + default_string_field_name: f"upserted_{i}", + } + for i in range(100, 250) + ] self.upsert(client, collection_name, upsert_rows) self.flush(client, collection_name) log.info("Upserted rows 100-249 (update 100-149, insert 200-249)") @@ -2931,36 +3287,32 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Restore - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify final state self.load_collection(client, restored_collection_name) # Expected: rows 50-249 = 200 rows - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 200, f"Expected 200 rows, got {res[0]['count(*)']}" # Deleted rows (0-49) should not exist - res, _ = self.query(client, restored_collection_name, - filter="id < 50", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id < 50", output_fields=["count(*)"]) assert res[0]["count(*)"] == 0, "Deleted rows should not exist" # Original rows (50-99) should have original values - res, _ = self.query(client, restored_collection_name, - filter="id == 75", output_fields=["varchar"]) + res, _ = self.query(client, restored_collection_name, filter="id == 75", output_fields=["varchar"]) assert res[0]["varchar"] == "original_75" # Upserted rows (100-149) should have updated values - res, _ = self.query(client, restored_collection_name, - filter="id == 125", output_fields=["varchar"]) + res, _ = self.query(client, restored_collection_name, filter="id == 125", output_fields=["varchar"]) assert res[0]["varchar"] == "upserted_125" # New rows (200-249) should exist - res, _ = self.query(client, restored_collection_name, - filter="id == 225", output_fields=["varchar"]) + res, _ = self.query(client, restored_collection_name, filter="id == 225", output_fields=["varchar"]) assert res[0]["varchar"] == "upserted_225" # Cleanup @@ -2988,16 +3340,18 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) # Insert data with categories - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "category": i % 10, - } for i in range(1000)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "category": i % 10, + } + for i in range(1000) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -3013,20 +3367,21 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): self.create_snapshot(client, collection_name, snapshot_name) # Restore - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 1000, f"Expected 1000 rows, got {res[0]['count(*)']}" # Verify category data integrity for cat in range(10): - res, _ = self.query(client, restored_collection_name, - filter=f"category == {cat}", output_fields=["count(*)"]) + res, _ = self.query( + client, restored_collection_name, filter=f"category == {cat}", output_fields=["count(*)"] + ) assert res[0]["count(*)"] == 100, f"Category {cat} should have 100 rows" # Cleanup @@ -3053,42 +3408,47 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index("vector", metric_type="COSINE") - self.create_collection(client, collection_name, schema=schema, - index_params=index_params) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) rng = np.random.default_rng(seed=19530) # Insert data with dynamic fields - rows = [{ - "id": i, - "vector": list(rng.random(default_dim)), - "dynamic_str": f"dynamic_{i}", - "dynamic_int": i * 100, - "dynamic_float": float(i * 0.5), - "dynamic_bool": i % 2 == 0, - } for i in range(100)] + rows = [ + { + "id": i, + "vector": list(rng.random(default_dim)), + "dynamic_str": f"dynamic_{i}", + "dynamic_int": i * 100, + "dynamic_float": float(i * 0.5), + "dynamic_bool": i % 2 == 0, + } + for i in range(100) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) # Create snapshot and restore self.create_snapshot(client, collection_name, snapshot_name) - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_collection_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) wait_for_restore_complete(client, job_id) # Verify dynamic field data self.load_collection(client, restored_collection_name) - res, _ = self.query(client, restored_collection_name, - filter="id == 50", - output_fields=["id", "dynamic_str", "dynamic_int", "dynamic_float", "dynamic_bool"]) + res, _ = self.query( + client, + restored_collection_name, + filter="id == 50", + output_fields=["id", "dynamic_str", "dynamic_int", "dynamic_float", "dynamic_bool"], + ) assert len(res) == 1 assert res[0]["dynamic_str"] == "dynamic_50" assert res[0]["dynamic_int"] == 5000 assert abs(res[0]["dynamic_float"] - 25.0) < 1e-5 - assert res[0]["dynamic_bool"] == True + assert res[0]["dynamic_bool"] is True # Verify all data count - res, _ = self.query(client, restored_collection_name, - filter="id >= 0", output_fields=["count(*)"]) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) assert res[0]["count(*)"] == 100 # Cleanup @@ -3096,7 +3456,6 @@ class TestMilvusClientSnapshotDataOperationsExtended(TestMilvusClientV2Base): self.drop_collection(client, restored_collection_name) - class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): """ Test concurrent operations for snapshot feature. @@ -3147,8 +3506,7 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): # Others should fail with "already exists" type error for err in errors: - assert "exist" in err.lower() or "duplicate" in err.lower(), \ - f"Unexpected error: {err}" + assert "exist" in err.lower() or "duplicate" in err.lower(), f"Unexpected error: {err}" # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -3168,10 +3526,13 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): rng = np.random.default_rng(seed=19530) # Insert initial data - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(1000)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(1000) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -3183,10 +3544,13 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): nonlocal insert_count batch_id = 0 while not stop_inserting.is_set(): - batch_rows = [{ - default_primary_key_field_name: 10000 + batch_id * 100 + i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(100)] + batch_rows = [ + { + default_primary_key_field_name: 10000 + batch_id * 100 + i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(100) + ] try: self.insert(client, collection_name, batch_rows) insert_count[0] += 100 @@ -3214,8 +3578,7 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): # Restore and verify consistency restored_name = cf.gen_unique_str(prefix + "_restored") - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, source_collection_name=collection_name) wait_for_restore_complete(client, job_id) self.load_collection(client, restored_name) @@ -3229,8 +3592,9 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): # Snapshot should not have more than total inserted at snapshot time # (may have less due to unflushed data) - assert restored_count <= insert_count[0], \ + assert restored_count <= insert_count[0], ( f"Should not exceed total inserted: {restored_count} > {insert_count[0]}" + ) # Cleanup self.drop_snapshot(client, snapshot_name, collection_name) @@ -3249,10 +3613,13 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): self.create_collection(client, collection_name, default_dim) rng = np.random.default_rng(seed=19530) - rows = [{ - default_primary_key_field_name: i, - default_vector_field_name: list(rng.random(default_dim)), - } for i in range(500)] + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] self.insert(client, collection_name, rows) self.flush(client, collection_name) @@ -3266,8 +3633,9 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): def restore_thread(idx): restored_name = cf.gen_unique_str(prefix + f"_concurrent_{idx}") try: - job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, - source_collection_name=collection_name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_name, source_collection_name=collection_name + ) with lock: job_ids.append(job_id) restored_names.append(restored_name) @@ -3294,3 +3662,1902 @@ class TestMilvusClientSnapshotConcurrency(TestMilvusClientV2Base): self.drop_snapshot(client, snapshot_name, collection_name) for name in restored_names: self.drop_collection(client, name) + + +class TestMilvusClientSnapshotLifecycle(TestMilvusClientV2Base): + """ + Test snapshot + collection lifecycle management edge cases. + + Covers race conditions and interactions between snapshot operations + and collection lifecycle operations (drop, rename, cross-db restore). + """ + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_drop_target_collection_during_restore(self): + """ + target: test dropping the target collection while restore is still in progress + method: start restore -> immediately drop the target collection -> check restore state + expected: restore job should eventually fail; no resource leak + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_collection_name = cf.gen_unique_str(prefix + "_restored") + + # 1. Create collection with data and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Start restore + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) + + # 3. Immediately drop the target collection while restore is in progress + try: + self.drop_collection(client, restored_collection_name, timeout=30) + except Exception as e: + log.info(f"Drop target collection during restore: {e}") + + # 4. Wait and check restore state - should eventually reach terminal state + timeout = 120 + start_time = time.time() + final_state = None + while time.time() - start_time < timeout: + state = client.get_restore_snapshot_state(job_id) + final_state = state.state + if final_state in ("RestoreSnapshotCompleted", "RestoreSnapshotFailed"): + break + time.sleep(2) + + log.info(f"Restore final state after dropping target collection: {final_state}") + # The restore should reach a terminal state (not hang forever) + assert final_state in ("RestoreSnapshotCompleted", "RestoreSnapshotFailed"), ( + f"Restore job should reach terminal state, got: {final_state}" + ) + + # 5. Log collection state for diagnostics (existence is timing-dependent) + collections = client.list_collections() + log.info(f"Collections after race (target may or may not exist): {collections}") + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + try: + self.drop_collection(client, restored_collection_name, timeout=30) + except Exception: + pass + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_rename_source_collection(self): + """ + target: test snapshot behavior after renaming the source collection + method: create snapshot -> rename collection -> list/describe/restore snapshot + expected: snapshot should still be usable; describe may show original collection name + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + new_collection_name = cf.gen_unique_str(prefix + "_renamed") + snapshot_name = cf.gen_unique_str(prefix) + restored_collection_name = cf.gen_unique_str(prefix + "_restored") + + # 1. Create collection with data and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Rename source collection + self.rename_collection(client, collection_name, new_collection_name) + + # 3. Verify snapshot is still discoverable + # list_snapshots with old name should fail (collection no longer exists) + error = {ct.err_code: 100, ct.err_msg: "collection not found"} + self.list_snapshots(client, collection_name=collection_name, check_task=CheckTasks.err_res, check_items=error) + + # list_snapshots with new name should find it + snapshots_new, _ = self.list_snapshots(client, collection_name=new_collection_name) + log.info(f"Snapshots listed with new name '{new_collection_name}': {snapshots_new}") + assert snapshot_name in snapshots_new, ( + f"Snapshot {snapshot_name} should be discoverable under new name '{new_collection_name}'" + ) + + # 4. Describe snapshot should still work (use new collection name) + info, _ = self.describe_snapshot(client, snapshot_name, new_collection_name) + assert info.name == snapshot_name + log.info(f"Snapshot collection_name after rename: {info.collection_name}") + + # 5. Restore should still work (snapshot data is independent of collection name) + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=new_collection_name + ) + wait_for_restore_complete(client, job_id) + + self.load_collection(client, restored_collection_name) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb, f"Restored collection should have {default_nb} rows" + + # Cleanup + self.drop_snapshot(client, snapshot_name, new_collection_name) + self.drop_collection(client, new_collection_name) + self.drop_collection(client, restored_collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_create_on_restoring_collection(self): + """ + target: test creating a snapshot on a collection that is being restored into + method: start restore -> immediately create snapshot on the target collection + expected: snapshot creation should either fail or capture incomplete data + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_collection_name = cf.gen_unique_str(prefix + "_restored") + snapshot_on_restored = cf.gen_unique_str(prefix + "_on_restored") + + # 1. Create collection with data and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Start restore + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_collection_name, source_collection_name=collection_name + ) + + # 3. Immediately try to create snapshot on the target collection + snapshot_created = False + try: + self.create_snapshot(client, restored_collection_name, snapshot_on_restored) + snapshot_created = True + log.info("Snapshot on restoring collection succeeded (captured partial state)") + except Exception as e: + log.info(f"Snapshot on restoring collection rejected: {e}") + + # 4. Wait for restore to complete regardless + wait_for_restore_complete(client, job_id, timeout=120) + + # 5. If snapshot was created during restore, verify it captured a subset of data + if snapshot_created: + restored_from_partial = cf.gen_unique_str(prefix + "_from_partial") + job_id2, _ = self.restore_snapshot( + client, snapshot_on_restored, restored_from_partial, source_collection_name=restored_collection_name + ) + wait_for_restore_complete(client, job_id2) + self.load_collection(client, restored_from_partial) + res, _ = self.query(client, restored_from_partial, filter="id >= 0", output_fields=["count(*)"]) + partial_count = res[0]["count(*)"] + log.info(f"Snapshot during restore captured {partial_count} rows (original: {default_nb})") + # Should have at most the original count (might be less if data was still copying) + assert partial_count <= default_nb + self.drop_collection(client, restored_from_partial) + self.drop_snapshot(client, snapshot_on_restored, restored_collection_name) + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, restored_collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_restore_failure_no_resource_leak(self): + """ + target: test that a failed restore does not leak resources + method: restore to an existing collection (will fail) -> verify no leftover resources + expected: restore fails cleanly, no orphan collections or jobs stuck in non-terminal state + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + existing_collection = cf.gen_unique_str(prefix + "_existing") + + # 1. Create collection with data and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Create the target collection so restore will fail (duplicate) + self.create_collection(client, existing_collection, default_dim) + + # 3. Restore to existing collection - should fail + error = {ct.err_code: 65535, ct.err_msg: "duplicate collection"} + self.restore_snapshot( + client, + snapshot_name, + existing_collection, + source_collection_name=collection_name, + check_task=CheckTasks.err_res, + check_items=error, + ) + + # 4. Verify the existing collection is untouched + self.load_collection(client, existing_collection) + res, _ = self.query(client, existing_collection, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == 0, "Existing collection should remain empty (untouched)" + + # 5. Verify snapshot is still usable after failed restore + info, _ = self.describe_snapshot(client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # 6. Successful restore to a new collection proves no state corruption + clean_restored = cf.gen_unique_str(prefix + "_clean") + job_id, _ = self.restore_snapshot(client, snapshot_name, clean_restored, source_collection_name=collection_name) + wait_for_restore_complete(client, job_id) + + self.load_collection(client, clean_restored) + res, _ = self.query(client, clean_restored, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, existing_collection) + self.drop_collection(client, clean_restored) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_concurrent_drop_same_snapshot(self): + """ + target: test concurrent drop of the same snapshot (idempotent) + method: drop the same snapshot from multiple threads simultaneously + expected: all threads should succeed (idempotent behavior), no errors + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + + # 1. Create collection and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(100) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Concurrent drop from multiple threads + results = [] + errors = [] + + def drop_thread(): + try: + client.drop_snapshot(snapshot_name, collection_name) + results.append("success") + except Exception as e: + errors.append(str(e)) + + threads = [threading.Thread(target=drop_thread) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + log.info(f"Concurrent drop results - successes: {len(results)}, errors: {len(errors)}") + log.info(f"Errors: {errors}") + + # At least one thread should succeed; others may succeed (idempotent) or + # hit transient errors (e.g., etcd write conflict under concurrent load) + assert len(results) >= 1, "At least one concurrent drop should succeed, got 0 successes" + assert len(results) + len(errors) == 5, "All 5 threads should have completed" + + # 3. Verify snapshot is gone + snapshots, _ = self.list_snapshots(client, collection_name=collection_name) + assert snapshot_name not in snapshots + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_create_during_drop_source_collection(self): + """ + target: test creating a snapshot while the source collection is being dropped + method: insert data -> flush -> start drop collection and create snapshot concurrently + expected: create snapshot should either succeed (before drop) or fail (after drop); + system should remain consistent + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + + # 1. Create collection with data + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + # 2. Concurrently drop collection and create snapshot + create_result = {"success": False, "error": None} + drop_result = {"success": False, "error": None} + + def create_snapshot_thread(): + try: + client.create_snapshot(collection_name, snapshot_name) + create_result["success"] = True + except Exception as e: + create_result["error"] = str(e) + + def drop_collection_thread(): + try: + client.drop_collection(collection_name) + drop_result["success"] = True + except Exception as e: + drop_result["error"] = str(e) + + t1 = threading.Thread(target=create_snapshot_thread) + t2 = threading.Thread(target=drop_collection_thread) + t1.start() + t2.start() + t1.join() + t2.join() + + log.info(f"Create snapshot: success={create_result['success']}, error={create_result['error']}") + log.info(f"Drop collection: success={drop_result['success']}, error={drop_result['error']}") + + # 3. Verify consistent state + # Note: with cascade delete (PR #48143), DropCollection triggers + # DropSnapshotsByCollection, so even if snapshot creation succeeded + # before drop, the snapshot may be cascade-deleted after drop completes. + if create_result["success"]: + # Snapshot was created before drop took effect; may be cascade-dropped + all_snapshots, _ = self.list_snapshots(client) + log.info(f"Snapshot created, all snapshots after race: {all_snapshots}") + if snapshot_name in all_snapshots: + # Still alive -- restore to verify data integrity + restored_name = cf.gen_unique_str(prefix + "_restored") + job_id, _ = self.restore_snapshot( + client, snapshot_name, restored_name, source_collection_name=collection_name + ) + wait_for_restore_complete(client, job_id) + self.load_collection(client, restored_name) + res, _ = self.query(client, restored_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb + self.drop_collection(client, restored_name) + # Cleanup snapshot + try: + client.drop_snapshot(snapshot_name, collection_name) + except Exception: + pass + else: + log.info("Snapshot was cascade-deleted with the source collection - OK") + else: + # Snapshot creation failed (drop happened first) - this is acceptable + log.info("Snapshot creation failed because collection was dropped first - OK") + + # Drop should always succeed + assert drop_result["success"], f"Drop collection should succeed, error: {drop_result['error']}" + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_restore_cross_database(self): + """ + target: test restoring a snapshot to a different database via db_name param + method: create snapshot in default db -> restore to target db via db_name + expected: restored collection should be created in the target db + note: requires pymilvus >= 2.7.0rc146 (fix: pass database context in snapshot APIs) + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + target_db = cf.gen_unique_str("test_db") + restored_collection_name = cf.gen_unique_str(prefix + "_cross_db") + + # 1. Create collection with data in default db and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Create target database + self.create_database(client, target_db) + + # 3. Restore snapshot to target db via target_db_name param + # SDK signature: restore_snapshot(snapshot_name, source_collection_name, + # target_collection_name, + # source_db_name="", target_db_name="", ...) + job_id = client.restore_snapshot( + snapshot_name, collection_name, restored_collection_name, target_db_name=target_db + ) + wait_for_restore_complete(client, job_id, timeout=120) + + # Note: MilvusClient does not honor per-call db_name kwarg for list/load/ + # query/drop_collection -- we must switch via using_database(). + try: + # 4. Verify collection is in target db + client.using_database(target_db) + target_collections = client.list_collections() + log.info(f"Collections in target db '{target_db}': {target_collections}") + assert restored_collection_name in target_collections, ( + f"Restored collection should be in target db '{target_db}'" + ) + + # 5. Verify collection is NOT in default db + client.using_database("default") + default_collections = client.list_collections() + log.info(f"Collections in default db: {default_collections}") + assert restored_collection_name not in default_collections, ( + "Restored collection should NOT be in default db" + ) + + # 6. Verify data integrity in target db + client.using_database(target_db) + client.load_collection(restored_collection_name) + res = client.query(restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb, f"Restored collection should have {default_nb} rows" + + # Cleanup target-db collection while still in target context + client.drop_collection(restored_collection_name) + finally: + client.using_database("default") + + # Cleanup remaining resources in default db + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, collection_name) + try: + client.drop_database(target_db) + except Exception as e: + log.warning(f"Failed to drop test database '{target_db}': {e}") + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_drop_and_restore_race(self): + """ + target: test race condition between DropSnapshot and RestoreSnapshot + method: start restore and drop snapshot concurrently from different threads + expected: either restore succeeds (drop blocked by ref count) or restore fails + (drop happened before restore registered ref); system should not hang + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_collection_name = cf.gen_unique_str(prefix + "_restored") + + # 1. Create collection with data and snapshot + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Start restore and drop concurrently + restore_result = {"job_id": None, "error": None} + drop_result = {"success": False, "error": None} + + def restore_thread(): + try: + # SDK positional args: (snapshot_name, source_collection_name, target_collection_name) + job_id = client.restore_snapshot(snapshot_name, collection_name, restored_collection_name, timeout=60) + restore_result["job_id"] = job_id + except Exception as e: + restore_result["error"] = str(e) + + def drop_thread(): + try: + client.drop_snapshot(snapshot_name, collection_name, timeout=60) + drop_result["success"] = True + except Exception as e: + drop_result["error"] = str(e) + + t_restore = threading.Thread(target=restore_thread, name="restore_thread") + t_drop = threading.Thread(target=drop_thread, name="drop_thread") + t_restore.start() + t_drop.start() + t_restore.join(timeout=90) + t_drop.join(timeout=90) + assert not t_restore.is_alive(), "restore_thread timed out" + assert not t_drop.is_alive(), "drop_thread timed out" + + log.info(f"Restore: job_id={restore_result['job_id']}, error={restore_result['error']}") + log.info(f"Drop: success={drop_result['success']}, error={drop_result['error']}") + + # 3. Analyze outcomes - two valid scenarios: + # + # Scenario A: Restore registered ref first -> drop blocked -> restore completes + # restore_result["job_id"] is not None, drop_result["error"] contains "is restoring" + # + # Scenario B: Drop succeeded first -> restore fails with "snapshot not found" + # drop_result["success"] is True, restore_result["error"] contains "not found" + # + # Scenario C: Both succeed in sequence (restore ref registered and released quickly) + # Both succeed - rare but possible if restore is very fast + + if restore_result["job_id"] is not None: + # Restore started - wait for it to complete + try: + wait_for_restore_complete(client, restore_result["job_id"], timeout=120) + log.info("Restore completed successfully") + + # Verify data + self.load_collection(client, restored_collection_name) + res, _ = self.query(client, restored_collection_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb + except Exception as e: + log.info(f"Restore ended with: {e}") + + if drop_result["error"]: + # Scenario A: drop was blocked by active pins/restores + # PR #48143 introduced explicit pin-based blocking: + # "active pins exist, unpin before dropping: snapshot is pinned" + log.info(f"Drop was blocked during restore: {drop_result['error']}") + err_lower = drop_result["error"].lower() + assert "pin" in err_lower or "restor" in err_lower, ( + f"Drop error should mention pin/restore, got: {drop_result['error']}" + ) + # Now drop should succeed + self.drop_snapshot(client, snapshot_name, collection_name) + else: + # Scenario C: drop also succeeded (restore was fast) + log.info("Both restore and drop succeeded") + else: + # Scenario B: restore failed (snapshot was dropped first) + log.info(f"Restore failed: {restore_result['error']}") + assert drop_result["success"], "If restore failed, drop should have succeeded" + + # Verify system is in a clean state after race condition: + # The restored collection should either not exist or be droppable + # within a reasonable timeout. If drop_collection hangs or times out, + # it indicates the server is stuck (e.g., broadcaster infinite retry loop). + collections = client.list_collections() + if restored_collection_name in collections: + log.info(f"Restored collection {restored_collection_name} exists, verifying it can be dropped") + self.drop_collection(client, restored_collection_name, timeout=30) + collections_after = client.list_collections() + assert restored_collection_name not in collections_after, ( + f"Restored collection {restored_collection_name} should be droppable after race condition, " + f"but drop_collection did not remove it. Server may be stuck in infinite retry loop." + ) + + # Cleanup snapshot (idempotent) + try: + client.drop_snapshot(snapshot_name, collection_name, timeout=30) + except Exception: + pass + + +class TestMilvusClientSnapshotAlias(TestMilvusClientV2Base): + """ + Test snapshot operations using collection aliases. + + Server resolves aliases via globalMetaCache.GetCollectionID() for + create_snapshot, list_snapshots, and list_restore_snapshot_jobs. + restore_snapshot takes a NEW collection name (not alias) as target. + """ + + def _create_collection_with_data(self, client, collection_name, nb=default_nb): + """Helper: create collection, insert data, flush.""" + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_create_via_alias(self): + """ + target: test creating a snapshot using collection alias instead of collection name + method: create collection -> create alias -> create snapshot via alias + expected: snapshot created successfully, describe shows real collection name + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + alias_name = cf.gen_unique_str(prefix + "_alias") + snapshot_name = cf.gen_unique_str(prefix) + + # 1. Create collection with data + self._create_collection_with_data(client, collection_name) + + # 2. Create alias + self.create_alias(client, collection_name, alias_name) + + # 3. Create snapshot using alias + self.create_snapshot(client, alias_name, snapshot_name) + + # 4. Describe snapshot via alias should show the real collection name + info, _ = self.describe_snapshot(client, snapshot_name, alias_name) + assert info.collection_name == collection_name, ( + f"Expected real collection name '{collection_name}', got '{info.collection_name}'" + ) + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_alias(client, alias_name) + self.drop_collection(client, collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_list_via_alias(self): + """ + target: test listing snapshots using collection alias + method: create snapshot with real name -> list snapshots via alias + expected: list returns the same snapshots as using real name + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + alias_name = cf.gen_unique_str(prefix + "_alias") + snapshot_name = cf.gen_unique_str(prefix) + + # 1. Create collection with data and snapshot + self._create_collection_with_data(client, collection_name) + self.create_alias(client, collection_name, alias_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. List snapshots using alias + snapshots_via_alias, _ = self.list_snapshots(client, collection_name=alias_name) + snapshots_via_name, _ = self.list_snapshots(client, collection_name=collection_name) + + assert snapshot_name in snapshots_via_alias, f"Snapshot not found via alias. Got: {snapshots_via_alias}" + assert snapshots_via_alias == snapshots_via_name, ( + f"Mismatch: via alias={snapshots_via_alias}, via name={snapshots_via_name}" + ) + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_alias(client, alias_name) + self.drop_collection(client, collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_restore_from_alias_created_snapshot(self): + """ + target: test restoring a snapshot that was created via alias + method: create snapshot via alias -> restore -> verify data + expected: restore succeeds with full data integrity + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + alias_name = cf.gen_unique_str(prefix + "_alias") + snapshot_name = cf.gen_unique_str(prefix) + restored_name = cf.gen_unique_str(prefix + "_restored") + + # 1. Create collection with data and alias + self._create_collection_with_data(client, collection_name) + self.create_alias(client, collection_name, alias_name) + + # 2. Create snapshot via alias + self.create_snapshot(client, alias_name, snapshot_name) + + # 3. Restore snapshot using alias as source (server resolves to real collection) + job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, source_collection_name=alias_name) + wait_for_restore_complete(client, job_id) + + # 4. Verify restored data + self.load_collection(client, restored_name) + res, _ = self.query(client, restored_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_alias(client, alias_name) + self.drop_collection(client, collection_name) + self.drop_collection(client, restored_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_list_restore_jobs_via_alias(self): + """ + target: test listing restore snapshot jobs using collection alias + method: restore snapshot to new collection -> create alias on restored + collection -> list restore jobs via alias + expected: list_restore_snapshot_jobs returns correct jobs via alias + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_name = cf.gen_unique_str(prefix + "_restored") + restored_alias = cf.gen_unique_str(prefix + "_restored_alias") + + # 1. Create collection with data and snapshot + self._create_collection_with_data(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Restore to new collection + job_id, _ = self.restore_snapshot(client, snapshot_name, restored_name, source_collection_name=collection_name) + wait_for_restore_complete(client, job_id) + + # 3. Create alias on restored collection and list jobs via alias + self.create_alias(client, restored_name, restored_alias) + jobs, _ = self.list_restore_snapshot_jobs(client, collection_name=restored_alias) + job_ids = [j.job_id for j in jobs] + assert job_id in job_ids, f"Restore job {job_id} not found via alias. Jobs: {job_ids}" + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_alias(client, restored_alias) + self.drop_collection(client, collection_name) + self.drop_collection(client, restored_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_drop_alias_then_create_snapshot(self): + """ + target: test that creating snapshot with dropped alias fails + method: create alias -> drop alias -> create snapshot with dropped alias + expected: create snapshot with dropped alias fails; real name still works + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + alias_name = cf.gen_unique_str(prefix + "_alias") + snapshot_name = cf.gen_unique_str(prefix) + + # 1. Create collection with data and alias + self._create_collection_with_data(client, collection_name) + self.create_alias(client, collection_name, alias_name) + + # 2. Drop alias + self.drop_alias(client, alias_name) + + # 3. Create snapshot should fail with dropped alias + error = {ct.err_code: 100, ct.err_msg: "not found"} + self.create_snapshot(client, alias_name, snapshot_name, check_task=CheckTasks.err_res, check_items=error) + + # 4. Create snapshot with real name should succeed + self.create_snapshot(client, collection_name, snapshot_name) + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_snapshot_alter_alias_then_list_snapshots(self): + """ + target: test that alias retarget affects snapshot listing + method: create alias on col_a -> create snapshot on col_a via alias -> + alter alias to col_b -> list snapshots via alias should show col_b snapshots + expected: alias retarget correctly affects which collection's snapshots are listed + """ + client = self._client() + col_a = cf.gen_collection_name_by_testcase_name() + "_a" + col_b = cf.gen_collection_name_by_testcase_name() + "_b" + alias_name = cf.gen_unique_str(prefix + "_alias") + snapshot_a = cf.gen_unique_str(prefix + "_a") + snapshot_b = cf.gen_unique_str(prefix + "_b") + + # 1. Create two collections with data + self._create_collection_with_data(client, col_a) + self._create_collection_with_data(client, col_b) + + # 2. Create alias pointing to col_a and create snapshot + self.create_alias(client, col_a, alias_name) + self.create_snapshot(client, alias_name, snapshot_a) + + # 3. Alter alias to point to col_b and create snapshot + self.alter_alias(client, col_b, alias_name) + self.create_snapshot(client, alias_name, snapshot_b) + + # 4. List snapshots via alias should show col_b's snapshots + snapshots_via_alias, _ = self.list_snapshots(client, collection_name=alias_name) + assert snapshot_b in snapshots_via_alias, ( + f"Snapshot_b not found via retargeted alias. Got: {snapshots_via_alias}" + ) + assert snapshot_a not in snapshots_via_alias, ( + f"Snapshot_a should not appear after alias retarget. Got: {snapshots_via_alias}" + ) + + # 5. List directly should show each collection's own snapshots + snapshots_a, _ = self.list_snapshots(client, collection_name=col_a) + snapshots_b, _ = self.list_snapshots(client, collection_name=col_b) + assert snapshot_a in snapshots_a + assert snapshot_b in snapshots_b + + # Cleanup + self.drop_snapshot(client, snapshot_a, col_a) + self.drop_snapshot(client, snapshot_b, col_b) + self.drop_alias(client, alias_name) + self.drop_collection(client, col_a) + self.drop_collection(client, col_b) + + @pytest.mark.tags(CaseLabel.L2) + def test_restore_target_name_equals_existing_alias_fails(self): + """ + target: test restoring a snapshot with target_collection_name equal to + an existing alias should fail (alias and collection share a namespace) + method: create col_src + snapshot -> create alias A pointing to col_src + -> restore snapshot to target_collection_name=A + expected: restore is synchronously rejected with an alias-conflict error; + source collection, snapshot, and alias all remain intact + note: the rejection happens in datacoord's broker.CreateCollection path + during RestoreCollection (snapshot_manager.go:833) + """ + client = self._client() + col_src = cf.gen_collection_name_by_testcase_name() + alias_name = cf.gen_unique_str(prefix + "_alias") + snapshot_name = cf.gen_unique_str(prefix) + + # 1. Create source collection + snapshot + alias + self._create_collection_with_data(client, col_src) + self.create_snapshot(client, col_src, snapshot_name) + self.create_alias(client, col_src, alias_name) + + # 2. Restore with target_collection_name = existing alias name must fail + error = {ct.err_code: 65535, ct.err_msg: "conflicts with an existing alias"} + self.restore_snapshot( + client, + snapshot_name, + alias_name, + source_collection_name=col_src, + check_task=CheckTasks.err_res, + check_items=error, + ) + + # 3. Verify source, snapshot, and alias are all untouched + snapshots, _ = self.list_snapshots(client, collection_name=col_src) + assert snapshot_name in snapshots + + # A restore succeeds with a fresh, non-conflicting target name — proves + # the alias-conflict was a clean rejection, not a corrupted state. + clean_target = cf.gen_unique_str(prefix + "_clean") + job_id, _ = self.restore_snapshot(client, snapshot_name, clean_target, source_collection_name=col_src) + wait_for_restore_complete(client, job_id) + + # Cleanup + self.drop_snapshot(client, snapshot_name, col_src) + self.drop_alias(client, alias_name) + self.drop_collection(client, col_src) + self.drop_collection(client, clean_target) + + +@pytest.mark.tags(CaseLabel.RBAC) +class TestMilvusClientSnapshotRbac(TestMilvusClientV2Base): + """ + Test RBAC v2 privilege enforcement for snapshot operations. + """ + + user_pre = "snap_user" + role_pre = "snap_role" + + def teardown_method(self, method): + """Clean up users, roles, snapshots and collections created during test.""" + log.info("[snapshot_rbac_teardown] Start teardown ...") + client = self._client() + + # drop all non-root users + users, _ = self.list_users(client) + for user in users: + if user != ct.default_user: + self.drop_user(client, user) + + # revoke privileges and drop non-builtin roles + # Must use revoke_privilege_v2 for privileges granted via v2 API, + # because v2-granted privileges carry db_name="*" which v1 revoke cannot match. + roles, _ = self.list_roles(client) + for role in roles: + if role not in ["admin", "public"]: + res, _ = self.describe_role(client, role) + if res and res.get("privileges"): + for privilege in res["privileges"]: + self.revoke_privilege_v2( + client, + role, + privilege["privilege"], + privilege.get("object_name", "*"), + privilege.get("db_name", "*"), + ) + self.drop_role(client, role) + + super().teardown_method(method) + + def _setup_user_with_role(self, root_client, host, port): + """Helper: create a user + role, assign role, return (user_client, role_name).""" + user_name = cf.gen_unique_str(self.user_pre) + role_name = cf.gen_unique_str(self.role_pre) + password = cf.gen_str_by_length(contain_numbers=True) + self.create_user(root_client, user_name=user_name, password=password) + self.create_role(root_client, role_name=role_name) + self.grant_role(root_client, user_name=user_name, role_name=role_name) + + uri = f"http://{host}:{port}" + user_client, _ = self.init_milvus_client(uri=uri, user=user_name, password=password) + return user_client, role_name + + def _prepare_collection_with_snapshot(self, client, nb=500): + """Helper: create collection, insert data, flush, create snapshot. Return (col_name, snap_name).""" + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + return collection_name, snapshot_name + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_create_denied_without_privilege(self, host, port): + """ + target: verify CreateSnapshot is denied without privilege + method: create user with empty role, attempt create_snapshot + expected: permission denied + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + self.create_collection(client, collection_name, default_dim) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # should be denied + snapshot_name = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, snapshot_name, check_task=CheckTasks.check_permission_deny) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_create_allowed_after_grant_v2(self, host, port): + """ + target: verify CreateSnapshot succeeds after granting privilege via v2 API + method: grant_privilege_v2 CreateSnapshot to role, then create snapshot + expected: create snapshot succeeds + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant Collection-level CreateSnapshot via v2 API (snapshot privileges + # moved from Global to Collection level in PR #48143) + self.grant_privilege_v2(client, role_name, "CreateSnapshot", "*", "*") + time.sleep(10) + + # should succeed + snapshot_name = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, snapshot_name) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_drop_denied_without_privilege(self, host, port): + """ + target: verify DropSnapshot is denied without privilege + method: create snapshot as root, attempt drop as unprivileged user + expected: permission denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # should be denied + self.drop_snapshot(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_drop_allowed_after_grant_v2(self, host, port): + """ + target: verify DropSnapshot succeeds after granting privilege via v2 API + method: grant_privilege_v2 DropSnapshot to role, then drop snapshot + expected: drop snapshot succeeds + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant via v2 API + self.grant_privilege_v2(client, role_name, "DropSnapshot", "*", "*") + time.sleep(10) + + # should succeed + self.drop_snapshot(user_client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_list_denied_without_privilege(self, host, port): + """ + target: verify ListSnapshots is denied without privilege + method: create snapshot as root, attempt list as unprivileged user + expected: permission denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # should be denied + self.list_snapshots(user_client, collection_name=collection_name, check_task=CheckTasks.check_permission_deny) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_list_allowed_after_grant_v2(self, host, port): + """ + target: verify ListSnapshots succeeds after granting privilege via v2 API + method: grant_privilege_v2 ListSnapshots to role, then list snapshots + expected: list snapshots succeeds and returns the snapshot + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant via v2 API + self.grant_privilege_v2(client, role_name, "ListSnapshots", "*", "*") + time.sleep(10) + + # should succeed + snapshots, _ = self.list_snapshots(user_client, collection_name=collection_name) + assert snapshot_name in snapshots + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_describe_denied_without_privilege(self, host, port): + """ + target: verify DescribeSnapshot is denied without privilege + method: create snapshot as root, attempt describe as unprivileged user + expected: permission denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # should be denied + self.describe_snapshot(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_describe_allowed_after_grant_v2(self, host, port): + """ + target: verify DescribeSnapshot succeeds after granting privilege via v2 API + method: grant_privilege_v2 DescribeSnapshot to role, then describe snapshot + expected: describe snapshot succeeds with correct info + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant via v2 API + self.grant_privilege_v2(client, role_name, "DescribeSnapshot", "*", "*") + time.sleep(10) + + # should succeed + info, _ = self.describe_snapshot(user_client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_restore_denied_without_privilege(self, host, port): + """ + target: verify RestoreSnapshot is denied without privilege + method: create snapshot as root, attempt restore as unprivileged user + expected: permission denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # should be denied + restored_name = cf.gen_unique_str(prefix + "_restored") + self.restore_snapshot( + user_client, + snapshot_name, + restored_name, + source_collection_name=collection_name, + check_task=CheckTasks.check_permission_deny, + ) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_restore_allowed_after_grant_v2(self, host, port): + """ + target: verify RestoreSnapshot succeeds after granting privilege via v2 API + method: grant_privilege_v2 RestoreSnapshot to role, then restore snapshot + expected: restore snapshot succeeds + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant via v2 API + self.grant_privilege_v2(client, role_name, "RestoreSnapshot", "*", "*") + time.sleep(10) + + # should succeed + restored_name = cf.gen_unique_str(prefix + "_restored") + job_id, _ = self.restore_snapshot( + user_client, snapshot_name, restored_name, source_collection_name=collection_name + ) + wait_for_restore_complete(client, job_id) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, restored_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_revoke_privilege_v2_then_denied(self, host, port): + """ + target: verify operation is denied after revoking privilege via v2 API + method: grant_privilege_v2 CreateSnapshot -> verify allowed -> revoke_privilege_v2 -> verify denied + expected: permission denied after revocation + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant via v2 and verify allowed + self.grant_privilege_v2(client, role_name, "CreateSnapshot", "*", "*") + time.sleep(10) + snapshot_name = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, snapshot_name) + self.drop_snapshot(client, snapshot_name, collection_name) + + # revoke via v2 and verify denied + self.revoke_privilege_v2(client, role_name, "CreateSnapshot", "*", "*") + time.sleep(10) + snapshot_name2 = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, snapshot_name2, check_task=CheckTasks.check_permission_deny) + + @pytest.mark.tags(CaseLabel.RBAC) + @pytest.mark.xfail( + reason="Built-in CollectionReadOnly group on deployed server does " + "not yet include DescribeSnapshot/ListSnapshots even though " + "util.CollectionReadOnlyPrivileges defines them. Tracking: " + "https://github.com/milvus-io/milvus/issues/47855" + ) + def test_snapshot_v2_privilege_group_collection_readonly(self, host, port): + """ + target: verify CollectionReadOnly v2 privilege group grants read snapshot ops + method: grant CollectionReadOnly, attempt describe and list + expected: describe/list succeed, create/drop/restore denied + note: snapshot privileges moved from Global (ClusterXxx) to Collection + level groups in PR #48143 (fixes #47855) + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant CollectionReadOnly (includes DescribeSnapshot + ListSnapshots) + self.grant_privilege_v2(client, role_name, "CollectionReadOnly", "*", "*") + time.sleep(10) + + # read ops should succeed + snapshots, _ = self.list_snapshots(user_client, collection_name=collection_name) + assert snapshot_name in snapshots + + info, _ = self.describe_snapshot(user_client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # write ops should be denied + new_snap = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, new_snap, check_task=CheckTasks.check_permission_deny) + self.drop_snapshot(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + restored_name = cf.gen_unique_str(prefix + "_restored") + self.restore_snapshot( + user_client, + snapshot_name, + restored_name, + source_collection_name=collection_name, + check_task=CheckTasks.check_permission_deny, + ) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + @pytest.mark.xfail( + reason="Built-in CollectionReadWrite group on deployed server does " + "not yet include CreateSnapshot/DropSnapshot. Tracking: " + "https://github.com/milvus-io/milvus/issues/47855" + ) + def test_snapshot_v2_privilege_group_collection_readwrite(self, host, port): + """ + target: verify CollectionReadWrite v2 privilege group grants CRUD snapshot ops + method: grant CollectionReadWrite, attempt create/drop/describe/list + expected: create/drop/describe/list succeed, restore denied + note: CollectionReadWrite = CollectionReadOnly + {CreateSnapshot, DropSnapshot}. + RestoreSnapshot is only granted by CollectionAdmin. + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant CollectionReadWrite (adds CreateSnapshot/DropSnapshot) + self.grant_privilege_v2(client, role_name, "CollectionReadWrite", "*", "*") + time.sleep(10) + + # create/list/describe/drop should succeed + snapshot_name = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, snapshot_name) + + snapshots, _ = self.list_snapshots(user_client, collection_name=collection_name) + assert snapshot_name in snapshots + + info, _ = self.describe_snapshot(user_client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # restore should still be denied (not in ReadWrite group) + restored_name = cf.gen_unique_str(prefix + "_restored") + self.restore_snapshot( + user_client, + snapshot_name, + restored_name, + source_collection_name=collection_name, + check_task=CheckTasks.check_permission_deny, + ) + + # drop should succeed + self.drop_snapshot(user_client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + @pytest.mark.xfail( + reason="Built-in CollectionAdmin group on deployed server does not " + "yet include Create/Drop/RestoreSnapshot. Tracking: " + "https://github.com/milvus-io/milvus/issues/47855" + ) + def test_snapshot_v2_privilege_group_collection_admin(self, host, port): + """ + target: verify CollectionAdmin v2 privilege group grants all snapshot ops + method: grant CollectionAdmin, attempt all snapshot ops + expected: all snapshot ops succeed + note: CollectionAdmin = CollectionReadWrite + {RestoreSnapshot, ...}. + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant CollectionAdmin - the highest collection-level privilege group + self.grant_privilege_v2(client, role_name, "CollectionAdmin", "*", "*") + time.sleep(10) + + # all snapshot ops should succeed + new_snap = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, new_snap) + + snapshots, _ = self.list_snapshots(user_client, collection_name=collection_name) + assert new_snap in snapshots + + info, _ = self.describe_snapshot(user_client, new_snap, collection_name) + assert info.name == new_snap + + restored_name = cf.gen_unique_str(prefix + "_restored") + job_id, _ = self.restore_snapshot(user_client, new_snap, restored_name, source_collection_name=collection_name) + wait_for_restore_complete(client, job_id) + + self.drop_snapshot(user_client, new_snap, collection_name) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, restored_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_admin_role_has_full_access(self, host, port): + """ + target: verify built-in admin role has full snapshot access + method: create user, assign admin role, test all snapshot ops + expected: all operations succeed + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + # create user with admin role + user_name = cf.gen_unique_str(self.user_pre) + password = cf.gen_str_by_length(contain_numbers=True) + self.create_user(client, user_name=user_name, password=password) + self.grant_role(client, user_name=user_name, role_name="admin") + + uri = f"http://{host}:{port}" + admin_client, _ = self.init_milvus_client(uri=uri, user=user_name, password=password) + + # all ops should succeed + snapshot_name = cf.gen_unique_str(prefix) + self.create_snapshot(admin_client, collection_name, snapshot_name) + + snapshots, _ = self.list_snapshots(admin_client, collection_name=collection_name) + assert snapshot_name in snapshots + + info, _ = self.describe_snapshot(admin_client, snapshot_name, collection_name) + assert info.name == snapshot_name + + restored_name = cf.gen_unique_str(prefix + "_restored") + job_id, _ = self.restore_snapshot( + admin_client, snapshot_name, restored_name, source_collection_name=collection_name + ) + wait_for_restore_complete(client, job_id) + + self.drop_snapshot(admin_client, snapshot_name, collection_name) + self.drop_collection(client, restored_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_multiple_privileges_granular_v2(self, host, port): + """ + target: verify granular privilege combination works correctly via v2 API + method: grant only ListSnapshots + DescribeSnapshot via v2, verify create/drop/restore denied + expected: only granted ops succeed, others denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant only read-related privileges via v2 API + self.grant_privilege_v2(client, role_name, "ListSnapshots", "*", "*") + self.grant_privilege_v2(client, role_name, "DescribeSnapshot", "*", "*") + time.sleep(10) + + # read ops should succeed + snapshots, _ = self.list_snapshots(user_client, collection_name=collection_name) + assert snapshot_name in snapshots + + info, _ = self.describe_snapshot(user_client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # write ops should be denied + new_snap = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, new_snap, check_task=CheckTasks.check_permission_deny) + self.drop_snapshot(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + restored_name = cf.gen_unique_str(prefix + "_restored") + self.restore_snapshot( + user_client, + snapshot_name, + restored_name, + source_collection_name=collection_name, + check_task=CheckTasks.check_permission_deny, + ) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_public_role_has_no_access(self, host, port): + """ + target: verify public role has no snapshot privileges by default + method: create user with only public role, attempt all snapshot ops + expected: all operations denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + # create user without any custom role (only default public role) + user_name = cf.gen_unique_str(self.user_pre) + password = cf.gen_str_by_length(contain_numbers=True) + self.create_user(client, user_name=user_name, password=password) + + uri = f"http://{host}:{port}" + user_client, _ = self.init_milvus_client(uri=uri, user=user_name, password=password) + + # all ops should be denied + new_snap = cf.gen_unique_str(prefix) + self.create_snapshot(user_client, collection_name, new_snap, check_task=CheckTasks.check_permission_deny) + self.list_snapshots(user_client, collection_name=collection_name, check_task=CheckTasks.check_permission_deny) + self.describe_snapshot(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + self.drop_snapshot(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + restored_name = cf.gen_unique_str(prefix + "_restored") + self.restore_snapshot( + user_client, + snapshot_name, + restored_name, + source_collection_name=collection_name, + check_task=CheckTasks.check_permission_deny, + ) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_pin_denied_without_privilege(self, host, port): + """ + target: verify PinSnapshotData is denied without the privilege + method: create snapshot as root; attempt pin as unprivileged user + expected: permission denied + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, _ = self._setup_user_with_role(client, host, port) + + self.pin_snapshot_data(user_client, snapshot_name, collection_name, check_task=CheckTasks.check_permission_deny) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_pin_allowed_after_grant_v2(self, host, port): + """ + target: verify PinSnapshotData succeeds after granting the privilege via v2 API + method: grant PinSnapshotData to role; attempt pin; immediately unpin as root + expected: pin succeeds with a non-zero pin_id + note: PinSnapshotData is a Global-level privilege + (pkg/util/constant.go:122-124) + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + self.grant_privilege_v2(client, role_name, "PinSnapshotData", "*", "*") + time.sleep(10) + + pin_id, _ = self.pin_snapshot_data(user_client, snapshot_name, collection_name, ttl_seconds=60) + assert isinstance(pin_id, int) and pin_id > 0, ( + f"pin_snapshot_data should return a positive int pin_id, got {pin_id!r}" + ) + + # Clean up with root (unpin doesn't need user privilege here) + try: + client.unpin_snapshot_data(pin_id) + except Exception as e: + log.warning(f"cleanup unpin failed, relying on TTL: {e}") + + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_snapshot_unpin_denied_without_privilege(self, host, port): + """ + target: verify UnpinSnapshotData is denied without the privilege + method: unprivileged user attempts unpin with an arbitrary pin_id + expected: permission denied (privilege check happens before pin lookup) + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, _ = self._setup_user_with_role(client, host, port) + + # Use an arbitrary pin_id — server must reject on privilege, not on lookup + self.unpin_snapshot_data(user_client, pin_id=123456789, check_task=CheckTasks.check_permission_deny) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.RBAC) + def test_restore_revoke_privilege_v2_then_denied(self, host, port): + """ + target: verify restore is denied after revoking RestoreSnapshot via v2 API + method: grant -> verify restore succeeds -> revoke -> verify restore denied + expected: permission denied after revocation + """ + client = self._client() + collection_name, snapshot_name = self._prepare_collection_with_snapshot(client) + + user_client, role_name = self._setup_user_with_role(client, host, port) + + # grant and verify allowed + self.grant_privilege_v2(client, role_name, "RestoreSnapshot", "*", "*") + time.sleep(10) + restored_ok = cf.gen_unique_str(prefix + "_ok") + job_id, _ = self.restore_snapshot( + user_client, snapshot_name, restored_ok, source_collection_name=collection_name + ) + wait_for_restore_complete(client, job_id) + + # revoke and verify denied + self.revoke_privilege_v2(client, role_name, "RestoreSnapshot", "*", "*") + time.sleep(10) + restored_denied = cf.gen_unique_str(prefix + "_denied") + self.restore_snapshot( + user_client, + snapshot_name, + restored_denied, + source_collection_name=collection_name, + check_task=CheckTasks.check_permission_deny, + ) + + # cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, restored_ok) + + +class TestMilvusClientSnapshotCreateParams(TestMilvusClientV2Base): + """Test create_snapshot parameter handling beyond basic lifecycle. + + Focus on the ``compaction_protection_seconds`` option introduced with + the snapshot feature (see ``internal/proxy/task_snapshot.go:118-126``). + """ + + @pytest.mark.tags(CaseLabel.L1) + def test_create_snapshot_with_compaction_protection_seconds(self): + """ + target: test create_snapshot accepts a positive compaction_protection_seconds + method: create snapshot with compaction_protection_seconds=3600 + expected: snapshot is created and subsequent describe/list succeed + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(500) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + # Call the SDK directly — wrapper forces description as kwarg ordering + client.create_snapshot(snapshot_name, collection_name, compaction_protection_seconds=3600) + + info, _ = self.describe_snapshot(client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.L1) + def test_create_snapshot_compaction_protection_negative(self): + """ + target: test create_snapshot rejects negative compaction_protection_seconds + method: pass compaction_protection_seconds=-1 + expected: server raises ParameterInvalid with "non-negative" in message + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + + self.create_collection(client, collection_name, default_dim) + + with pytest.raises(Exception) as exc_info: + client.create_snapshot(snapshot_name, collection_name, compaction_protection_seconds=-1) + msg = str(exc_info.value).lower() + assert "non-negative" in msg or "compaction_protection_seconds" in msg, ( + f"Expected compaction_protection error, got: {exc_info.value}" + ) + + @pytest.mark.tags(CaseLabel.L1) + def test_create_snapshot_compaction_protection_exceeds_max(self): + """ + target: test create_snapshot rejects compaction_protection_seconds > server max + method: pass compaction_protection_seconds well above the default 604800s cap + expected: server raises ParameterInvalid with "must not exceed" in message + note: default max is 604800s (7 days) per dataCoord.snapshot.maxCompactionProtectionSeconds + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + + self.create_collection(client, collection_name, default_dim) + + with pytest.raises(Exception) as exc_info: + client.create_snapshot(snapshot_name, collection_name, compaction_protection_seconds=999_999_999) + msg = str(exc_info.value).lower() + assert "not exceed" in msg or "exceeds" in msg or "compaction_protection_seconds" in msg, ( + f"Expected exceeds-max error, got: {exc_info.value}" + ) + + +class TestMilvusClientSnapshotRestoreParams(TestMilvusClientV2Base): + """Test restore_snapshot parameter handling (cross-db source, etc.).""" + + @pytest.mark.tags(CaseLabel.L2) + def test_restore_snapshot_source_db_name_explicit(self): + """ + target: test restore_snapshot honors explicit source_db_name + method: create snapshot in DB X; from default-db context call restore + with source_db_name=X and target_db_name="default" + expected: restore completes; target collection is created in default db + with identical row count; source remains in DB X untouched + note: complements the existing cross-db test which only exercises + target_db_name. This test exercises source_db_name explicitly. + """ + client = self._client() + source_db = cf.gen_unique_str("test_src_db") + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + restored_name = cf.gen_unique_str(prefix + "_restored") + + # 1. Create source db and populate collection + snapshot inside it + self.create_database(client, source_db) + try: + client.using_database(source_db) + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(default_nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + + # 2. Switch to default db and restore via explicit source_db_name + client.using_database("default") + job_id = client.restore_snapshot( + snapshot_name, collection_name, restored_name, source_db_name=source_db, target_db_name="default" + ) + wait_for_restore_complete(client, job_id, timeout=120) + + # 3. Target collection lives in default db + default_collections = client.list_collections() + assert restored_name in default_collections, ( + f"Restored collection should be in default db, got {default_collections}" + ) + + client.load_collection(restored_name) + res = client.query(restored_name, filter="id >= 0", output_fields=["count(*)"]) + assert res[0]["count(*)"] == default_nb, f"Restored collection should have {default_nb} rows" + + # 4. Source still exists in source_db + client.using_database(source_db) + source_collections = client.list_collections() + assert collection_name in source_collections, ( + f"Source collection missing in {source_db}: {source_collections}" + ) + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + self.drop_collection(client, collection_name) + client.using_database("default") + self.drop_collection(client, restored_name) + finally: + client.using_database("default") + try: + client.drop_database(source_db) + except Exception as e: + log.warning(f"Failed to drop source db '{source_db}': {e}") + + +class TestMilvusClientSnapshotPin(TestMilvusClientV2Base): + """Test pin_snapshot_data / unpin_snapshot_data. + + These APIs exist in both the SDK (``pymilvus/milvus_client/milvus_client.py``) + and the server (``internal/proxy/task_snapshot.go:839-1029``) but are not + exercised by existing tests. They are the admin-facing hooks for holding + snapshot segments against compaction/GC during out-of-band copy-out. + """ + + def _prepare(self, client, nb=500): + """Helper: create collection + snapshot, return (collection, snapshot).""" + collection_name = cf.gen_collection_name_by_testcase_name() + snapshot_name = cf.gen_unique_str(prefix) + self.create_collection(client, collection_name, default_dim) + rng = np.random.default_rng(seed=19530) + rows = [ + { + default_primary_key_field_name: i, + default_vector_field_name: list(rng.random(default_dim)), + } + for i in range(nb) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + self.create_snapshot(client, collection_name, snapshot_name) + return collection_name, snapshot_name + + @pytest.mark.tags(CaseLabel.L1) + def test_pin_snapshot_data_basic(self): + """ + target: test basic pin → unpin flow + method: pin with ttl=60; assert pin_id > 0; unpin; drop snapshot still works + expected: pin returns a positive int pin_id; unpin is side-effect free + """ + client = self._client() + collection_name, snapshot_name = self._prepare(client) + + pin_id, _ = self.pin_snapshot_data(client, snapshot_name, collection_name, ttl_seconds=60) + assert isinstance(pin_id, int) and pin_id > 0, ( + f"pin_snapshot_data should return a positive pin_id, got {pin_id!r}" + ) + + self.unpin_snapshot_data(client, pin_id) + + # Snapshot must still be intact after pin/unpin cycle + info, _ = self.describe_snapshot(client, snapshot_name, collection_name) + assert info.name == snapshot_name + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.L2) + def test_pin_snapshot_blocks_drop(self): + """ + target: test that a pin blocks drop_snapshot until unpinned + method: pin with ttl=300s; attempt drop → expect failure; unpin → drop succeeds + expected: drop fails while pin is active; error mentions pin; drop works after unpin + note: mirrors the pin-based protection exercised indirectly by + ``test_snapshot_drop_and_restore_race`` in TestMilvusClientSnapshotLifecycle. + """ + client = self._client() + collection_name, snapshot_name = self._prepare(client) + + pin_id, _ = self.pin_snapshot_data(client, snapshot_name, collection_name, ttl_seconds=300) + + # drop should be rejected while pinned + with pytest.raises(Exception) as exc_info: + client.drop_snapshot(snapshot_name, collection_name, timeout=30) + err_msg = str(exc_info.value).lower() + assert "pin" in err_msg, f"Expected drop error to mention pin, got: {exc_info.value}" + + # Unpin releases the hold + self.unpin_snapshot_data(client, pin_id) + + # drop should now succeed + self.drop_snapshot(client, snapshot_name, collection_name) + + snapshots, _ = self.list_snapshots(client, collection_name=collection_name) + assert snapshot_name not in snapshots + + @pytest.mark.tags(CaseLabel.L1) + def test_pin_snapshot_invalid_ttl_negative(self): + """ + target: test pin rejects negative ttl_seconds + method: call pin with ttl_seconds=-1 + expected: server returns ParameterInvalid with "non-negative" + """ + client = self._client() + collection_name, snapshot_name = self._prepare(client) + + with pytest.raises(Exception) as exc_info: + client.pin_snapshot_data(snapshot_name, collection_name, ttl_seconds=-1) + msg = str(exc_info.value).lower() + assert "non-negative" in msg or "ttl_seconds" in msg, f"Expected ttl_seconds error, got: {exc_info.value}" + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.L1) + def test_pin_snapshot_invalid_ttl_exceeds_max(self): + """ + target: test pin rejects ttl_seconds beyond the 30-day cap + method: call pin with ttl_seconds > 2592000 (30 days) + expected: server returns ParameterInvalid with "exceeds maximum" + note: cap defined in internal/proxy/task_snapshot.go:891 (maxPinTTLSeconds) + """ + client = self._client() + collection_name, snapshot_name = self._prepare(client) + + with pytest.raises(Exception) as exc_info: + client.pin_snapshot_data(snapshot_name, collection_name, ttl_seconds=2_592_001) + msg = str(exc_info.value).lower() + assert "exceed" in msg or "ttl_seconds" in msg, f"Expected exceeds-max error, got: {exc_info.value}" + + # Cleanup + self.drop_snapshot(client, snapshot_name, collection_name) + + @pytest.mark.tags(CaseLabel.L1) + def test_pin_nonexistent_snapshot(self): + """ + target: test pin on a non-existent snapshot fails cleanly + method: pin a snapshot_name that was never created + expected: server returns a clear error (typically not found / snapshot metadata missing) + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + self.create_collection(client, collection_name, default_dim) + + with pytest.raises(Exception) as exc_info: + client.pin_snapshot_data(cf.gen_unique_str("ghost"), collection_name, ttl_seconds=60) + # Accept any error — the server may return "not found", "snapshot", or ParameterInvalid + log.info(f"pin non-existent snapshot error: {exc_info.value}") + assert str(exc_info.value), "pin_snapshot_data on non-existent snapshot must raise" + + @pytest.mark.tags(CaseLabel.L2) + def test_unpin_invalid_pin_id(self): + """ + target: test unpin with an unknown / never-issued pin_id + method: call unpin with a random int not produced by pin_snapshot_data + expected: server either ignores idempotently or returns a clear error; + no hang, no system inconsistency + """ + client = self._client() + + # Either the call raises a clear error, or it silently no-ops. + # Both are acceptable — the key is it must not hang or corrupt state. + try: + client.unpin_snapshot_data(pin_id=987_654_321, timeout=30) + log.info("unpin with unknown pin_id was idempotent (no error)") + except Exception as e: + log.info(f"unpin with unknown pin_id rejected: {e}") + assert str(e), "unpin error, if any, must carry a message"