Set Authenticated User and Role Fields from Client (#490)

Addresses #489

Author: kraftp

dbos/_client.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import time
 import uuid
 from typing import (
@@ -63,6 +64,8 @@ class EnqueueOptions(_EnqueueOptionsRequired, total=False):
     priority: int
     max_recovery_attempts: int
     queue_partition_key: str
+    authenticated_user: str
+    authenticated_roles: list[str]
 
 
 def validate_enqueue_options(options: EnqueueOptions) -> None:
@@ -189,6 +192,13 @@ def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
             "queue_partition_key": options.get("queue_partition_key"),
         }
 
+        authenticated_user = options.get("authenticated_user")
+        authenticated_roles = (
+            json.dumps(options.get("authenticated_roles"))
+            if options.get("authenticated_roles")
+            else None
+        )
+
         inputs: WorkflowInputs = {
             "args": args,
             "kwargs": kwargs,
@@ -202,9 +212,9 @@ def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
             "queue_name": queue_name,
             "app_version": enqueue_options_internal["app_version"],
             "config_name": None,
-            "authenticated_user": None,
+            "authenticated_user": authenticated_user,
             "assumed_role": None,
-            "authenticated_roles": None,
+            "authenticated_roles": authenticated_roles,
             "output": None,
             "error": None,
             "created_at": None,

tests/test_client.py

@@ -593,3 +593,35 @@ def test_enqueue_with_priority(dbos: DBOS, client: DBOSClient) -> None:
 def test_client_bad_url() -> None:
     with pytest.raises(DBAPIError) as exc_info:
         DBOSClient("postgresql://postgres:fakepassword@localhost:5433/fake_database")
+
+
+def test_client_auth(dbos: DBOS, client: DBOSClient) -> None:
+    run_client_collateral()
+
+    johnDoe: Person = {"first": "John", "last": "Doe", "age": 30}
+    wfid = str(uuid.uuid4())
+
+    user = "testuser"
+    roles = ["role1", "role2"]
+
+    options: EnqueueOptions = {
+        "queue_name": "test_queue",
+        "workflow_name": "enqueue_test",
+        "workflow_id": wfid,
+        "authenticated_user": user,
+        "authenticated_roles": roles,
+    }
+
+    handle: WorkflowHandle[str] = client.enqueue(options, 42, "test", johnDoe)
+    result = handle.get_result()
+    assert result == '42-test-{"first": "John", "last": "Doe", "age": 30}'
+
+    list_results = client.list_workflows()
+    assert len(list_results) == 1
+    assert list_results[0].workflow_id == wfid
+    assert list_results[0].status == "SUCCESS"
+    assert list_results[0].output == result
+    assert list_results[0].input is not None
+    assert list_results[0].authenticated_user == user
+    assert list_results[0].authenticated_roles == roles
+    assert list_results[0].assumed_role is None