Unable to manually enter the ndb context

[jacek-jablonski]

Environment details

1. OS type and version: Mac OS Sonoma 14.1.1
2. Python version: 3.11.3
3. google-cloud-ndb version: 2.2.2

Steps to reproduce

1. Run code from sample.
2. Exception occurs: google.cloud.ndb.exceptions.ContextError: No current context. NDB calls must be made in context established by google.cloud.ndb.Client.context.

Code example

from google.cloud import ndb


class SomeModel(ndb.Model):
    someproperty = ndb.StringProperty()


class Repo:
    def __init__(self):
        self._ndb = ndb.Client()

    def __enter__(self) -> "Repo":
        self._ndb.context().__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("exit")

    def get_some_data(self):
        query = SomeModel.query().filter(SomeModel.someproperty == "tests")
        for entry in query:
            print(entry)


with Repo() as repo:
    repo.get_some_data()

Stack trace

Traceback (most recent call last):
  File "/Users/jacek/Projects/rtb_apps_resource_usage/test.py", line 26, in <module>
    repo.get_some_data()
  File "/Users/jacek/Projects/rtb_apps_resource_usage/test.py", line 20, in get_some_data
    query = SomeModel.query().filter(SomeModel.someproperty == "tests")
            ^^^^^^^^^^^^^^^^^
  File "/Users/jacek/Projects/.venv/lib/python3.11/site-packages/google/cloud/ndb/utils.py", line 118, in wrapper
    return wrapped(*args, **new_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacek/Projects/.venv/lib/python3.11/site-packages/google/cloud/ndb/model.py", line 5534, in _query
    query = query_module.Query(
            ^^^^^^^^^^^^^^^^^^^
  File "/Users/jacek/Projects/.venv/lib/python3.11/site-packages/google/cloud/ndb/query.py", line 1386, in __init__
    database = context_module.get_context().client.database or None
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacek/Projects/.venv/lib/python3.11/site-packages/google/cloud/ndb/context.py", line 119, in get_context
    raise exceptions.ContextError()
google.cloud.ndb.exceptions.ContextError: No current context. NDB calls must be made in context established by google.cloud.ndb.Client.context.

How can I manually enter the context, so I can design class like that?
I'm quite surprised that there is difference in lib behaviour between calling with self._ndb.context(): and manually entering the contxt: self._ndb.context().__enter__()

[jacek-jablonski]

I've noticed my mistake. Solution:

from google.cloud import ndb


class SomeModel(ndb.Model):
    someproperty = ndb.StringProperty()


class Repo:
    def __init__(self):
        self._ndb = ndb.Client()
        self._ctx = None

    def __enter__(self) -> "Repo":
        self._ctx = self._ndb.context()
        self._ctx.__enter__()

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._ctx.__exit__(exc_type, exc_value, traceback)

    def get_some_data(self):
        query = SomeModel.query().filter(SomeModel.someproperty == "tests")
        for entry in query:
            print(entry)


with Repo() as repo:
    repo.get_some_data()