From 2fe5be31784a036062180f9c0f2c7b5eda978123 Mon Sep 17 00:00:00 2001
From: Chris Rossi <chris@archimedeanco.com>
Date: Thu, 9 Jan 2020 14:06:36 -0500
Subject: [PATCH] fix: Handle `int` for DateTimeProperty (#285)

In Datastore, projection queries involving entities with DateTime
properties return integer timestamps instead of `datetime.datetime`
objects. This fix handles that case and returns `datetime.datetime`
objects regardless of the query type.

Fixes #261.
---
 google/cloud/ndb/model.py  | 10 +++++++++-
 tests/system/test_query.py | 34 ++++++++++++++++++++++++++++++++++
 tests/unit/test_model.py   |  8 ++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py
index b23b2379..19048c4d 100644
--- a/google/cloud/ndb/model.py
+++ b/google/cloud/ndb/model.py
@@ -3601,7 +3601,10 @@ def _from_base_type(self, value):
         """Convert a value from the "base" value type for this property.
 
         Args:
-            value (datetime.datetime): The value to be converted.
+            value (Union[int, datetime.datetime]): The value to be converted.
+                The value will be `int` for entities retrieved by a projection
+                query and is a timestamp as the number of nanoseconds since the
+                epoch.
 
         Returns:
             Optional[datetime.datetime]: If ``tzinfo`` is set on this property,
@@ -3609,6 +3612,11 @@ def _from_base_type(self, value):
                 returns the value without ``tzinfo`` or ``None`` if value did
                 not have ``tzinfo`` set.
         """
+        if isinstance(value, six.integer_types):
+            # Projection query, value is integer nanoseconds
+            seconds = value / 1e6
+            value = datetime.datetime.fromtimestamp(seconds, pytz.utc)
+
         if self._tzinfo is not None:
             return value.astimezone(self._tzinfo)
 
diff --git a/tests/system/test_query.py b/tests/system/test_query.py
index 0908c631..582b766d 100644
--- a/tests/system/test_query.py
+++ b/tests/system/test_query.py
@@ -16,11 +16,13 @@
 System tests for queries.
 """
 
+import datetime
 import functools
 import operator
 
 import grpc
 import pytest
+import pytz
 
 import test_utils.system
 
@@ -194,6 +196,38 @@ class SomeKind(ndb.Model):
         results[1].bar
 
 
+@pytest.mark.usefixtures("client_context")
+def test_projection_datetime(ds_entity):
+    """Regression test for Issue #261
+
+    https://github.com/googleapis/python-ndb/issues/261
+    """
+    entity_id = test_utils.system.unique_resource_id()
+    ds_entity(
+        KIND,
+        entity_id,
+        foo=datetime.datetime(2010, 5, 12, 2, 42, tzinfo=pytz.UTC),
+    )
+    entity_id = test_utils.system.unique_resource_id()
+    ds_entity(
+        KIND,
+        entity_id,
+        foo=datetime.datetime(2010, 5, 12, 2, 43, tzinfo=pytz.UTC),
+    )
+
+    class SomeKind(ndb.Model):
+        foo = ndb.DateTimeProperty()
+        bar = ndb.StringProperty()
+
+    query = SomeKind.query(projection=("foo",))
+    results = eventually(query.fetch, _length_equals(2))
+
+    results = sorted(results, key=operator.attrgetter("foo"))
+
+    assert results[0].foo == datetime.datetime(2010, 5, 12, 2, 42)
+    assert results[1].foo == datetime.datetime(2010, 5, 12, 2, 43)
+
+
 @pytest.mark.usefixtures("client_context")
 def test_distinct_on(ds_entity):
     for i in range(6):
diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py
index ff3a0616..11e161ae 100644
--- a/tests/unit/test_model.py
+++ b/tests/unit/test_model.py
@@ -2707,6 +2707,14 @@ def test__from_base_type_convert_timezone():
             2010, 5, 11, 20, tzinfo=timezone(-4)
         )
 
+    @staticmethod
+    def test__from_base_type_int():
+        prop = model.DateTimeProperty(name="dt_val")
+        value = 1273632120000000
+        assert prop._from_base_type(value) == datetime.datetime(
+            2010, 5, 12, 2, 42
+        )
+
     @staticmethod
     def test__to_base_type_noop():
         prop = model.DateTimeProperty(name="dt_val", tzinfo=timezone(-4))