#!/usr/bin/env python

#

# Copyright 2007 Google Inc.

#

# 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.

#

"""The Python datastore API used by app developers.

Defines Entity, Query, and Iterator classes, as well as methods for all of the

datastore's calls. Also defines conversions between the Python classes and

their PB counterparts.

The datastore errors are defined in the datastore_errors module. That module is

only required to avoid circular imports. datastore imports datastore_types,

which needs BadValueError, so it can't be defined in datastore.

"""

import heapq

import itertools

import logging

import re

import string

import sys

import traceback

from xml.sax import saxutils

from google.appengine.api import api_base_pb

from google.appengine.api import apiproxy_stub_map

from google.appengine.api import capabilities

from google.appengine.api import datastore_errors

from google.appengine.api import datastore_types

from google.appengine.datastore import datastore_index

from google.appengine.datastore import datastore_pb

from google.appengine.runtime import apiproxy_errors

from google.appengine.datastore import entity_pb

try:

  __import__('google.appengine.api.labs.taskqueue.taskqueue_service_pb')

  taskqueue_service_pb = sys.modules.get(

      'google.appengine.api.labs.taskqueue.taskqueue_service_pb')

except ImportError:

  from google.appengine.api.taskqueue import taskqueue_service_pb

MAX_ALLOWABLE_QUERIES = 30

MAXIMUM_RESULTS = 1000

DEFAULT_TRANSACTION_RETRIES = 3

READ_CAPABILITY = capabilities.CapabilitySet('datastore_v3')

WRITE_CAPABILITY = capabilities.CapabilitySet(

    'datastore_v3',

    capabilities=['write'])

_MAX_INDEXED_PROPERTIES = 5000

_MAX_ID_BATCH_SIZE = 1000 * 1000 * 1000

Key = datastore_types.Key

typename = datastore_types.typename

_txes = {}

def NormalizeAndTypeCheck(arg, types):

  """Normalizes and type checks the given argument.

  Args:

    arg: an instance, tuple, list, iterator, or generator of the given type(s)

    types: allowed type or tuple of types

  Returns:

    A (list, bool) tuple. The list is a normalized, shallow copy of the

    argument. The boolean is True if the argument was a sequence, False

    if it was a single object.

  Raises:

    AssertionError: types includes list or tuple.

    BadArgumentError: arg is not an instance or sequence of one of the given

    types.

  """

  if not isinstance(types, (list, tuple)):

    types = (types,)

  assert list not in types and tuple not in types

  if isinstance(arg, types):

    return ([arg], False)

  else:

    try:

      for val in arg:

        if not isinstance(val, types):

          raise datastore_errors.BadArgumentError(

              'Expected one of %s; received %s (a %s).' %

              (types, val, typename(val)))

    except TypeError:

      raise datastore_errors.BadArgumentError(

          'Expected an instance or sequence of %s; received %s (a %s).' %

          (types, arg, typename(arg)))

    return (list(arg), True)

def NormalizeAndTypeCheckKeys(keys):

  """Normalizes and type checks that the given argument is a valid key or keys.

  A wrapper around NormalizeAndTypeCheck() that accepts strings, Keys, and

  Entities, and normalizes to Keys.

  Args:

    keys: a Key or sequence of Keys

  Returns:

    A (list of Keys, bool) tuple. See NormalizeAndTypeCheck.

  Raises:

    BadArgumentError: arg is not an instance or sequence of one of the given

    types.

  """

  keys, multiple = NormalizeAndTypeCheck(keys, (basestring, Entity, Key))

  keys = [_GetCompleteKeyOrError(key) for key in keys]

  return (keys, multiple)

def Put(entities):

  """Store one or more entities in the datastore.

  The entities may be new or previously existing. For new entities, Put() will

  fill in the app id and key assigned by the datastore.

  If the argument is a single Entity, a single Key will be returned. If the

  argument is a list of Entity, a list of Keys will be returned.

  Args:

    entities: Entity or list of Entities

  Returns:

    Key or list of Keys

  Raises:

    TransactionFailedError, if the Put could not be committed.

  """

  entities, multiple = NormalizeAndTypeCheck(entities, Entity)

  if multiple and not entities:

    return []

  for entity in entities:

    if not entity.kind() or not entity.app_id_namespace():

      raise datastore_errors.BadRequestError(

          'App and kind must not be empty, in entity: %s' % entity)

  req = datastore_pb.PutRequest()

  req.entity_list().extend([e._ToPb() for e in entities])

  keys = [e.key() for e in entities]

  tx = _MaybeSetupTransaction(req, keys)

  resp = datastore_pb.PutResponse()

  try:

    apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Put', req, resp)

  except apiproxy_errors.ApplicationError, err:

    raise _ToDatastoreError(err)

  keys = resp.key_list()

  num_keys = len(keys)

  num_entities = len(entities)

  if num_keys != num_entities:

    raise datastore_errors.InternalError(

        'Put accepted %d entities but returned %d keys.' %

        (num_entities, num_keys))

  for entity, key in zip(entities, keys):

    entity._Entity__key._Key__reference.CopyFrom(key)

  if tx:

    tx.entity_group = entities[0].entity_group()

  if multiple:

    return [Key._FromPb(k) for k in keys]

  else:

    return Key._FromPb(resp.key(0))

def Get(keys):

  """Retrieves one or more entities from the datastore.

  Retrieves the entity or entities with the given key(s) from the datastore

  and returns them as fully populated Entity objects, as defined below. If

  there is an error, raises a subclass of datastore_errors.Error.

  If keys is a single key or string, an Entity will be returned, or

  EntityNotFoundError will be raised if no existing entity matches the key.

  However, if keys is a list or tuple, a list of entities will be returned

  that corresponds to the sequence of keys. It will include entities for keys

  that were found and None placeholders for keys that were not found.

  Args:

    # the primary key(s) of the entity(ies) to retrieve

    keys: Key or string or list of Keys or strings

  Returns:

    Entity or list of Entity objects

  """

  keys, multiple = NormalizeAndTypeCheckKeys(keys)

  if multiple and not keys:

    return []

  req = datastore_pb.GetRequest()

  req.key_list().extend([key._Key__reference for key in keys])

  _MaybeSetupTransaction(req, keys)

  resp = datastore_pb.GetResponse()

  try:

    apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Get', req, resp)

  except apiproxy_errors.ApplicationError, err:

    raise _ToDatastoreError(err)

  entities = []

  for group in resp.entity_list():

    if group.has_entity():

      entities.append(Entity._FromPb(group.entity()))

    else:

      entities.append(None)

  if multiple:

    return entities

  else:

    if entities[0] is None:

      raise datastore_errors.EntityNotFoundError()

    return entities[0]

def Delete(keys):

  """Deletes one or more entities from the datastore. Use with care!

  Deletes the given entity(ies) from the datastore. You can only delete

  entities from your app. If there is an error, raises a subclass of

  datastore_errors.Error.

  Args:

    # the primary key(s) of the entity(ies) to delete

    keys: Key or string or list of Keys or strings

  Raises:

    TransactionFailedError, if the Delete could not be committed.

  """

  keys, multiple = NormalizeAndTypeCheckKeys(keys)

  if multiple and not keys:

    return

  req = datastore_pb.DeleteRequest()

  req.key_list().extend([key._Key__reference for key in keys])

  tx = _MaybeSetupTransaction(req, keys)

  resp = datastore_pb.DeleteResponse()

  try:

    apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Delete', req, resp)

  except apiproxy_errors.ApplicationError, err:

    raise _ToDatastoreError(err)

class Entity(dict):

  """A datastore entity.

  Includes read-only accessors for app id, kind, and primary key. Also

  provides dictionary-style access to properties.

  """

  def __init__(self, kind, parent=None, _app=None, name=None, id=None,

               unindexed_properties=[], _namespace=None):

    """Constructor. Takes the kind and transaction root, which cannot be

    changed after the entity is constructed, and an optional parent. Raises

    BadArgumentError or BadKeyError if kind is invalid or parent is not an

    existing Entity or Key in the datastore.

    Args:

      # this entity's kind

      kind: string

      # if provided, this entity's parent. Its key must be complete.

      parent: Entity or Key

      # if provided, this entity's name.

      name: string

      # if provided, this entity's id.

      id: integer

      # if provided, a sequence of property names that should not be indexed

      # by the built-in single property indices.

      unindexed_properties: list or tuple of strings

    """

    ref = entity_pb.Reference()

    _app_namespace = datastore_types.ResolveAppIdNamespace(_app, _namespace)

    ref.set_app(_app_namespace.to_encoded())

    datastore_types.ValidateString(kind, 'kind',

                                   datastore_errors.BadArgumentError)

    if parent is not None:

      parent = _GetCompleteKeyOrError(parent)

      if _app_namespace != parent.app_id_namespace():

        raise datastore_errors.BadArgumentError(

            " %s doesn't match parent's app_namespace %s" %

            (_app_namespace, parent.app_id_namespace()))

      ref.CopyFrom(parent._Key__reference)

    last_path = ref.mutable_path().add_element()

    last_path.set_type(kind.encode('utf-8'))

    if name is not None and id is not None:

      raise datastore_errors.BadArgumentError(

          "Cannot set both name and id on an Entity")

    if name is not None:

      datastore_types.ValidateString(name, 'name')

      last_path.set_name(name.encode('utf-8'))

    if id is not None:

      datastore_types.ValidateInteger(id, 'id')

      last_path.set_id(id)

    unindexed_properties, multiple = NormalizeAndTypeCheck(unindexed_properties, basestring)

    if not multiple:

      raise datastore_errors.BadArgumentError(

        'unindexed_properties must be a sequence; received %s (a %s).' %

        (unindexed_properties, typename(unindexed_properties)))

    for prop in unindexed_properties:

      datastore_types.ValidateProperty(prop, None)

    self.__unindexed_properties = frozenset(unindexed_properties)

    self.__key = Key._FromPb(ref)

  def app(self):

    """Returns the name of the application that created this entity, a

    string or None if not set.

    """

    return self.__key.app()

  def namespace(self):

    """Returns the namespace of this entity, a string or None.

    """

    return self.__key.namespace()

  def app_id_namespace(self):

    """Returns the AppIdNamespace of this entity or None if not set.

    """

    return self.__key.app_id_namespace()

  def kind(self):

    """Returns this entity's kind, a string.

    """

    return self.__key.kind()

  def is_saved(self):

    """Returns if this entity has been saved to the datastore

    """

    last_path = self.__key._Key__reference.path().element_list()[-1]

    return ((last_path.has_name() ^ last_path.has_id()) and

            self.__key.has_id_or_name())

  def key(self):

    """Returns this entity's primary key, a Key instance.

    """

    return self.__key

  def parent(self):

    """Returns this entity's parent, as a Key. If this entity has no parent,

    returns None.

    """

    return self.key().parent()

  def entity_group(self):

    """Returns this entity's entity group as a Key.

    Note that the returned Key will be incomplete if this is a a root entity

    and its key is incomplete.

    """

    return self.key().entity_group()

  def unindexed_properties(self):

    """Returns this entity's unindexed properties, as a frozenset of strings."""

    return getattr(self, '_Entity__unindexed_properties', [])

  def __setitem__(self, name, value):

    """Implements the [] operator. Used to set property value(s).

    If the property name is the empty string or not a string, raises

    BadPropertyError. If the value is not a supported type, raises

    BadValueError.

    """

    datastore_types.ValidateProperty(name, value)

    dict.__setitem__(self, name, value)

  def setdefault(self, name, value):

    """If the property exists, returns its value. Otherwise sets it to value.

    If the property name is the empty string or not a string, raises

    BadPropertyError. If the value is not a supported type, raises

    BadValueError.

    """

    datastore_types.ValidateProperty(name, value)

    return dict.setdefault(self, name, value)

  def update(self, other):

    """Updates this entity's properties from the values in other.

    If any property name is the empty string or not a string, raises

    BadPropertyError. If any value is not a supported type, raises

    BadValueError.

    """

    for name, value in other.items():

      self.__setitem__(name, value)

  def copy(self):

    """The copy method is not supported.

    """

    raise NotImplementedError('Entity does not support the copy() method.')

  def ToXml(self):

    """Returns an XML representation of this entity. Atom and gd:namespace

    properties are converted to XML according to their respective schemas. For

    more information, see:

      http://www.atomenabled.org/developers/syndication/

      http://code.google.com/apis/gdata/common-elements.html

    This is *not* optimized. It shouldn't be used anywhere near code that's

    performance-critical.

    """

    xml = u'<entity kind=%s' % saxutils.quoteattr(self.kind())

    if self.__key.has_id_or_name():

      xml += ' key=%s' % saxutils.quoteattr(str(self.__key))

    xml += '>'

    if self.__key.has_id_or_name():

      xml += '\n  <key>%s</key>' % self.__key.ToTagUri()

    properties = self.keys()

    if properties:

      properties.sort()

      xml += '\n  ' + '\n  '.join(self._PropertiesToXml(properties))

    xml += '\n</entity>\n'

    return xml

  def _PropertiesToXml(self, properties):

    """ Returns a list of the XML representations of each of the given

    properties. Ignores properties that don't exist in this entity.

    Arg:

      properties: string or list of strings

    Returns:

      list of strings

    """

    xml_properties = []

    for propname in properties:

      if not self.has_key(propname):

        continue

      propname_xml = saxutils.quoteattr(propname)

      values = self[propname]

      if not isinstance(values, list):

        values = [values]

      proptype = datastore_types.PropertyTypeName(values[0])

      proptype_xml = saxutils.quoteattr(proptype)

      escaped_values = self._XmlEscapeValues(propname)

      open_tag = u'<property name=%s type=%s>' % (propname_xml, proptype_xml)

      close_tag = u'</property>'

      xml_properties += [open_tag + val + close_tag for val in escaped_values]

    return xml_properties

  def _XmlEscapeValues(self, property):

    """ Returns a list of the XML-escaped string values for the given property.

    Raises an AssertionError if the property doesn't exist.

    Arg:

      property: string

    Returns:

      list of strings

    """

    assert self.has_key(property)

    xml = []

    values = self[property]

    if not isinstance(values, list):

      values = [values]

    for val in values:

      if hasattr(val, 'ToXml'):

        xml.append(val.ToXml())

      else:

        if val is None:

          xml.append('')

        else: