Example #1
0
  def testSetOptions(self):
    datastore._conn = None
    self.mox.StubOutWithMock(os, 'getenv')
    self.mox.StubOutWithMock(helper, 'get_credentials_from_env')
    os.getenv('DATASTORE_HOST').AndReturn('http://localhost:8080')
    helper.get_credentials_from_env().AndReturn(FakeCredentialsFromEnv())
    self.mox.ReplayAll()

    datastore.set_options(dataset='bar')
    conn = datastore.get_default_connection()
    self.assertEqual('http://localhost:8080/datastore/v1beta1/datasets/bar/',
                     conn._url)
    self.assertEqual(FakeCredentialsFromEnv, type(conn._credentials))
    self.mox.VerifyAll()
def get_default_connection():
  """Return the default datastore connection.

  Defaults dataset to os.getenv('DATASTORE_DATASET'), host to
  os.getenv('DATASTORE_HOST'), and credentials to
  helper.get_credentials_from_env().

  Use set_options to override defaults.
  """
  tid = id(threading.current_thread())
  conn = _conn_holder.get(tid)
  if not conn:
    with(_rlock):
      # No other thread would insert a value in our slot, so no need
      # to recheck existence inside the lock.
      if 'dataset' not in _options:
        _options['dataset'] = os.getenv('DATASTORE_DATASET')
      if 'host' not in _options:
        _options['host'] = os.getenv('DATASTORE_HOST')
      if 'credentials' not in _options:
        _options['credentials'] = helper.get_credentials_from_env()
      # We still need the lock when caching the thread local connection so we
      # don't race with _conn_holder.clear() in set_options().
      _conn_holder[tid] = conn = Datastore(**_options)
  return conn
Example #3
0
def get_default_connection():
    """Return the default datastore connection.

  dataset defaults to helper.get_dataset_from_env(), host to
  os.getenv('DATASTORE_HOST'), and credentials to
  helper.get_credentials_from_env().

  Use set_options to override defaults.
  """
    tid = id(threading.current_thread())
    conn = _conn_holder.get(tid)
    if not conn:
        with (_rlock):
            # No other thread would insert a value in our slot, so no need
            # to recheck existence inside the lock.
            if 'dataset' not in _options:
                _options['dataset'] = os.getenv('DATASTORE_DATASET')
            if 'host' not in _options:
                _options['host'] = os.getenv('DATASTORE_HOST')
            if 'credentials' not in _options:
                _options['credentials'] = helper.get_credentials_from_env()
            # We still need the lock when caching the thread local connection so we
            # don't race with _conn_holder.clear() in set_options().
            _conn_holder[tid] = conn = Datastore(**_options)
    return conn
def _make_cloud_datastore_context(app_id, external_app_ids=(), host=None):
    """Creates a new context to connect to a remote Cloud Datastore instance.

  This should only be used outside of Google App Engine.

  Args:
    app_id: The application id to connect to. This differs from the project
      id as it may have an additional prefix, e.g. "s~" or "e~".
    external_app_ids: A list of apps that may be referenced by data in your
      application. For example, if you are connected to s~my-app and store keys
      for s~my-other-app, you should include s~my-other-app in the external_apps
      list.
    host: The hostname to provide to the datastore connection. If None, the
      default is used.
  Returns:
    An ndb.Context that can connect to a Remote Cloud Datastore. You can use
    this context by passing it to ndb.set_context.
  """
    from . import model  # Late import to deal with circular imports.

    # Late import since it might not exist.
    if not datastore_pbs._CLOUD_DATASTORE_ENABLED:
        raise datastore_errors.BadArgumentError(datastore_pbs.MISSING_CLOUD_DATASTORE_MESSAGE)
    import googledatastore

    try:
        from google.appengine.datastore import cloud_datastore_v1_remote_stub
    except ImportError:
        from google3.apphosting.datastore import cloud_datastore_v1_remote_stub

    current_app_id = os.environ.get("APPLICATION_ID", None)
    if current_app_id and current_app_id != app_id:
        # TODO(pcostello): We should support this so users can connect to different
        # applications.
        raise ValueError(
            "Cannot create a Cloud Datastore context that connects "
            "to an application (%s) that differs from the application "
            "already connected to (%s)." % (app_id, current_app_id)
        )
    os.environ["APPLICATION_ID"] = app_id

    id_resolver = datastore_pbs.IdResolver((app_id,) + tuple(external_app_ids))
    project_id = id_resolver.resolve_project_id(app_id)

    from googledatastore import helper

    credentials = helper.get_credentials_from_env()
    datastore = googledatastore.Datastore(project_id, credentials=credentials, host=host)

    conn = model.make_connection(_api_version=datastore_rpc._CLOUD_DATASTORE_V1, _id_resolver=id_resolver)

    # If necessary, install the stubs
    try:
        stub = cloud_datastore_v1_remote_stub.CloudDatastoreV1RemoteStub(datastore)
        apiproxy_stub_map.apiproxy.RegisterStub(datastore_rpc._CLOUD_DATASTORE_V1, stub)
    except:
        pass  # The stub is already installed.
    # TODO(pcostello): Ensure the current stub is connected to the right project.
    return make_context(conn=conn)
 def testGetCredentialsFromEnvCompute(self):
   self.mox.StubOutWithMock(gce, 'AppAssertionCredentials')
   credentials = self.mox.CreateMockAnything()
   gce.AppAssertionCredentials(connection.SCOPE).AndReturn(credentials)
   credentials.authorize(mox.IsA(httplib2.Http))
   credentials.refresh(mox.IsA(httplib2.Http))
   self.mox.ReplayAll()
   self.assertIs(credentials, helper.get_credentials_from_env())
   self.mox.VerifyAll()
 def testGetCredentialsFromEnvCompute(self):
     self.mox.StubOutWithMock(gce, 'AppAssertionCredentials')
     credentials = self.mox.CreateMockAnything()
     gce.AppAssertionCredentials(connection.SCOPE).AndReturn(credentials)
     credentials.authorize(mox.IsA(httplib2.Http))
     credentials.refresh(mox.IsA(httplib2.Http))
     self.mox.ReplayAll()
     self.assertIs(credentials, helper.get_credentials_from_env())
     self.mox.VerifyAll()
 def testGetCredentialsFromEnvJwt(self):
   self.mox.StubOutWithMock(os, 'getenv')
   self.mox.StubOutWithMock(client, 'SignedJwtAssertionCredentials')
   credentials = self.mox.CreateMockAnything()
   os.getenv('DATASTORE_SERVICE_ACCOUNT').AndReturn('*****@*****.**')
   os.getenv('DATASTORE_PRIVATE_KEY_FILE').AndReturn(self.certificate.name)
   client.SignedJwtAssertionCredentials('*****@*****.**',
                                        'not-a-secret-key',
                                        connection.SCOPE).AndReturn(
                                            credentials)
   self.mox.ReplayAll()
   self.assertIs(credentials, helper.get_credentials_from_env())
   self.mox.VerifyAll()
 def testGetCredentialsFromEnvJwt(self):
     self.mox.StubOutWithMock(os, 'getenv')
     self.mox.StubOutWithMock(client, 'SignedJwtAssertionCredentials')
     credentials = self.mox.CreateMockAnything()
     os.getenv('DATASTORE_SERVICE_ACCOUNT').AndReturn('*****@*****.**')
     os.getenv('DATASTORE_PRIVATE_KEY_FILE').AndReturn(
         self.certificate.name)
     client.SignedJwtAssertionCredentials(
         '*****@*****.**', 'not-a-secret-key',
         connection.SCOPE).AndReturn(credentials)
     self.mox.ReplayAll()
     self.assertIs(credentials, helper.get_credentials_from_env())
     self.mox.VerifyAll()
def get_default_connection():
  """Return the default datastore connection.

  Defaults dataset to os.getenv('DATASTORE_DATASET'), host to
  os.getenv('DATASTORE_HOST'), and credentials to
  helper.get_credentials_from_env().

  Use set_options to override defaults.
  """
  global _conn
  if not _conn:
    if 'dataset' not in _options:
      _options['dataset'] = os.getenv('DATASTORE_DATASET')
    if 'host' not in _options:
      _options['host'] = os.getenv('DATASTORE_HOST')
    if 'credentials' not in _options:
      _options['credentials'] = helper.get_credentials_from_env()
    _conn = Datastore(**_options)
  return _conn
def get_default_connection():
  """Return the default datastore connection.

  Defaults dataset to os.getenv('DATASTORE_DATASET'), host to
  os.getenv('DATASTORE_HOST'), and credentials to
  helper.get_credentials_from_env().

  Use set_options to override defaults.
  """
  global _conn
  if not _conn:
    if 'dataset' not in _options:
      _options['dataset'] = os.getenv('DATASTORE_DATASET')
    if 'host' not in _options:
      _options['host'] = os.getenv('DATASTORE_HOST')
    if 'credentials' not in _options:
      _options['credentials'] = helper.get_credentials_from_env()
    _conn = Datastore(**_options)
  return _conn
    def testSetOptions(self):
        other_thread_conn = []
        lock1 = threading.Lock()
        lock2 = threading.Lock()
        lock1.acquire()
        lock2.acquire()

        def target():
            # Grab two connections
            other_thread_conn.append(datastore.get_default_connection())
            other_thread_conn.append(datastore.get_default_connection())
            lock1.release(
            )  # Notify that we have grabbed the first 2 connections.
            lock2.acquire()  # Wait for the signal to grab the 3rd.
            other_thread_conn.append(datastore.get_default_connection())

        other_thread = threading.Thread(target=target)

        # Resetting options and state.
        datastore._options = {}
        datastore.set_options(dataset='foo')

        self.mox.StubOutWithMock(os, 'getenv')
        self.mox.StubOutWithMock(helper, 'get_credentials_from_env')
        os.getenv('DATASTORE_HOST').AndReturn('http://localhost:8080')
        os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)
        os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)
        os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)
        os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)

        helper.get_credentials_from_env().AndReturn(FakeCredentialsFromEnv())
        self.mox.ReplayAll()

        # Start the thread and wait for the first lock.
        other_thread.start()
        lock1.acquire()

        t1_conn1 = datastore.get_default_connection()
        t2_conn1, t2_conn1b = other_thread_conn
        other_thread_conn = []
        # The two threads get different connections.
        self.assertIsNot(t1_conn1, t2_conn1)
        # Multiple calls on the same thread get the same connection.
        self.assertIs(t1_conn1, datastore.get_default_connection())
        self.assertIs(t2_conn1, t2_conn1b)

        # Change the global options and grab the connections again.
        datastore.set_options(dataset='bar')
        lock2.release()
        other_thread.join()
        t1_conn2 = datastore.get_default_connection()
        t2_conn2 = other_thread_conn[0]

        # Changing the options causes all threads to create new connections.
        self.assertIsNot(t1_conn1, t1_conn2)
        self.assertIsNot(t2_conn1, t2_conn2)
        # The new connections are still different for each thread.
        self.assertIsNot(t1_conn2, t2_conn2)
        # The old connections has the old settings.
        self.assertEqual(
            'http://localhost:8080/datastore/v1beta2/datasets/foo/',
            t1_conn1._url)
        self.assertEqual(
            'http://localhost:8080/datastore/v1beta2/datasets/foo/',
            t2_conn1._url)
        # The new connections has the new settings.
        self.assertEqual(
            'http://localhost:8080/datastore/v1beta2/datasets/bar/',
            t1_conn2._url)
        self.assertEqual(
            'http://localhost:8080/datastore/v1beta2/datasets/bar/',
            t2_conn2._url)
        self.assertEqual(FakeCredentialsFromEnv, type(t1_conn2._credentials))
        self.assertEqual(FakeCredentialsFromEnv, type(t2_conn2._credentials))
        self.mox.VerifyAll()
    def testSetOptions(self):
        other_thread_conn = []
        lock1 = threading.Lock()
        lock2 = threading.Lock()
        lock1.acquire()
        lock2.acquire()

        def target():
            # Grab two connections
            other_thread_conn.append(datastore.get_default_connection())
            other_thread_conn.append(datastore.get_default_connection())
            lock1.release()  # Notify that we have grabbed the first 2 connections.
            lock2.acquire()  # Wait for the signal to grab the 3rd.
            other_thread_conn.append(datastore.get_default_connection())

        other_thread = threading.Thread(target=target)

        # Resetting options and state.
        datastore._options = {}
        datastore.set_options(project_id="foo")

        self.mox.StubOutWithMock(helper, "get_credentials_from_env")
        self.mox.StubOutWithMock(helper, "get_project_endpoint_from_env")
        endpoint = "http://localhost:8080/datastore/v1/projects/%s"
        helper.get_project_endpoint_from_env(project_id="foo").AndReturn(endpoint % "foo")
        helper.get_project_endpoint_from_env(project_id="foo").AndReturn(endpoint % "foo")
        helper.get_project_endpoint_from_env(project_id="bar").AndReturn(endpoint % "bar")
        helper.get_project_endpoint_from_env(project_id="bar").AndReturn(endpoint % "bar")

        helper.get_credentials_from_env().AndReturn(FakeCredentialsFromEnv())
        self.mox.ReplayAll()

        # Start the thread and wait for the first lock.
        other_thread.start()
        lock1.acquire()

        t1_conn1 = datastore.get_default_connection()
        t2_conn1, t2_conn1b = other_thread_conn
        other_thread_conn = []
        # The two threads get different connections.
        self.assertIsNot(t1_conn1, t2_conn1)
        # Multiple calls on the same thread get the same connection.
        self.assertIs(t1_conn1, datastore.get_default_connection())
        self.assertIs(t2_conn1, t2_conn1b)

        # Change the global options and grab the connections again.
        datastore.set_options(project_id="bar")
        lock2.release()
        other_thread.join()
        t1_conn2 = datastore.get_default_connection()
        t2_conn2 = other_thread_conn[0]

        # Changing the options causes all threads to create new connections.
        self.assertIsNot(t1_conn1, t1_conn2)
        self.assertIsNot(t2_conn1, t2_conn2)
        # The new connections are still different for each thread.
        self.assertIsNot(t1_conn2, t2_conn2)
        # The old connections has the old settings.
        self.assertEqual("http://localhost:8080/datastore/v1/projects/foo", t1_conn1._url)
        self.assertEqual("http://localhost:8080/datastore/v1/projects/foo", t2_conn1._url)
        # The new connections has the new settings.
        self.assertEqual("http://localhost:8080/datastore/v1/projects/bar", t1_conn2._url)
        self.assertEqual("http://localhost:8080/datastore/v1/projects/bar", t2_conn2._url)
        self.assertEqual(FakeCredentialsFromEnv, type(t1_conn2._credentials))
        self.assertEqual(FakeCredentialsFromEnv, type(t2_conn2._credentials))
        self.mox.VerifyAll()
  def testSetOptions(self):
    other_thread_conn = []
    lock1 = threading.Lock()
    lock2 = threading.Lock()
    lock1.acquire()
    lock2.acquire()
    def target():
      # Grab two connections
      other_thread_conn.append(datastore.get_default_connection())
      other_thread_conn.append(datastore.get_default_connection())
      lock1.release()  # Notify that we have grabbed the first 2 connections.
      lock2.acquire()  # Wait for the signal to grab the 3rd.
      other_thread_conn.append(datastore.get_default_connection())
    other_thread = threading.Thread(target=target)

    # Resetting options and state.
    datastore._options = {}
    datastore.set_options(dataset='foo')

    self.mox.StubOutWithMock(os, 'getenv')
    self.mox.StubOutWithMock(helper, 'get_credentials_from_env')
    os.getenv('DATASTORE_HOST').AndReturn('http://localhost:8080')
    os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)
    os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)
    os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)
    os.getenv('DATASTORE_URL_INTERNAL_OVERRIDE').AndReturn(None)

    helper.get_credentials_from_env().AndReturn(FakeCredentialsFromEnv())
    self.mox.ReplayAll()

    # Start the thread and wait for the first lock.
    other_thread.start()
    lock1.acquire()

    t1_conn1 = datastore.get_default_connection()
    t2_conn1, t2_conn1b = other_thread_conn
    other_thread_conn = []
    # The two threads get different connections.
    self.assertIsNot(t1_conn1, t2_conn1)
    # Multiple calls on the same thread get the same connection.
    self.assertIs(t1_conn1, datastore.get_default_connection())
    self.assertIs(t2_conn1, t2_conn1b)

    # Change the global options and grab the connections again.
    datastore.set_options(dataset='bar')
    lock2.release()
    other_thread.join()
    t1_conn2 = datastore.get_default_connection()
    t2_conn2 = other_thread_conn[0]

    # Changing the options causes all threads to create new connections.
    self.assertIsNot(t1_conn1, t1_conn2)
    self.assertIsNot(t2_conn1, t2_conn2)
    # The new connections are still different for each thread.
    self.assertIsNot(t1_conn2, t2_conn2)
    # The old connections has the old settings.
    self.assertEqual('http://localhost:8080/datastore/v1beta2/datasets/foo/',
                     t1_conn1._url)
    self.assertEqual('http://localhost:8080/datastore/v1beta2/datasets/foo/',
                     t2_conn1._url)
    # The new connections has the new settings.
    self.assertEqual('http://localhost:8080/datastore/v1beta2/datasets/bar/',
                     t1_conn2._url)
    self.assertEqual('http://localhost:8080/datastore/v1beta2/datasets/bar/',
                     t2_conn2._url)
    self.assertEqual(FakeCredentialsFromEnv, type(t1_conn2._credentials))
    self.assertEqual(FakeCredentialsFromEnv, type(t2_conn2._credentials))
    self.mox.VerifyAll()
def _make_cloud_datastore_context(app_id, external_app_ids=(), host=None):
    """Creates a new context to connect to a remote Cloud Datastore instance.

  This should only be used outside of Google App Engine.

  Args:
    app_id: The application id to connect to. This differs from the project
      id as it may have an additional prefix, e.g. "s~" or "e~".
    external_app_ids: A list of apps that may be referenced by data in your
      application. For example, if you are connected to s~my-app and store keys
      for s~my-other-app, you should include s~my-other-app in the external_apps
      list.
    host: The hostname to provide to the datastore connection. If None, the
      default is used.
  Returns:
    An ndb.Context that can connect to a Remote Cloud Datastore. You can use
    this context by passing it to ndb.set_context.
  """
    from . import model  # Late import to deal with circular imports.
    # Late import since it might not exist.
    if not datastore_pbs._CLOUD_DATASTORE_ENABLED:
        raise datastore_errors.BadArgumentError(
            datastore_pbs.MISSING_CLOUD_DATASTORE_MESSAGE)
    import googledatastore
    try:
        from google.appengine.datastore import cloud_datastore_v1_remote_stub
    except ImportError:
        from google3.apphosting.datastore import cloud_datastore_v1_remote_stub

    current_app_id = os.environ.get('APPLICATION_ID', None)
    if current_app_id and current_app_id != app_id:
        # TODO(pcostello): We should support this so users can connect to different
        # applications.
        raise ValueError(
            'Cannot create a Cloud Datastore context that connects '
            'to an application (%s) that differs from the application '
            'already connected to (%s).' % (app_id, current_app_id))
    os.environ['APPLICATION_ID'] = app_id

    id_resolver = datastore_pbs.IdResolver((app_id, ) +
                                           tuple(external_app_ids))
    project_id = id_resolver.resolve_project_id(app_id)

    from googledatastore import helper
    credentials = helper.get_credentials_from_env()
    datastore = googledatastore.Datastore(project_id,
                                          credentials=credentials,
                                          host=host)

    conn = model.make_connection(
        _api_version=datastore_rpc._CLOUD_DATASTORE_V1,
        _id_resolver=id_resolver)

    # If necessary, install the stubs
    try:
        stub = cloud_datastore_v1_remote_stub.CloudDatastoreV1RemoteStub(
            datastore)
        apiproxy_stub_map.apiproxy.RegisterStub(
            datastore_rpc._CLOUD_DATASTORE_V1, stub)
    except:
        pass  # The stub is already installed.
    # TODO(pcostello): Ensure the current stub is connected to the right project.
    return make_context(conn=conn)