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__(