diff --git a/shared/datastore/src/datastore_interface.py b/shared/datastore/src/datastore_interface.py
index 97aa365..7cd5445 100644
--- a/shared/datastore/src/datastore_interface.py
+++ b/shared/datastore/src/datastore_interface.py
@@ -1,5 +1,7 @@
 """Definition of datastore interface."""
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 
@@ -25,8 +27,9 @@ class DatastoreInterface(ABC):
     @abstractmethod
     def __enter__(
         self,
-    ) -> None:
+    ) -> DatastoreInterface:
         self.connect()
+        return self
 
     @abstractmethod
     def __exit__(