test_system.py

# Copyright 2014 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
import os
import string
import unittest

import requests
import six

from google.cloud._helpers import UTC
from google.cloud import datastore
from google.cloud.datastore.helpers import GeoPoint
from google.cloud.environment_vars import GCD_DATASET
from google.cloud.exceptions import Conflict

from test_utils.system import EmulatorCreds
from test_utils.system import unique_resource_id

from tests.system.utils import clear_datastore
from tests.system.utils import populate_datastore


class Config(object):
    """Run-time configuration to be modified at set-up.

    This is a mutable stand-in to allow test set-up to modify
    global state.
    """

    CLIENT = None
    TO_DELETE = []


def clone_client(client):
    return datastore.Client(
        project=client.project,
        namespace=client.namespace,
        credentials=client._credentials,
        _http=client._http,
    )


def setUpModule():
    emulator_dataset = os.getenv(GCD_DATASET)
    # Isolated namespace so concurrent test runs don't collide.
    test_namespace = "ns" + unique_resource_id()
    if emulator_dataset is None:
        Config.CLIENT = datastore.Client(namespace=test_namespace)
    else:
        credentials = EmulatorCreds()
        http = requests.Session()  # Un-authorized.
        Config.CLIENT = datastore.Client(
            project=emulator_dataset,
            namespace=test_namespace,
            credentials=credentials,
            _http=http,
        )


def tearDownModule():
    keys = [entity.key for entity in Config.TO_DELETE]
    with Config.CLIENT.transaction():
        Config.CLIENT.delete_multi(keys)


class TestDatastore(unittest.TestCase):
    def setUp(self):
        self.case_entities_to_delete = []

    def tearDown(self):
        with Config.CLIENT.transaction():
            keys = [entity.key for entity in self.case_entities_to_delete]
            Config.CLIENT.delete_multi(keys)


class TestDatastoreAllocateIDs(TestDatastore):
    def test_allocate_ids(self):
        num_ids = 10
        allocated_keys = Config.CLIENT.allocate_ids(Config.CLIENT.key("Kind"), num_ids)
        self.assertEqual(len(allocated_keys), num_ids)

        unique_ids = set()
        for key in allocated_keys:
            unique_ids.add(key.id)
            self.assertIsNone(key.name)
            self.assertNotEqual(key.id, None)

        self.assertEqual(len(unique_ids), num_ids)


class TestDatastoreSave(TestDatastore):
    @classmethod
    def setUpClass(cls):
        cls.PARENT = Config.CLIENT.key("Blog", "PizzaMan")

    def _get_post(self, id_or_name=None, post_content=None):
        post_content = post_content or {
            "title": u"How to make the perfect pizza in your grill",
            "tags": [u"pizza", u"grill"],
            "publishedAt": datetime.datetime(2001, 1, 1, tzinfo=UTC),
            "author": u"Silvano",
            "isDraft": False,
            "wordCount": 400,
            "rating": 5.0,
        }
        # Create an entity with the given content.
        # NOTE: Using a parent to ensure consistency for query
        #       in `test_empty_kind`.
        key = Config.CLIENT.key("Post", parent=self.PARENT)
        entity = datastore.Entity(key=key)
        entity.update(post_content)

        # Update the entity key.
        if id_or_name is not None:
            entity.key = entity.key.completed_key(id_or_name)

        return entity

    def _generic_test_post(self, name=None, key_id=None):
        entity = self._get_post(id_or_name=(name or key_id))
        Config.CLIENT.put(entity)

        # Register entity to be deleted.
        self.case_entities_to_delete.append(entity)

        if name is not None:
            self.assertEqual(entity.key.name, name)
        if key_id is not None:
            self.assertEqual(entity.key.id, key_id)
        retrieved_entity = Config.CLIENT.get(entity.key)
        # Check the given and retrieved are the the same.
        self.assertEqual(retrieved_entity, entity)

    def test_post_with_name(self):
        self._generic_test_post(name="post1")

    def test_post_with_id(self):
        self._generic_test_post(key_id=123456789)

    def test_post_with_generated_id(self):
        self._generic_test_post()

    def test_save_multiple(self):
        with Config.CLIENT.transaction() as xact:
            entity1 = self._get_post()
            xact.put(entity1)
            # Register entity to be deleted.
            self.case_entities_to_delete.append(entity1)

            second_post_content = {
                "title": u"How to make the perfect homemade pasta",
                "tags": [u"pasta", u"homemade"],
                "publishedAt": datetime.datetime(2001, 1, 1),
                "author": u"Silvano",
                "isDraft": False,
                "wordCount": 450,
                "rating": 4.5,
            }
            entity2 = self._get_post(post_content=second_post_content)
            xact.put(entity2)
            # Register entity to be deleted.
            self.case_entities_to_delete.append(entity2)

        keys = [entity1.key, entity2.key]
        matches = Config.CLIENT.get_multi(keys)
        self.assertEqual(len(matches), 2)

    def test_empty_kind(self):
        query = Config.CLIENT.query(kind="Post")
        query.ancestor = self.PARENT
        posts = list(query.fetch(limit=2))
        self.assertEqual(posts, [])

    def test_all_value_types(self):
        key = Config.CLIENT.key("TestPanObject", 1234)
        entity = datastore.Entity(key=key)
        entity["timestamp"] = datetime.datetime(2014, 9, 9, tzinfo=UTC)
        key_stored = Config.CLIENT.key("SavedKey", "right-here")
        entity["key"] = key_stored
        entity["truthy"] = True
        entity["float"] = 2.718281828
        entity["int"] = 3735928559
        entity["words"] = u"foo"
        entity["blob"] = b"seekretz"
        entity_stored = datastore.Entity(key=key_stored)
        entity_stored["hi"] = "bye"
        entity["nested"] = entity_stored
        entity["items"] = [1, 2, 3]
        entity["geo"] = GeoPoint(1.0, 2.0)
        entity["nothing_here"] = None

        # Store the entity.
        self.case_entities_to_delete.append(entity)
        Config.CLIENT.put(entity)

        # Check the original and retrieved are the the same.
        retrieved_entity = Config.CLIENT.get(entity.key)
        self.assertEqual(retrieved_entity, entity)


class TestDatastoreSaveKeys(TestDatastore):
    def test_save_key_self_reference(self):
        parent_key = Config.CLIENT.key("Residence", "NewYork")
        key = Config.CLIENT.key("Person", "name", parent=parent_key)
        entity = datastore.Entity(key=key)
        entity["fullName"] = u"Full name"
        entity["linkedTo"] = key  # Self reference.

        Config.CLIENT.put(entity)
        self.case_entities_to_delete.append(entity)

        query = Config.CLIENT.query(kind="Person")
        # Adding ancestor to ensure consistency.
        query.ancestor = parent_key
        query.add_filter("linkedTo", "=", key)

        stored_persons = list(query.fetch(limit=2))
        self.assertEqual(stored_persons, [entity])


class TestDatastoreQuery(TestDatastore):
    @classmethod
    def setUpClass(cls):
        cls.CLIENT = clone_client(Config.CLIENT)
        # Remove the namespace from the cloned client, since these
        # query tests rely on the entities to be already stored and indexed,
        # hence ``test_namespace`` set at runtime can't be used.
        cls.CLIENT.namespace = None

        # In the emulator, re-populating the datastore is cheap.
        if os.getenv(GCD_DATASET) is not None:
            # Populate the datastore with the cloned client.
            populate_datastore.add_characters(client=cls.CLIENT)

        cls.CHARACTERS = populate_datastore.CHARACTERS
        # Use the client for this test instead of the global.
        cls.ANCESTOR_KEY = cls.CLIENT.key(*populate_datastore.ANCESTOR)

    @classmethod
    def tearDownClass(cls):
        # In the emulator, destroy the query entities.
        if os.getenv(GCD_DATASET) is not None:
            # Use the client for this test instead of the global.
            clear_datastore.remove_all_entities(client=cls.CLIENT)

    def _base_query(self):
        # Use the client for this test instead of the global.
        return self.CLIENT.query(kind="Character", ancestor=self.ANCESTOR_KEY)

    def test_limit_queries(self):
        limit = 5
        query = self._base_query()

        # Fetch characters.
        iterator = query.fetch(limit=limit)
        page = six.next(iterator.pages)
        character_entities = list(page)
        cursor = iterator.next_page_token
        self.assertEqual(len(character_entities), limit)

        # Check cursor after fetch.
        self.assertIsNotNone(cursor)

        # Fetch remaining characters.
        new_character_entities = list(query.fetch(start_cursor=cursor))
        characters_remaining = len(self.CHARACTERS) - limit
        self.assertEqual(len(new_character_entities), characters_remaining)

    def test_query_simple_filter(self):
        query = self._base_query()
        query.add_filter("appearances", ">=", 20)
        expected_matches = 6
        # We expect 6, but allow the query to get 1 extra.
        entities = list(query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

    def test_query_multiple_filters(self):
        query = self._base_query()
        query.add_filter("appearances", ">=", 26)
        query.add_filter("family", "=", "Stark")
        expected_matches = 4
        # We expect 4, but allow the query to get 1 extra.
        entities = list(query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

    def test_ancestor_query(self):
        filtered_query = self._base_query()

        expected_matches = 8
        # We expect 8, but allow the query to get 1 extra.
        entities = list(filtered_query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

    def test_query_key_filter(self):
        # Use the client for this test instead of the global.
        rickard_key = self.CLIENT.key(*populate_datastore.RICKARD)

        query = self._base_query()
        query.key_filter(rickard_key)
        expected_matches = 1
        # We expect 1, but allow the query to get 1 extra.
        entities = list(query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

    def test_ordered_query(self):
        query = self._base_query()
        query.order = "appearances"
        expected_matches = 8
        # We expect 8, but allow the query to get 1 extra.
        entities = list(query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

        # Actually check the ordered data returned.
        self.assertEqual(entities[0]["name"], self.CHARACTERS[0]["name"])
        self.assertEqual(entities[7]["name"], self.CHARACTERS[3]["name"])

    def test_projection_query(self):
        filtered_query = self._base_query()
        filtered_query.projection = ["name", "family"]
        filtered_query.order = ["name", "family"]

        # NOTE: There are 9 responses because of Catelyn. She has both
        #       Stark and Tully as her families, hence occurs twice in
        #       the results.
        expected_matches = 9
        # We expect 9, but allow the query to get 1 extra.
        entities = list(filtered_query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

        arya_entity = entities[0]
        catelyn_tully_entity = entities[3]
        sansa_entity = entities[8]

        arya_dict = dict(arya_entity)
        self.assertEqual(arya_dict, {"name": "Arya", "family": "Stark"})

        catelyn_stark_entity = entities[2]
        catelyn_stark_dict = dict(catelyn_stark_entity)
        self.assertEqual(catelyn_stark_dict, {"name": "Catelyn", "family": "Stark"})

        catelyn_tully_dict = dict(catelyn_tully_entity)
        self.assertEqual(catelyn_tully_dict, {"name": "Catelyn", "family": "Tully"})

        # Check both Catelyn keys are the same.
        self.assertEqual(catelyn_stark_entity.key, catelyn_tully_entity.key)

        sansa_dict = dict(sansa_entity)
        self.assertEqual(sansa_dict, {"name": "Sansa", "family": "Stark"})

    def test_query_paginate_simple_uuid_keys(self):

        # See issue #4264
        page_query = self.CLIENT.query(kind="uuid_key")
        iterator = page_query.fetch()

        seen = set()
        page_count = 0
        for page in iterator.pages:
            page_count += 1
            for entity in page:
                uuid_str = entity.key.name
                self.assertNotIn(uuid_str, seen, uuid_str)
                seen.add(uuid_str)

        self.assertTrue(page_count > 1)

    def test_query_paginate_simple_timestamp_keys(self):

        # See issue #4264
        page_query = self.CLIENT.query(kind="timestamp_key")
        iterator = page_query.fetch()

        seen = set()
        page_count = 0
        for page in iterator.pages:
            page_count += 1
            for entity in page:
                timestamp = entity.key.id
                self.assertNotIn(timestamp, seen, timestamp)
                seen.add(timestamp)

        self.assertTrue(page_count > 1)

    def test_query_offset_timestamp_keys(self):
        # See issue #4675
        max_all = 10000
        offset = 1
        max_offset = max_all - offset
        query = self.CLIENT.query(kind="timestamp_key")
        all_w_limit = list(query.fetch(limit=max_all))
        self.assertEqual(len(all_w_limit), max_all)

        offset_w_limit = list(query.fetch(offset=offset, limit=max_offset))
        self.assertEqual(offset_w_limit, all_w_limit[offset:])

    def test_query_paginate_with_offset(self):
        page_query = self._base_query()
        page_query.order = "appearances"
        offset = 2
        limit = 3
        iterator = page_query.fetch(limit=limit, offset=offset)

        # Fetch characters.
        page = six.next(iterator.pages)
        entities = list(page)
        cursor = iterator.next_page_token
        self.assertEqual(len(entities), limit)
        self.assertEqual(entities[0]["name"], "Robb")
        self.assertEqual(entities[1]["name"], "Bran")
        self.assertEqual(entities[2]["name"], "Catelyn")

        # Fetch next set of characters.
        new_iterator = page_query.fetch(limit=limit, offset=0, start_cursor=cursor)
        entities = list(new_iterator)
        self.assertEqual(len(entities), limit)
        self.assertEqual(entities[0]["name"], "Sansa")
        self.assertEqual(entities[1]["name"], "Jon Snow")
        self.assertEqual(entities[2]["name"], "Arya")

    def test_query_paginate_with_start_cursor(self):
        page_query = self._base_query()
        page_query.order = "appearances"
        limit = 3
        offset = 2
        iterator = page_query.fetch(limit=limit, offset=offset)

        # Fetch characters.
        page = six.next(iterator.pages)
        entities = list(page)
        cursor = iterator.next_page_token
        self.assertEqual(len(entities), limit)

        # Use cursor to create a fresh query.
        fresh_query = self._base_query()
        fresh_query.order = "appearances"

        new_entities = list(fresh_query.fetch(start_cursor=cursor, limit=limit))
        characters_remaining = len(self.CHARACTERS) - limit - offset
        self.assertEqual(len(new_entities), characters_remaining)
        self.assertEqual(new_entities[0]["name"], "Sansa")
        self.assertEqual(new_entities[2]["name"], "Arya")

    def test_query_distinct_on(self):
        query = self._base_query()
        query.distinct_on = ["alive"]

        expected_matches = 2
        # We expect 2, but allow the query to get 1 extra.
        entities = list(query.fetch(limit=expected_matches + 1))
        self.assertEqual(len(entities), expected_matches)

        self.assertEqual(entities[0]["name"], "Catelyn")
        self.assertEqual(entities[1]["name"], "Arya")


class TestDatastoreQueryOffsets(TestDatastore):
    TOTAL_OBJECTS = 2500
    NAMESPACE = "LargeCharacterEntity"
    KIND = "LargeCharacter"

    @classmethod
    def setUpClass(cls):
        cls.CLIENT = clone_client(Config.CLIENT)
        # Remove the namespace from the cloned client, since these
        # query tests rely on the entities to be already stored
        # cls.CLIENT.namespace = cls.NAMESPACE
        cls.CLIENT.namespace = None

        # Populating the datastore if necessary.
        populate_datastore.add_large_character_entities(client=cls.CLIENT)

    @classmethod
    def tearDownClass(cls):
        # In the emulator, destroy the query entities.
        if os.getenv(GCD_DATASET) is not None:
            # Use the client for this test instead of the global.
            clear_datastore.remove_all_entities(client=cls.CLIENT)

    def _base_query(self):
        # Use the client for this test instead of the global.
        return self.CLIENT.query(kind=self.KIND, namespace=self.NAMESPACE)

    def _verify(self, limit, offset, expected):
        # Query used for all tests
        page_query = self._base_query()
        page_query.add_filter("family", "=", "Stark")
        page_query.add_filter("alive", "=", False)

        iterator = page_query.fetch(limit=limit, offset=offset)
        entities = [e for e in iterator]
        if len(entities) != expected:
            self.fail(f"{limit}, {offset}, {expected}. Returned: {len(entities)}")

    def test_query_in_bounds_offsets(self):
        # Verify that with no offset there are the correct # of results
        self._verify(limit=None, offset=None, expected=self.TOTAL_OBJECTS)

        # Verify that with no limit there are results (offset provided)")
        self._verify(limit=None, offset=900, expected=self.TOTAL_OBJECTS - 900)

        # Offset beyond items larger Verify 200 items found")
        self._verify(limit=200, offset=1100, expected=200)

    def test_query_partially_out_of_bounds_offsets(self):
        # Offset within range, expect 50 despite larger limit")
        self._verify(limit=100, offset=self.TOTAL_OBJECTS - 50, expected=50)

    def test_query_out_of_bounds_offsets(self):
        # Offset beyond items larger Verify no items found")
        self._verify(limit=200, offset=self.TOTAL_OBJECTS + 1000, expected=0)


class TestDatastoreTransaction(TestDatastore):
    def test_transaction_via_with_statement(self):
        entity = datastore.Entity(key=Config.CLIENT.key("Company", "Google"))
        entity["url"] = u"www.google.com"

        with Config.CLIENT.transaction() as xact:
            result = Config.CLIENT.get(entity.key)
            if result is None:
                xact.put(entity)
                self.case_entities_to_delete.append(entity)

        # This will always return after the transaction.
        retrieved_entity = Config.CLIENT.get(entity.key)
        self.case_entities_to_delete.append(retrieved_entity)
        self.assertEqual(retrieved_entity, entity)

    def test_transaction_via_explicit_begin_get_commit(self):
        # See
        # github.com/GoogleCloudPlatform/google-cloud-python/issues/1859
        # Note that this example lacks the threading which provokes the race
        # condition in that issue:  we are basically just exercising the
        # "explict" path for using transactions.
        BEFORE_1 = 100
        BEFORE_2 = 0
        TRANSFER_AMOUNT = 40
        key1 = Config.CLIENT.key("account", "123")
        account1 = datastore.Entity(key=key1)
        account1["balance"] = BEFORE_1
        key2 = Config.CLIENT.key("account", "234")
        account2 = datastore.Entity(key=key2)
        account2["balance"] = BEFORE_2
        Config.CLIENT.put_multi([account1, account2])
        self.case_entities_to_delete.append(account1)
        self.case_entities_to_delete.append(account2)

        xact = Config.CLIENT.transaction()
        xact.begin()
        from_account = Config.CLIENT.get(key1, transaction=xact)
        to_account = Config.CLIENT.get(key2, transaction=xact)
        from_account["balance"] -= TRANSFER_AMOUNT
        to_account["balance"] += TRANSFER_AMOUNT

        xact.put(from_account)
        xact.put(to_account)
        xact.commit()

        after1 = Config.CLIENT.get(key1)
        after2 = Config.CLIENT.get(key2)
        self.assertEqual(after1["balance"], BEFORE_1 - TRANSFER_AMOUNT)
        self.assertEqual(after2["balance"], BEFORE_2 + TRANSFER_AMOUNT)

    def test_failure_with_contention(self):
        contention_prop_name = "baz"
        local_client = clone_client(Config.CLIENT)

        # Insert an entity which will be retrieved in a transaction
        # and updated outside it with a contentious value.
        key = local_client.key("BreakTxn", 1234)
        orig_entity = datastore.Entity(key=key)
        orig_entity["foo"] = u"bar"
        local_client.put(orig_entity)
        self.case_entities_to_delete.append(orig_entity)

        with self.assertRaises(Conflict):
            with local_client.transaction() as txn:
                entity_in_txn = local_client.get(key)

                # Update the original entity outside the transaction.
                orig_entity[contention_prop_name] = u"outside"
                Config.CLIENT.put(orig_entity)

                # Try to update the entity which we already updated outside the
                # transaction.
                entity_in_txn[contention_prop_name] = u"inside"
                txn.put(entity_in_txn)

    def test_empty_array_put(self):
        local_client = clone_client(Config.CLIENT)

        key = local_client.key("EmptyArray", 1234)
        local_client = datastore.Client()
        entity = datastore.Entity(key=key)
        entity["children"] = []
        local_client.put(entity)
        retrieved = local_client.get(entity.key)

        self.assertEqual(entity["children"], retrieved["children"])