Allow custom entity names (#105)

* Allow custom entity names

Author: andystaples

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 CHANGED
 
+- Allow entities with custom names
 - Add/update type-hinting for various worker methods
 
 ## v1.2.0

durabletask/task.py

@@ -615,6 +615,14 @@ def retry_timeout(self) -> Optional[timedelta]:
         return self._retry_timeout
 
 
+def get_entity_name(fn: Entity) -> str:
+    if hasattr(fn, "__durable_entity_name__"):
+        return getattr(fn, "__durable_entity_name__")
+    if isinstance(fn, type) and issubclass(fn, DurableEntity):
+        return fn.__name__
+    return get_name(fn)
+
+
 def get_name(fn: Callable) -> str:
     """Returns the name of the provided function"""
     name = fn.__name__

durabletask/worker.py

@@ -188,16 +188,14 @@ def add_named_activity(self, name: str, fn: task.Activity[TInput, TOutput]) -> N
     def get_activity(self, name: str) -> Optional[task.Activity[Any, Any]]:
         return self.activities.get(name)
 
-    def add_entity(self, fn: task.Entity) -> str:
+    def add_entity(self, fn: task.Entity, name: Optional[str] = None) -> str:
         if fn is None:
             raise ValueError("An entity function argument is required.")
 
-        if isinstance(fn, type) and issubclass(fn, DurableEntity):
-            name = fn.__name__
-            self.add_named_entity(name, fn)
-        else:
-            name = task.get_name(fn)
-            self.add_named_entity(name, fn)
+        if name is None:
+            name = task.get_entity_name(fn)
+
+        self.add_named_entity(name, fn)
         return name
 
     def add_named_entity(self, name: str, fn: task.Entity) -> None:
@@ -378,13 +376,13 @@ def add_activity(self, fn: task.Activity) -> str:
             )
         return self._registry.add_activity(fn)
 
-    def add_entity(self, fn: task.Entity) -> str:
+    def add_entity(self, fn: task.Entity, name: Optional[str] = None) -> str:
         """Registers an entity function with the worker."""
         if self._is_running:
             raise RuntimeError(
                 "Entities cannot be added while the worker is running."
             )
-        return self._registry.add_entity(fn)
+        return self._registry.add_entity(fn, name)
 
     def use_versioning(self, version: VersioningOptions) -> None:
         """Initializes versioning options for sub-orchestrators and activities."""

tests/durabletask-azuremanaged/test_dts_class_based_entities_e2e.py

@@ -17,7 +17,7 @@
 endpoint = os.getenv("ENDPOINT", "http://localhost:8080")
 
 
-def test_client_signal_class_entity():
+def test_client_signal_class_entity_and_custom_name():
     invoked = False
 
     class EmptyEntity(entities.DurableEntity):
@@ -28,12 +28,12 @@ def do_nothing(self, _):
     # Start a worker, which will connect to the sidecar in a background thread
     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
                                     taskhub=taskhub_name, token_credential=None) as w:
-        w.add_entity(EmptyEntity)
+        w.add_entity(EmptyEntity, name="EntityNameCustom")
         w.start()
 
         c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
                                        taskhub=taskhub_name, token_credential=None)
-        entity_id = entities.EntityInstanceId("EmptyEntity", "testEntity")
+        entity_id = entities.EntityInstanceId("EntityNameCustom", "testEntity")
         c.signal_entity(entity_id, "do_nothing")
         time.sleep(2)  # wait for the signal to be processed
 
@@ -70,7 +70,7 @@ def do_nothing(self, _):
     assert invoked
 
 
-def test_orchestration_signal_class_entity():
+def test_orchestration_signal_class_entity_and_custom_name():
     invoked = False
 
     class EmptyEntity(entities.DurableEntity):
@@ -79,14 +79,14 @@ def do_nothing(self, _):
             invoked = True
 
     def empty_orchestrator(ctx: task.OrchestrationContext, _):
-        entity_id = entities.EntityInstanceId("EmptyEntity", "testEntity")
+        entity_id = entities.EntityInstanceId("EntityNameCustom", "testEntity")
         ctx.signal_entity(entity_id, "do_nothing")
 
     # Start a worker, which will connect to the sidecar in a background thread
     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
                                     taskhub=taskhub_name, token_credential=None) as w:
         w.add_orchestrator(empty_orchestrator)
-        w.add_entity(EmptyEntity)
+        w.add_entity(EmptyEntity, name="EntityNameCustom")
         w.start()
 
         c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,

tests/durabletask-azuremanaged/test_dts_function_based_entities_e2e.py

@@ -17,7 +17,7 @@
 endpoint = os.getenv("ENDPOINT", "http://localhost:8080")
 
 
-def test_client_signal_entity():
+def test_client_signal_entity_and_custom_name():
     invoked = False
 
     def empty_entity(ctx: entities.EntityContext, _):
@@ -28,12 +28,12 @@ def empty_entity(ctx: entities.EntityContext, _):
     # Start a worker, which will connect to the sidecar in a background thread
     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
                                     taskhub=taskhub_name, token_credential=None) as w:
-        w.add_entity(empty_entity)
+        w.add_entity(empty_entity, name="EntityNameCustom")
         w.start()
 
         c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
                                        taskhub=taskhub_name, token_credential=None)
-        entity_id = entities.EntityInstanceId("empty_entity", "testEntity")
+        entity_id = entities.EntityInstanceId("EntityNameCustom", "testEntity")
         c.signal_entity(entity_id, "do_nothing")
         time.sleep(2)  # wait for the signal to be processed
 
@@ -70,7 +70,7 @@ def empty_entity(ctx: entities.EntityContext, _):
     assert invoked
 
 
-def test_orchestration_signal_entity():
+def test_orchestration_signal_entity_and_custom_name():
     invoked = False
 
     def empty_entity(ctx: entities.EntityContext, _):
@@ -79,14 +79,14 @@ def empty_entity(ctx: entities.EntityContext, _):
             invoked = True
 
     def empty_orchestrator(ctx: task.OrchestrationContext, _):
-        entity_id = entities.EntityInstanceId("empty_entity", f"{ctx.instance_id}_testEntity")
+        entity_id = entities.EntityInstanceId("EntityNameCustom", f"{ctx.instance_id}_testEntity")
         ctx.signal_entity(entity_id, "do_nothing")
 
     # Start a worker, which will connect to the sidecar in a background thread
     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
                                     taskhub=taskhub_name, token_credential=None) as w:
         w.add_orchestrator(empty_orchestrator)
-        w.add_entity(empty_entity)
+        w.add_entity(empty_entity, name="EntityNameCustom")
         w.start()
 
         c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,