diff --git a/examples/async_DDL_test.py b/examples/async_DDL_test.py index d0db8282d..646dbc3a0 100644 --- a/examples/async_DDL_test.py +++ b/examples/async_DDL_test.py @@ -1,4 +1,5 @@ import asyncio +import time import numpy as np from pymilvus import AsyncMilvusClient, DataType @@ -11,6 +12,8 @@ test_db = "test_database" partition_name = "test_partition" group_name = "test_privilege_group" +rg1_name = "test_resource_group" +rg2_name = "test_resource_group" async def create_resources(client): print(fmt.format("Creating Resources")) @@ -19,10 +22,9 @@ async def create_resources(client): await client.create_database(test_db) print(f"Database {test_db} created") - print("Creating user and role...") + print("Creating user...") await client.create_user(test_user, "password123") - await client.create_role(test_role) - print(f"User {test_user} and role {test_role} created") + print(f"User {test_user} created") print("Creating privilege group...") await client.create_privilege_group(group_name) @@ -52,6 +54,12 @@ async def create_resources(client): print("Creating alias...") await client.create_alias(collection_name, test_alias) print(f"Alias {test_alias} created") + + print("Creating resource groups...") + await client.create_resource_group(rg1_name) + print(f"Resource group {rg1_name} created") + await client.create_resource_group(rg2_name) + print(f"Resource group {rg2_name} created") print("Loading collection...") await client.load_collection(collection_name) @@ -112,14 +120,6 @@ async def test_functionality(client): await client.update_password(test_user, "password123", "newpassword123") print("Password updated") - print("Testing role management...") - print(f"list_roles: {await client.list_roles()}") - print(f"describe_role: {await client.describe_role(test_role)}") - - await client.grant_role(test_user, test_role) - await client.revoke_role(test_user, test_role) - print("Role grant/revoke test completed") - print("Testing alias operations...") print(f"describe_alias: {await client.describe_alias(test_alias)}") print(f"list_aliases: {await client.list_aliases(collection_name)}") @@ -181,6 +181,36 @@ async def test_functionality(client): await client.remove_privileges_from_group(group_name, ["Delete"]) print("Privilege group operations completed") + print(fmt.format("Testing create_field_schema")) + field1 = client.create_field_schema("test_id", DataType.INT64, desc="test int64 field", is_primary=True, auto_id=True) + print(f"Field1: {field1}") + field2 = client.create_field_schema("test_vector", DataType.FLOAT_VECTOR, desc="test vector field", dim=dim) + print(f"Field2: {field2}") + field3 = client.create_field_schema("test_text", DataType.VARCHAR, desc="test varchar field", max_length=256) + print(f"Field3: {field3}") + del field1 + del field2 + del field3 + print("create_field_schema test completed") + + print(fmt.format("Testing Resource Group Operations")) + + # List resource groups + print("\nListing resource groups...") + rgs = await client.list_resource_groups() + print(f"Resource groups: {rgs}") + assert rg1_name in rgs + assert rg2_name in rgs + + # Describe resource groups + print(f"\nDescribing resource group {rg1_name}...") + rg1_info = await client.describe_resource_group(rg1_name) + print(f"Resource group {rg1_name} info: {rg1_info}") + + print(f"\nDescribing resource group {rg2_name}...") + rg2_info = await client.describe_resource_group(rg2_name) + print(f"Resource group {rg2_name} info: {rg2_info}") + async def cleanup_resources(client): print(fmt.format("Cleaning Up Resources")) @@ -192,20 +222,33 @@ async def cleanup_resources(client): await client.drop_collection(collection_name) print(f"Collection {collection_name} dropped") - print("Dropping user and role...") + print("Dropping user...") await client.drop_user(test_user) - await client.drop_role(test_role) - print(f"User {test_user} and role {test_role} dropped") + print(f"User {test_user} dropped") print("Dropping privilege group...") await client.drop_privilege_group(group_name) print(f"Privilege group {group_name} dropped") + print("Dropping database...") + await client.drop_database(test_db) + print(f"Database {test_db} dropped") + + print("\nDropping resource groups...") + if rg1_name in await client.list_resource_groups(): + await client.drop_resource_group(rg1_name) + print(f"Resource group {rg1_name} dropped") + if rg2_name in await client.list_resource_groups(): + await client.drop_resource_group(rg2_name) + print(f"Resource group {rg2_name} dropped") + async def main(): client = AsyncMilvusClient("http://localhost:19530") await create_resources(client) await test_functionality(client) + time.sleep(10) await cleanup_resources(client) + time.sleep(10) await client.close() print(fmt.format("Test Completed")) diff --git a/examples/async_previledge_test.py b/examples/async_previledge_test.py new file mode 100644 index 000000000..cfebcebf8 --- /dev/null +++ b/examples/async_previledge_test.py @@ -0,0 +1,212 @@ +import asyncio + +from pymilvus import AsyncMilvusClient, DataType + +# Super user credentials for admin operations +super_user = "" +super_password = "" + +fmt = "\n=== {:30} ===\n" +collection_name = "test_privilege_collection" +test_user = "test_privilege_user" +test_role_v1 = "test_role_v1" +test_role_v2 = "test_role_v2" +test_db = "test_privilege_database" +privilege_group_name = "test_privilege_group" + +# Test data for privilege operations +db_rw_privileges = [ + {"object_type": "Collection", "privilege": "Insert", "object_name": collection_name}, + {"object_type": "Collection", "privilege": "Delete", "object_name": collection_name}, + {"object_type": "Collection", "privilege": "Search", "object_name": collection_name}, + {"object_type": "Collection", "privilege": "Query", "object_name": collection_name}, +] + +db_ro_privileges = [ + {"object_type": "Collection", "privilege": "Search", "object_name": collection_name}, + {"object_type": "Collection", "privilege": "Query", "object_name": collection_name}, +] + + +async def create_resources(client: AsyncMilvusClient): + print(fmt.format("Creating Resources")) + + print("Creating database...") + await client.create_database(test_db) + print(f"Database {test_db} created") + + print("Creating user and roles...") + await client.create_user(test_user, "password123") + await client.create_role(test_role_v1) + await client.create_role(test_role_v2) + print(f"User {test_user} and roles {test_role_v1}, {test_role_v2} created") + + print("Creating privilege group...") + await client.create_privilege_group(privilege_group_name) + print(f"Privilege group {privilege_group_name} created") + + print("Adding privileges to group...") + await client.add_privileges_to_group(privilege_group_name, ["Search", "Query", "Insert"]) + print("Privileges added to group") + + print("Creating collection...") + schema = client.create_schema( + auto_id=True, description="Test privilege collection", enable_dynamic_field=True + ) + schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True) + schema.add_field("vector", DataType.FLOAT_VECTOR, dim=128) + schema.add_field("text", DataType.VARCHAR, max_length=512) + + await client.create_collection(collection_name, schema=schema, using_database=test_db) + print(f"Collection {collection_name} created in database {test_db}") + + return True + + +async def test_grant_revoke_privilege(client: AsyncMilvusClient): + print(fmt.format("Testing grant_privilege / revoke_privilege")) + + # Test grant_privilege for different privileges + print("Testing grant_privilege...") + for privilege_item in db_rw_privileges: + await client.grant_privilege( + role_name=test_role_v1, + object_type=privilege_item["object_type"], + privilege=privilege_item["privilege"], + object_name=privilege_item["object_name"], + db_name=test_db, + ) + print( + f"Granted {privilege_item['privilege']} on {privilege_item['object_name']} to {test_role_v1}" + ) + + # Check role privileges + print(f"\nChecking role privileges for {test_role_v1}...") + role_info = await client.describe_role(test_role_v1) + print(f"Role {test_role_v1} privileges: {role_info}") + + # Grant role to user + print(f"\nGranting role {test_role_v1} to user {test_user}...") + await client.grant_role(test_user, test_role_v1) + print("Role granted to user") + + # Check user info + user_info = await client.describe_user(test_user) + print(f"User {test_user} info: {user_info}") + + # Test revoke_privilege + print("\nTesting revoke_privilege...") + for privilege_item in db_ro_privileges: + await client.revoke_privilege( + role_name=test_role_v1, + object_type=privilege_item["object_type"], + privilege=privilege_item["privilege"], + object_name=privilege_item["object_name"], + db_name=test_db, + ) + print( + f"Revoked {privilege_item['privilege']} on {privilege_item['object_name']} from {test_role_v1}" + ) + + # Check role privileges after revoke + print(f"\nChecking role privileges after revoke for {test_role_v1}...") + role_info = await client.describe_role(test_role_v1) + print(f"Role {test_role_v1} privileges after revoke: {role_info}") + + # Revoke role from user + print(f"\nRevoking role {test_role_v1} from user {test_user}...") + await client.revoke_role(test_user, test_role_v1) + print("Role revoked from user") + + +async def test_grant_revoke_privilege_v2(client: AsyncMilvusClient): + print(fmt.format("Testing grant_privilege_v2 / revoke_privilege_v2")) + + # Test grant_privilege_v2 with custom privilege group + print("Testing grant_privilege_v2 with custom privilege group...") + await client.grant_privilege_v2( + role_name=test_role_v2, privilege=privilege_group_name, collection_name="*", db_name=test_db + ) + print( + f"Granted privilege group {privilege_group_name} on all collections in {test_db} to {test_role_v2}" + ) + + # Test grant_privilege_v2 with built-in privilege groups + print("\nTesting grant_privilege_v2 with built-in privilege groups...") + + # Grant collection-level privileges + await client.grant_privilege_v2( + role_name=test_role_v2, + privilege="CollectionReadWrite", + collection_name=collection_name, + db_name=test_db, + ) + print(f"Granted CollectionReadWrite on {collection_name} to {test_role_v2}") + + # Grant database-level privileges + await client.grant_privilege_v2( + role_name=test_role_v2, privilege="DatabaseReadOnly", collection_name="*", db_name=test_db + ) + print(f"Granted DatabaseReadOnly on database {test_db} to {test_role_v2}") + + # Check role privileges + print(f"\nChecking role privileges for {test_role_v2}...") + role_info = await client.describe_role(test_role_v2) + print(f"Role {test_role_v2} privileges: {role_info}") + + # Grant role to user + print(f"\nGranting role {test_role_v2} to user {test_user}...") + await client.grant_role(test_user, test_role_v2) + print("Role granted to user") + + # Check user info + user_info = await client.describe_user(test_user) + print(f"User {test_user} info: {user_info}") + + # Test revoke_privilege_v2 + print("\nTesting revoke_privilege_v2...") + + # Revoke custom privilege group + await client.revoke_privilege_v2( + role_name=test_role_v2, privilege=privilege_group_name, collection_name="*", db_name=test_db + ) + print(f"Revoked privilege group {privilege_group_name} from {test_role_v2}") + + # Revoke built-in privilege group + await client.revoke_privilege_v2( + role_name=test_role_v2, privilege="DatabaseReadOnly", collection_name="*", db_name=test_db + ) + print(f"Revoked DatabaseReadOnly from {test_role_v2}") + + # Revoke collection-level privilege + await client.revoke_privilege_v2( + role_name=test_role_v2, + privilege="CollectionReadWrite", + collection_name=collection_name, + db_name=test_db, + ) + print(f"Revoked CollectionReadWrite from {test_role_v2}") + + # Check role privileges after revoke + print(f"\nChecking role privileges after revoke for {test_role_v2}...") + role_info = await client.describe_role(test_role_v2) + print(f"Role {test_role_v2} privileges after revoke: {role_info}") + + # Revoke role from user + print(f"\nRevoking role {test_role_v2} from user {test_user}...") + await client.revoke_role(test_user, test_role_v2) + print("Role revoked from user") + + +async def main(): + client = AsyncMilvusClient("http://localhost:19530", user=super_user, password=super_password) + + await create_resources(client) + await test_grant_revoke_privilege(client) + await test_grant_revoke_privilege_v2(client) + print(fmt.format("All Privilege Tests Completed Successfully")) + await client.close() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/pymilvus/client/async_grpc_handler.py b/pymilvus/client/async_grpc_handler.py index f47bf3107..bbb4d1b15 100644 --- a/pymilvus/client/async_grpc_handler.py +++ b/pymilvus/client/async_grpc_handler.py @@ -1822,7 +1822,7 @@ async def compact( ) check_status(response.status) - req = Prepare.manual_compaction(response.collectionID, collection_name, is_clustering) + req = Prepare.manual_compaction(collection_name, is_clustering, response.collectionID) response = await self._async_stub.ManualCompaction(req, timeout=timeout, metadata=meta) check_status(response.status)