This is an automated email from the ASF dual-hosted git repository.
udim pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 96d947a [BEAM-7828] Fixes Key type conversion from/to client entity
in Python Datastore IO (#9174)
96d947a is described below
commit 96d947a5d2df9260e1fed1087feb1e3578c6dae2
Author: Ludovic Post <[email protected]>
AuthorDate: Thu Aug 1 17:17:13 2019 -0700
[BEAM-7828] Fixes Key type conversion from/to client entity in Python
Datastore IO (#9174)
---
.../io/gcp/datastore/v1new/datastoreio.py | 11 ++++++++--
.../io/gcp/datastore/v1new/datastoreio_test.py | 16 ++++++++++++++
.../apache_beam/io/gcp/datastore/v1new/types.py | 25 +++++++++++++++-------
.../io/gcp/datastore/v1new/types_test.py | 13 +++++++++--
4 files changed, 53 insertions(+), 12 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
index c04daa2..7ecd1fc 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
@@ -365,7 +365,9 @@ class WriteToDatastore(_Mutate):
:class:`~apache_beam.io.gcp.datastore.v1new.types.Entity` to Cloud Datastore.
Entity keys must be complete. The ``project`` field in each key must match
the
- project ID passed to this transform.
+ project ID passed to this transform. If ``project`` field in entity or
+ property key is empty then it is filled with the project ID passed to this
+ transform.
"""
def __init__(self, project):
@@ -382,6 +384,8 @@ class WriteToDatastore(_Mutate):
if not isinstance(element, types.Entity):
raise
ValueError('apache_beam.io.gcp.datastore.v1new.datastoreio.Entity'
' expected, got: %s' % type(element))
+ if not element.key.project:
+ element.key.project = self._project
client_entity = element.to_client_entity()
if client_entity.key.is_partial:
raise ValueError('Entities to be written to Cloud Datastore must '
@@ -403,7 +407,8 @@ class DeleteFromDatastore(_Mutate):
Datastore.
Keys must be complete. The ``project`` field in each key must match the
- project ID passed to this transform.
+ project ID passed to this transform. If ``project`` field in key is empty
then
+ it is filled with the project ID passed to this transform.
"""
def __init__(self, project):
"""Initialize the `DeleteFromDatastore` transform.
@@ -420,6 +425,8 @@ class DeleteFromDatastore(_Mutate):
if not isinstance(element, types.Key):
raise ValueError('apache_beam.io.gcp.datastore.v1new.datastoreio.Key'
' expected, got: %s' % type(element))
+ if not element.project:
+ element.project = self._project
client_key = element.to_client_key()
if client_key.is_partial:
raise ValueError('Keys to be deleted from Cloud Datastore must be '
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
index e3975c7..79d43fe 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
@@ -37,6 +37,7 @@ try:
from apache_beam.io.gcp.datastore.v1new.datastoreio import
DeleteFromDatastore
from apache_beam.io.gcp.datastore.v1new.datastoreio import ReadFromDatastore
from apache_beam.io.gcp.datastore.v1new.datastoreio import WriteToDatastore
+ from apache_beam.io.gcp.datastore.v1new.types import Key
from google.cloud.datastore import client
from google.cloud.datastore import entity
from google.cloud.datastore import helpers
@@ -203,6 +204,14 @@ class DatastoreioTest(DatastoreioTestBase):
entities = helper.create_entities(num_entities)
expected_entities = [entity.to_client_entity() for entity in entities]
+ # Infer project from write fn project arg.
+ if num_entities:
+ key = Key(['k1', 1234], project=self._PROJECT)
+ expected_key = key.to_client_key()
+ key.project = None
+ entities[0].key = key
+ expected_entities[0].key = expected_key
+
all_batch_entities = []
commit_count = [0]
self._mock_client.batch.side_effect = (
@@ -274,6 +283,13 @@ class DatastoreioTest(DatastoreioTestBase):
keys = [entity.key for entity in helper.create_entities(10)]
expected_keys = [key.to_client_key() for key in keys]
+ # Infer project from delete fn project arg.
+ key = Key(['k1', 1234], project=self._PROJECT)
+ expected_key = key.to_client_key()
+ key.project = None
+ keys.append(key)
+ expected_keys.append(expected_key)
+
all_batch_keys = []
self._mock_client.batch.side_effect = (
lambda: FakeBatch(all_batch_items=all_batch_keys))
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
b/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
index c80fe04..d162bf3 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
@@ -140,6 +140,8 @@ class Key(object):
return False
if self.path_elements != other.path_elements:
return False
+ if self.project != other.project:
+ return False
if self.parent is not None and other.parent is not None:
return self.parent == other.parent
@@ -181,21 +183,28 @@ class Entity(object):
@staticmethod
def from_client_entity(client_entity):
- key = Key.from_client_key(client_entity.key)
- entity = Entity(
- key, exclude_from_indexes=set(client_entity.exclude_from_indexes))
- entity.set_properties(client_entity)
- return entity
+ res = Entity(
+ Key.from_client_key(client_entity.key),
+ exclude_from_indexes=set(client_entity.exclude_from_indexes))
+ for name, value in client_entity.items():
+ if isinstance(value, key.Key):
+ value = Key.from_client_key(value)
+ res.properties[name] = value
+ return res
def to_client_entity(self):
"""
Returns a :class:`google.cloud.datastore.entity.Entity` instance that
represents this entity.
"""
- key = self.key.to_client_key()
- res = entity.Entity(key=key,
+ res = entity.Entity(key=self.key.to_client_key(),
exclude_from_indexes=tuple(self.exclude_from_indexes))
- res.update(self.properties)
+ for name, value in self.properties.items():
+ if isinstance(value, Key):
+ if not value.project:
+ value.project = self.key.project
+ value = value.to_client_key()
+ res[name] = value
return res
def __eq__(self, other):
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
b/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
index 7ba82c5..aa6464b 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
@@ -51,18 +51,23 @@ class TypesTest(unittest.TestCase):
kc = k.to_client_key()
exclude_from_indexes = ('efi1', 'efi2')
e = Entity(k, exclude_from_indexes=exclude_from_indexes)
- e.set_properties({'efi1': 'value', 'property': 'value'})
+ ref = Key(['kind2', 1235])
+ e.set_properties({'efi1': 'value', 'property': 'value', 'ref': ref})
ec = e.to_client_entity()
self.assertEqual(kc, ec.key)
self.assertSetEqual(set(exclude_from_indexes), ec.exclude_from_indexes)
self.assertEqual('kind', ec.kind)
self.assertEqual(1234, ec.id)
+ self.assertEqual('kind2', ec['ref'].kind)
+ self.assertEqual(1235, ec['ref'].id)
+ self.assertEqual(self._PROJECT, ec['ref'].project)
def testEntityFromClientEntity(self):
k = Key(['kind', 1234], project=self._PROJECT)
exclude_from_indexes = ('efi1', 'efi2')
e = Entity(k, exclude_from_indexes=exclude_from_indexes)
- e.set_properties({'efi1': 'value', 'property': 'value'})
+ ref = Key(['kind2', 1235])
+ e.set_properties({'efi1': 'value', 'property': 'value', 'ref': ref})
efc = Entity.from_client_entity(e.to_client_entity())
self.assertEqual(e, efc)
@@ -101,6 +106,10 @@ class TypesTest(unittest.TestCase):
kfc3 = Key.from_client_key(kfc2.to_client_key())
self.assertEqual(kfc2, kfc3)
+ kfc4 = Key.from_client_key(kfc2.to_client_key())
+ kfc4.project = 'other'
+ self.assertNotEqual(kfc2, kfc4)
+
def testKeyFromClientKeyNoNamespace(self):
k = Key(['k1', 1234], project=self._PROJECT)
ck = k.to_client_key()
