- Fixes for memory issues in the Python 2.7 runtime.

pcostell · pcostell · commit 0da63f9ae78c · 2016-05-16T15:18:43.000-07:00
- Test cleanup.
diff --git a/ndb/context_test.py b/ndb/context_test.py
@@ -1536,19 +1536,16 @@ def MakeContext(self, *args, **kwargs):
 
 
 @real_unittest.skipUnless(datastore_pbs._CLOUD_DATASTORE_ENABLED,
-                          "V1 must be supported to run V1 tests.")
-class ContextV1WithRemoteAPITests(ContextTestMixin,
-                                 ContextMemcacheTestMixin,
-                                 ContextTaskQueueTestMixin,
-                                 test_utils.NDBCloudDatastoreV1Test):
+                          'V1 must be supported to run V1 tests.')
+class ContextV1Tests(ContextTestMixin,
+                     test_utils.NDBCloudDatastoreV1Test):
   """Context tests that use a Cloud Datastore V1 connection.
 
-  These tests run with memcache and taskqueue stubs available."""
+  These tests run with memcache and taskqueue stubs available.
+  """
 
   def setUp(self):
-    super(ContextV1WithRemoteAPITests, self).setUp()
-    self.testbed.init_memcache_stub()
-    self.testbed.init_taskqueue_stub()
+    super(ContextV1Tests, self).setUp()
     self.HRTest()
     MyAutoBatcher.reset_log()
     self.ctx = self.MakeContext(default_model=model.Expando,
@@ -1558,75 +1555,74 @@ def setUp(self):
   def make_bad_transaction(*arg, **kwargs):
     return ''
 
+  def testContext_AutoBatcher_Errors(self):
+    # Performs tests via direct memcache call, which is disabled by default
+    # in V1.
+    pass
+
+  def testContext_AllocateIds(self):
+    # V1 does not support Allocate id range.
+    pass
+
   def testContext_TransactionAddTask(self):
     # Transactional AddTask still will be unavailable.
-    self.ExpectWarnings()
-    key = model.Key('Foo', 1)
-
-    @tasklets.tasklet
     def foo():
-      ent = model.Expando(key=key, bar=1)
-
-      @tasklets.tasklet
-      def callback():
-        ctx = tasklets.get_context()
-        yield ctx.put(ent)
-        taskqueue.add(url='/', transactional=True)
-      yield self.ctx.transaction(callback)
-    self.assertRaises(ValueError, foo().check_success)
+      taskqueue.add(url='/', transactional=True)
+    self.assertRaises(ValueError, model.transaction, foo)
 
   def MakeContext(self, *args, **kwargs):
-    ctx = super(ContextV1WithRemoteAPITests, self).MakeContext(*args, **kwargs)
-    # Re-enable default cache policy.
+    ctx = super(ContextV1Tests, self).MakeContext(*args, **kwargs)
+    # Re-enable in-context cache. Memcache must remain off because
+    # the stub is not enabled.
     ctx.set_cache_policy(None)
-    ctx.set_memcache_policy(None)
     return ctx
 
+
 @real_unittest.skipUnless(datastore_pbs._CLOUD_DATASTORE_ENABLED,
-                          "V1 must be supported to run V1 tests.")
-class ContextV1Tests(ContextTestMixin,
-                     test_utils.NDBCloudDatastoreV1Test):
+                          'V1 must be supported to run V1 tests.')
+class ContextV1WithRemoteAPITests(ContextV1Tests,
+                                  ContextMemcacheTestMixin,
+                                  ContextTaskQueueTestMixin):
   """Context tests that use a Cloud Datastore V1 connection.
 
-  These tests run with memcache and taskqueue stubs available."""
+  These tests run with memcache and taskqueue stubs available.
+  """
 
   def setUp(self):
-    super(ContextV1Tests, self).setUp()
+    # testbed needs to get set up first.
+    super(ContextV1WithRemoteAPITests, self).setUp()
     self.testbed.init_memcache_stub()
     self.testbed.init_taskqueue_stub()
-    self.HRTest()
-    MyAutoBatcher.reset_log()
-    self.ctx = self.MakeContext(default_model=model.Expando,
-                                auto_batcher_class=MyAutoBatcher)
-    tasklets.set_context(self.ctx)
-
-  def make_bad_transaction(*arg, **kwargs):
-    return ''
 
-  def testContext_TransactionAddTask(self):
-    # Transactional AddTask still will be unavailable.
+  def testContext_AutoBatcher_Errors(self):
+    # Test that errors are properly distributed over all Futures.
     self.ExpectWarnings()
-    key = model.Key('Foo', 1)
-
-    @tasklets.tasklet
-    def foo():
-      ent = model.Expando(key=key, bar=1)
 
-      @tasklets.tasklet
-      def callback():
-        ctx = tasklets.get_context()
-        yield ctx.put(ent)
-        taskqueue.add(url='/', transactional=True)
-      yield self.ctx.transaction(callback)
-    self.assertRaises(ValueError, foo().check_success)
+    class Blobby(model.Model):
+      blob = model.BlobProperty()
+    ent1 = Blobby()
+    ent2 = Blobby(blob='x' * 2000000)
+    fut1 = self.ctx.put(ent1)
+    fut2 = self.ctx.put(ent2)  # Error
+    err1 = fut1.get_exception()
+    err2 = fut2.get_exception()
+    self.assertTrue(isinstance(err1, datastore_errors.BadRequestError))
+    self.assertTrue(err1 is err2)
+    # Try memcache as well (different tasklet, different error).
+    fut1 = self.ctx.memcache_set('key1', 'x')
+    fut2 = self.ctx.memcache_set('key2', 'x' * 1000001)
+    err1 = fut1.get_exception()
+    err2 = fut1.get_exception()
+    self.assertTrue(isinstance(err1, ValueError))
+    self.assertTrue(err1 is err2)
 
   def MakeContext(self, *args, **kwargs):
-    ctx = super(ContextV1Tests, self).MakeContext(*args, **kwargs)
-    # Re-enable in-context cache. Memcache must remain off because
-    # the stub is not enabled.
-    ctx.set_cache_policy(None)
+    ctx = super(ContextV1WithRemoteAPITests, self).MakeContext(*args, **kwargs)
+    # Re-enable memcache.
+    ctx.set_memcache_policy(None)
     return ctx
 
+
 class ContextFutureCachingTests(test_utils.NDBTest):
   # See issue 62.  http://goo.gl/5zLkK
 
diff --git a/ndb/google_imports.py b/ndb/google_imports.py
@@ -88,9 +88,17 @@ def set_appengine_imports():
   from google.appengine.ext import db
   from google.appengine.ext import gql
   try:
+    # For the python-compat runtime.
     from google.appengine.ext.vmruntime import callback
   except ImportError:
-    callback = None
+    # For the python 2.7 runtime.
+    try:
+      from google.appengine.runtime import apiproxy as callback
+      # Python 2.5 and dev_appserver is not supported.
+      if not hasattr(callback, 'SetRequestEndCallback'):
+        callback = None
+    except ImportError:
+      callback = None
   from google.appengine.runtime import apiproxy_errors
   from google.net.proto import ProtocolBuffer
 else:
diff --git a/ndb/tasklets.py b/ndb/tasklets.py
@@ -106,6 +106,12 @@ def foo():
 
 _CALLBACK_KEY = '__CALLBACK__'
 
+# Python 2.5 compatability.
+if hasattr(weakref, 'WeakSet'):
+  _set = weakref.WeakSet
+else:
+  _set = set
+
 
 def _is_generator(obj):
   """Helper to test for a generator object.
@@ -122,20 +128,25 @@ class _State(utils.threading_local):
   def __init__(self):
     super(_State, self).__init__()
     self.current_context = None
-    self.all_generators = weakref.WeakSet()
+    self.all_generators = _set()
     self.all_pending = set()
 
   def set_context(self, ctx):
+    self.current_context = ctx
+
+  def add_generator(self, gen):
     if _request_callback and _CALLBACK_KEY not in os.environ:
       _request_callback.SetRequestEndCallback(self.reset)
       os.environ[_CALLBACK_KEY] = '1'
-    self.current_context = ctx
 
-  def add_generator(self, gen):
     _logging_debug('all_generators: add %s', gen)
     self.all_generators.add(gen)
 
   def add_pending(self, fut):
+    if _request_callback and _CALLBACK_KEY not in os.environ:
+      _request_callback.SetRequestEndCallback(self.reset)
+      os.environ[_CALLBACK_KEY] = '1'
+
     _logging_debug('all_pending: add %s', fut)
     self.all_pending.add(fut)
 
@@ -172,7 +183,7 @@ def dump_all_pending(self, verbose=False):
       pending.append(line)
     return '\n'.join(pending)
 
-  def reset(self):
+  def reset(self, unused_req_id):
     self.current_context = None
     ev = eventloop.get_event_loop()
     ev.clear()
diff --git a/ndb/tasklets_test.py b/ndb/tasklets_test.py
@@ -855,31 +855,50 @@ class CloudDatastoreContextCreationTests(test_utils.NDBCloudDatastoreV1Test):
   def setUp(self):
     super(CloudDatastoreContextCreationTests, self).setUp()
     self.HRTest()
-    self.ctx = tasklets._make_cloud_datastore_context(self.APP_ID)
-    tasklets.set_context(self.ctx)
-
-  def testCloudDatastoreContext(self):
+    self.old_env = dict(os.environ)
     del os.environ['APPLICATION_ID']  # Clear APPLICATION_ID for test.
+
+  def tearDown(self):
+    super(CloudDatastoreContextCreationTests, self).setUp()
+    os.environ.clear()
+    os.environ.update(self.old_env)
+
+  def testCloudDatastoreContext_ProjectIdNotSufficient(self):
     os.environ['DATASTORE_PROJECT_ID'] = 'project'
     # Setting just a project id isn't enough info.
     self.assertRaisesRegexp(ValueError,
                             'Could not determine app id\..*',
                             tasklets.make_default_context)
-    # Unless the override is specified.
-    os.environ['DATASTORE_USE_PROJECT_ID_AS_APP_ID'] = 'True'
+
+  def testCloudDatastoreContext_UseProjectIdAsAppId(self):
+    os.environ['DATASTORE_PROJECT_ID'] = 'project'
+    os.environ['DATASTORE_USE_PROJECT_ID_AS_APP_ID'] = 'true'
     ctx = tasklets.make_default_context()
     self.assertEqual(datastore_rpc._CLOUD_DATASTORE_V1, ctx._conn._api_version)
+
+  def testCloudDatastoreContext_AppIdAndUseProjectIdAsAppId(self):
+    os.environ['DATASTORE_PROJECT_ID'] = 'project'
+    os.environ['DATASTORE_USE_PROJECT_ID_AS_APP_ID'] = 'true'
     os.environ['DATASTORE_APP_ID'] = 's~project'
     # App id and override should not both be specified
     self.assertRaisesRegexp(ValueError,
                             'App id was provided .* but .* was set to true\.',
                             tasklets.make_default_context)
-    del os.environ['DATASTORE_USE_PROJECT_ID_AS_APP_ID']
+
+  def testCloudDatastoreContext_AppId(self):
+    os.environ['DATASTORE_PROJECT_ID'] = 'project'
+    os.environ['DATASTORE_APP_ID'] = 's~project'
     ctx = tasklets.make_default_context()
     self.assertEqual(datastore_rpc._CLOUD_DATASTORE_V1, ctx._conn._api_version)
+
+  def testCloudDatastoreContext_DifferentAppId(self):
+    os.environ['DATASTORE_PROJECT_ID'] = 'project'
     os.environ['DATASTORE_APP_ID'] = 's~anotherproject'
     # The provided app id must resolve to the provided project id.
     self.assertRaises(ValueError, tasklets.make_default_context)
+
+  def testCloudDatastoreContext_ManyAppIds(self):
+    os.environ['DATASTORE_PROJECT_ID'] = 'project'
     os.environ['DATASTORE_APP_ID'] = 's~project'
     os.environ['DATASTORE_ADDITIONAL_APP_IDS'] = 's~anotherapp,s~andanother'
     ctx = tasklets.make_default_context()
@@ -895,13 +914,19 @@ def testCloudDatastoreContextWithExistingApplicationId(self):
     self.assertRaises(ValueError, tasklets.make_default_context)
 
   def testContext_MemcacheUnavailable(self):
+    os.environ['DATASTORE_PROJECT_ID'] = self.APP_ID
+    os.environ['DATASTORE_USE_PROJECT_ID_AS_APP_ID'] = 'true'
+    tasklets.set_context(tasklets.make_default_context())
     key = model.Key('Foo', 1)
     ent = model.Expando(key=key, bar=1)
     ctx = tasklets.get_context()
     ctx.set_memcache_policy(True)
     self.assertRaises(NotImplementedError, ctx.put(ent).get_result)
 
   def testContext_TaskQueueUnavailable(self):
+    os.environ['DATASTORE_PROJECT_ID'] = self.APP_ID
+    os.environ['DATASTORE_USE_PROJECT_ID_AS_APP_ID'] = 'true'
+    tasklets.set_context(tasklets.make_default_context())
     self.assertRaises(NotImplementedError, taskqueue.add, url='/')
 
 
diff --git a/ndb/test_utils.py b/ndb/test_utils.py
@@ -22,6 +22,7 @@
 
 import logging
 import tempfile
+import os
 
 from .google_imports import apiproxy_stub_map
 from .google_imports import datastore
@@ -178,6 +179,10 @@ def tearDownClass(cls):
 
   def setUp(self):
     super(NDBCloudDatastoreV1Test, self).setUp()
+    # The host doesn't get used since we override the stub, however the
+    # existence of this environment variable can break if we try to get
+    # credentials.
+    os.environ['DATASTORE_EMULATOR_HOST'] = 'localhost:1234'
     self.datastore.Clear()
     stub = cloud_datastore_v1_remote_stub.CloudDatastoreV1RemoteStub(
         self.datastore.GetDatastore())
