Merge pull request #219 from authzed/add-more-realistic-bulk-import-e…

…xample

Add more realistic bulk import example

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(

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)