beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From u...@apache.org
Subject [beam] branch master updated: [BEAM-7828] Fixes Key type conversion from/to client entity in Python Datastore IO (#9174)
Date Fri, 02 Aug 2019 00:17:28 GMT
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 <ludo@thepost.family>
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()


Mime
View raw message