diff --git a/README.md b/README.md index 6e80c78..b73945e 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,8 @@ Optional, only used in [com.atproto.repo](https://arroba.readthedocs.io/en/stabl _Breaking changes:_ +* `datastore_storage`: + * `DatastoreStorage`: add new `ndb_context_kwargs` constructor kwarg. * `repo`: * `apply_commit`, `apply_writes`: raise an exception if the repo is inactive. * `storage`: diff --git a/arroba/datastore_storage.py b/arroba/datastore_storage.py index fb0228d..dd405f4 100644 --- a/arroba/datastore_storage.py +++ b/arroba/datastore_storage.py @@ -402,24 +402,26 @@ class DatastoreStorage(Storage): See :class:`Storage` for method details. """ ndb_client = None + ndb_context_kwargs = None - def __init__(self, *, ndb_client=None): + def __init__(self, *, ndb_client=None, ndb_context_kwargs=None): """Constructor. Args: - ndb_client (google.cloud.ndb.Client): used in :meth:`read_blocks_by_seq`; - it's used in the `subscribeRepos` event subscription, so lexrpc calls - it on a different thread, so it needs its own ndb client context. + ndb_client (google.cloud.ndb.Client): used when there isn't already + an ndb context active + ndb_context_kwargs (dict): optional, used when creating a new ndb context """ super().__init__() self.ndb_client = ndb_client + self.ndb_context_kwargs = ndb_context_kwargs or {} def ndb_context(fn): @wraps(fn) def decorated(self, *args, **kwargs): ctx = context.get_context(raise_context_error=False) - with ctx.use() if ctx else self.ndb_client.context(): + with ctx.use() if ctx else self.ndb_client.context(**self.ndb_context_kwargs): ret = fn(self, *args, **kwargs) return ret @@ -541,7 +543,7 @@ def read_blocks_by_seq(self, start=0, repo=None): while True: ctx = context.get_context(raise_context_error=False) - with ctx.use() if ctx else self.ndb_client.context(): + with ctx.use() if ctx else self.ndb_client.context(**self.ndb_context_kwargs): # lexrpc event subscription handlers like subscribeRepos call this # on a different thread, so if we're there, we need to create a new # ndb context