@@ -8369,88 +8369,3 @@ async def test_previous_run_failure(client: Client):
83698369 )
83708370 result = await handle .result ()
83718371 assert result == "Done"
8372-
8373- class EncryptionCodec (PayloadCodec ):
8374- def __init__ (self , key_id : str = "test-key-id" , key : bytes = b"test-key-test-key-test-key-test!" ) -> None :
8375- super ().__init__ ()
8376- self .key_id = key_id
8377-
8378- async def encode (self , payloads : Iterable [Payload ]) -> List [Payload ]:
8379- # We blindly encode all payloads with the key and set the metadata
8380- # saying which key we used
8381- return [
8382- Payload (
8383- metadata = {
8384- "encoding" : b"binary/encrypted" ,
8385- "encryption-key-id" : self .key_id .encode (),
8386- },
8387- data = self .encrypt (p .SerializeToString ()),
8388- )
8389- for p in payloads
8390- ]
8391-
8392- async def decode (self , payloads : Iterable [Payload ]) -> List [Payload ]:
8393- ret : List [Payload ] = []
8394- for p in payloads :
8395- # Ignore ones w/out our expected encoding
8396- if p .metadata .get ("encoding" , b"" ).decode () != "binary/encrypted" :
8397- ret .append (p )
8398- continue
8399- # Confirm our key ID is the same
8400- key_id = p .metadata .get ("encryption-key-id" , b"" ).decode ()
8401- if key_id != self .key_id :
8402- raise ValueError (f"Unrecognized key ID { key_id } . Current key ID is { self .key_id } ." )
8403- # Decrypt and append
8404- ret .append (Payload .FromString (self .decrypt (p .data )))
8405- return ret
8406-
8407- def encrypt (self , data : bytes ) -> bytes :
8408- nonce = os .urandom (12 )
8409- return data
8410-
8411- def decrypt (self , data : bytes ) -> bytes :
8412- return data
8413-
8414-
8415- @workflow .defn (name = "Workflow" )
8416- class GreetingWorkflow :
8417- @workflow .run
8418- async def run (self , name : str ) -> str :
8419- print (
8420- await workflow .execute_child_workflow (
8421- workflow = ChildWorkflow .run ,
8422- arg = name ,
8423- id = f"child-{ name } " ,
8424- search_attributes = workflow .info ().typed_search_attributes ,
8425- )
8426- )
8427- return f"Hello, { name } "
8428-
8429-
8430- @workflow .defn (name = "ChildWorkflow" )
8431- class ChildWorkflow :
8432- @workflow .run
8433- async def run (self , name : str ) -> str :
8434- return f"Hello from child, { name } "
8435-
8436- async def test_search_attribute_codec (client : Client ):
8437-
8438- config = client .config ()
8439- config ["data_converter" ] = dataclasses .replace (temporalio .converter .default (), payload_codec = EncryptionCodec ())
8440- client = Client (** config )
8441- # Run a worker for the workflow
8442- async with Worker (
8443- client ,
8444- task_queue = "encryption-task-queue" ,
8445- workflows = [GreetingWorkflow , ChildWorkflow ],
8446- ):
8447- # Run workflow
8448- result = await client .execute_workflow (
8449- GreetingWorkflow .run ,
8450- "Temporal" ,
8451- id = f"encryption-workflow-id" ,
8452- task_queue = "encryption-task-queue" ,
8453- search_attributes = TypedSearchAttributes (
8454- [SearchAttributePair (SearchAttributeKey .for_keyword ("show_name" ), "test_show" )]
8455- ),
8456- )
0 commit comments