diff --git a/examples/v1/bulk_import_export_relationships.py b/examples/v1/bulk_import_export_relationships.py index f133131..68282ba 100644 --- a/examples/v1/bulk_import_export_relationships.py +++ b/examples/v1/bulk_import_export_relationships.py @@ -1,3 +1,9 @@ +""" +This demonstrates the simplest possible usage of the bulk +import and export functionality. See import_bulk_relationships.py +for a more realistic example of batching. +""" + from authzed.api.v1 import ( BulkExportRelationshipsRequest, BulkImportRelationshipsRequest, @@ -60,7 +66,7 @@ ) ] -import_reps = client.BulkImportRelationships(((req for req in reqs))) +import_reps = client.BulkImportRelationships((req for req in reqs)) assert import_reps.num_loaded == 2 export_resp = client.BulkExportRelationships( diff --git a/examples/v1/import_bulk_relationships.py b/examples/v1/import_bulk_relationships.py new file mode 100644 index 0000000..45cafc6 --- /dev/null +++ b/examples/v1/import_bulk_relationships.py @@ -0,0 +1,79 @@ +""" +This is intended to be a (slightly) more real-world example +that demonstrates the two levels of batching in making BulkImportRelationships +requests. + +This example makes use of itertools.batched to break up an iterator of relationships +into chunks. Documentation for itertools.batched can be found here: +https://docs.python.org/3/library/itertools.html#itertools.batched +""" + +from itertools import batched + +from authzed.api.v1 import ( + Client, + ObjectReference, + Relationship, + SubjectReference, + WriteSchemaRequest, +) +from authzed.api.v1.permission_service_pb2 import ImportBulkRelationshipsRequest +from grpcutil import insecure_bearer_token_credentials + +TOKEN = "sometoken" + +# Stand up a client +client = Client( + "localhost:50051", + insecure_bearer_token_credentials(TOKEN), +) + +# Write a simple schema +schema = """ + definition user {} + definition resource { + relation viewer: user + permission view = viewer + } +""" + +client.WriteSchema(WriteSchemaRequest(schema=schema)) + + +# A generator that we can use to create an arbitrarily-long list of relationships +# In your own application, this would be whatever's generating the list of imported +# relationships. +def relationship_generator(num_relationships): + idx = 0 + while idx < num_relationships: + idx += 1 + yield Relationship( + resource=ObjectReference(object_type="resource", object_id=str(idx)), + relation="viewer", + subject=SubjectReference( + object=ObjectReference(object_type="user", object_id="our_user") + ), + ) + + +TOTAL_RELATIONSHIPS_TO_WRITE = 1_000 + +RELATIONSHIPS_PER_TRANSACTION = 100 +RELATIONSHIPS_PER_REQUEST_CHUNK = 10 + +# NOTE: batched takes a larger iterator and makes an iterator of smaller chunks out of it. +# We iterate over chunks of size RELATIONSHIPS_PER_TRANSACTION, and then we break each request into +# chunks of size RELATIONSHIPS_PER_REQUEST_CHUNK. +transaction_chunks = batched( + relationship_generator(TOTAL_RELATIONSHIPS_TO_WRITE), RELATIONSHIPS_PER_TRANSACTION +) +for relationships_for_request in transaction_chunks: + request_chunks = batched(relationships_for_request, RELATIONSHIPS_PER_REQUEST_CHUNK) + response = client.ImportBulkRelationships( + ( + ImportBulkRelationshipsRequest(relationships=relationships_chunk) + for relationships_chunk in request_chunks + ) + ) + print("request successful") + print(response.num_loaded)