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

#

"""

In-memory persistent stub for the Python datastore API. Gets, queries,

and searches are implemented as in-memory scans over all entities.

Stores entities across sessions as pickled proto bufs in a single file. On

startup, all entities are read from the file and loaded into memory. On

every Put(), the file is wiped and all entities are written from scratch.

Clients can also manually Read() and Write() the file themselves.

Transactions are serialized through __tx_lock. Each transaction acquires it

when it begins and releases it when it commits or rolls back. This is

important, since there are other member variables like __tx_snapshot that are

per-transaction, so they should only be used by one tx at a time.

"""

import datetime

import logging

import md5

import os

import struct

import sys

import tempfile

import threading

import warnings

import cPickle as pickle

from google.appengine.api import api_base_pb

from google.appengine.api import apiproxy_stub

from google.appengine.api import datastore

from google.appengine.api import datastore_admin

from google.appengine.api import datastore_errors

from google.appengine.api import datastore_types

from google.appengine.api import users

from google.appengine.datastore import datastore_pb

from google.appengine.datastore import datastore_index

from google.appengine.runtime import apiproxy_errors

from google.net.proto import ProtocolBuffer

from google.appengine.datastore import entity_pb

warnings.filterwarnings('ignore', 'tempnam is a potential security risk')

entity_pb.Reference.__hash__ = lambda self: hash(self.Encode())

datastore_pb.Query.__hash__ = lambda self: hash(self.Encode())

_MAXIMUM_RESULTS = 1000

_MAX_QUERY_OFFSET = 1000

_MAX_QUERY_COMPONENTS = 100

_BATCH_SIZE = 20

class _StoredEntity(object):

  """Simple wrapper around an entity stored by the stub.

  Public properties:

    protobuf: Native protobuf Python object, entity_pb.EntityProto.

    encoded_protobuf: Encoded binary representation of above protobuf.

    native: datastore.Entity instance.

  """

  def __init__(self, entity):

    """Create a _StoredEntity object and store an entity.

    Args:

      entity: entity_pb.EntityProto to store.

    """

    self.protobuf = entity

    self.encoded_protobuf = entity.Encode()

    self.native = datastore.Entity._FromPb(entity)

class _Cursor(object):

  """A query cursor.

  Public properties:

    cursor: the integer cursor

    count: the original total number of results

    keys_only: whether the query is keys_only

  Class attributes:

    _next_cursor: the next cursor to allocate

    _next_cursor_lock: protects _next_cursor

    _offset: the internal index for where we are in the results

  """

  _next_cursor = 1

  _next_cursor_lock = threading.Lock()

  def __init__(self, results, keys_only):

    """Constructor.

    Args:

      # the query results, in order, such that results[self.offset:] is

      # the next result

      results: list of entity_pb.EntityProto

      keys_only: integer

    """

    self.__results = results

    self.count = len(results)

    self.keys_only = keys_only

    self._offset = 0

    self._next_cursor_lock.acquire()

    try:

      self.cursor = _Cursor._next_cursor

      _Cursor._next_cursor += 1

    finally:

      self._next_cursor_lock.release()

  def PopulateQueryResult(self, result, count, offset=None, compiled=False):

    """Populates a QueryResult with this cursor and the given number of results.

    Args:

      result: datastore_pb.QueryResult

      count: integer

      offset: integer, overrides the internal offset

      compiled: boolean, whether we are compiling this query

    """

    _offset = offset

    if offset is None:

      _offset = self._offset

    if count > _MAXIMUM_RESULTS:

      count = _MAXIMUM_RESULTS

    result.mutable_cursor().set_cursor(self.cursor)

    result.set_keys_only(self.keys_only)

    self.count = len(self.__results[_offset:_offset + count])

    results_pbs = [r._ToPb() for r in self.__results[_offset:_offset + count]]

    result.result_list().extend(results_pbs)

    if offset is None:

      self._offset += self.count

    result.set_more_results(len(self.__results[_offset + self.count:]) > 0)

    if compiled and result.more_results():

      compiled_query = _FakeCompiledQuery(cursor=self.cursor,

                                          offset=_offset + self.count,

                                          keys_only=self.keys_only)

      result.mutable_compiled_query().CopyFrom(compiled_query._ToPb())

class _FakeCompiledQuery(object):

  def __init__(self, cursor, offset, keys_only=False):

    self.cursor = cursor

    self.offset = offset

    self.keys_only = keys_only

  def _ToPb(self):

    compiled_pb = datastore_pb.CompiledQuery()

    compiled_pb.mutable_primaryscan()

    compiled_pb.set_keys_only(self.keys_only)

    compiled_pb.set_limit(self.cursor)

    compiled_pb.set_offset(self.offset)

    return compiled_pb

class DatastoreFileStub(apiproxy_stub.APIProxyStub):

  """ Persistent stub for the Python datastore API.

  Stores all entities in memory, and persists them to a file as pickled

  protocol buffers. A DatastoreFileStub instance handles a single app's data

  and is backed by files on disk.

  """

  _PROPERTY_TYPE_TAGS = {

    datastore_types.Blob: entity_pb.PropertyValue.kstringValue,

    bool: entity_pb.PropertyValue.kbooleanValue,

    datastore_types.Category: entity_pb.PropertyValue.kstringValue,

    datetime.datetime: entity_pb.PropertyValue.kint64Value,

    datastore_types.Email: entity_pb.PropertyValue.kstringValue,

    float: entity_pb.PropertyValue.kdoubleValue,

    datastore_types.GeoPt: entity_pb.PropertyValue.kPointValueGroup,

    datastore_types.IM: entity_pb.PropertyValue.kstringValue,

    int: entity_pb.PropertyValue.kint64Value,

    datastore_types.Key: entity_pb.PropertyValue.kReferenceValueGroup,

    datastore_types.Link: entity_pb.PropertyValue.kstringValue,

    long: entity_pb.PropertyValue.kint64Value,

    datastore_types.PhoneNumber: entity_pb.PropertyValue.kstringValue,

    datastore_types.PostalAddress: entity_pb.PropertyValue.kstringValue,

    datastore_types.Rating: entity_pb.PropertyValue.kint64Value,

    str: entity_pb.PropertyValue.kstringValue,

    datastore_types.Text: entity_pb.PropertyValue.kstringValue,

    type(None): 0,

    unicode: entity_pb.PropertyValue.kstringValue,

    users.User: entity_pb.PropertyValue.kUserValueGroup,

    }

  WRITE_ONLY = entity_pb.CompositeIndex.WRITE_ONLY

  READ_WRITE = entity_pb.CompositeIndex.READ_WRITE

  DELETED = entity_pb.CompositeIndex.DELETED

  ERROR = entity_pb.CompositeIndex.ERROR

  _INDEX_STATE_TRANSITIONS = {

    WRITE_ONLY: frozenset((READ_WRITE, DELETED, ERROR)),

    READ_WRITE: frozenset((DELETED,)),

    ERROR: frozenset((DELETED,)),

    DELETED: frozenset((ERROR,)),

  }

  def __init__(self,

               app_id,

               datastore_file,

               history_file=None,

               require_indexes=False,

               service_name='datastore_v3',

               trusted=False):

    """Constructor.

    Initializes and loads the datastore from the backing files, if they exist.

    Args:

      app_id: string

      datastore_file: string, stores all entities across sessions.  Use None

          not to use a file.

      history_file: DEPRECATED. No-op.

      require_indexes: bool, default False.  If True, composite indexes must

          exist in index.yaml for queries that need them.

      service_name: Service name expected for all calls.

      trusted: bool, default False.  If True, this stub allows an app to

        access the data of another app.

    """

    super(DatastoreFileStub, self).__init__(service_name)

    assert isinstance(app_id, basestring) and app_id != ''

    self.__app_id = app_id

    self.__datastore_file = datastore_file

    self.SetTrusted(trusted)

    self.__entities = {}

    self.__schema_cache = {}

    self.__tx_snapshot = {}

    self.__queries = {}

    self.__transactions = {}

    self.__indexes = {}

    self.__require_indexes = require_indexes

    self.__query_history = {}

    self.__next_id = 1

    self.__next_tx_handle = 1

    self.__next_index_id = 1

    self.__id_lock = threading.Lock()

    self.__tx_handle_lock = threading.Lock()

    self.__index_id_lock = threading.Lock()

    self.__tx_lock = threading.Lock()

    self.__entities_lock = threading.Lock()

    self.__file_lock = threading.Lock()

    self.__indexes_lock = threading.Lock()

    self.Read()

  def Clear(self):

    """ Clears the datastore by deleting all currently stored entities and

    queries. """

    self.__entities = {}

    self.__queries = {}

    self.__transactions = {}

    self.__query_history = {}

    self.__schema_cache = {}

  def SetTrusted(self, trusted):

    """Set/clear the trusted bit in the stub.

    This bit indicates that the app calling the stub is trusted. A

    trusted app can write to datastores of other apps.

    Args:

      trusted: boolean.

    """

    self.__trusted = trusted

  def __ValidateAppId(self, app_id):

    """Verify that this is the stub for app_id.

    Args:

      app_id: An application ID.

    Raises:

      datastore_errors.BadRequestError: if this is not the stub for app_id.

    """

    if not self.__trusted and app_id != self.__app_id:

      raise datastore_errors.BadRequestError(

          'app %s cannot access app %s\'s data' % (self.__app_id, app_id))

  def __ValidateKey(self, key):

    """Validate this key.

    Args:

      key: entity_pb.Reference

    Raises:

      datastore_errors.BadRequestError: if the key is invalid

    """

    assert isinstance(key, entity_pb.Reference)

    self.__ValidateAppId(key.app())

    for elem in key.path().element_list():

      if elem.has_id() == elem.has_name():

        raise datastore_errors.BadRequestError(

          'each key path element should have id or name but not both: %r' % key)

  def _AppIdNamespaceKindForKey(self, key):

    """ Get (app, kind) tuple from given key.

    The (app, kind) tuple is used as an index into several internal

    dictionaries, e.g. __entities.

    Args:

      key: entity_pb.Reference

    Returns:

      Tuple (app, kind), both are unicode strings.

    """

    last_path = key.path().element_list()[-1]

    return key.app(), last_path.type()

  def _StoreEntity(self, entity):

    """ Store the given entity.

    Args:

      entity: entity_pb.EntityProto

    """

    key = entity.key()

    app_kind = self._AppIdNamespaceKindForKey(key)

    if app_kind not in self.__entities:

      self.__entities[app_kind] = {}

    self.__entities[app_kind][key] = _StoredEntity(entity)

    if app_kind in self.__schema_cache:

      del self.__schema_cache[app_kind]

  READ_PB_EXCEPTIONS = (ProtocolBuffer.ProtocolBufferDecodeError, LookupError,

                        TypeError, ValueError)

  READ_ERROR_MSG = ('Data in %s is corrupt or a different version. '

                    'Try running with the --clear_datastore flag.\n%r')

  READ_PY250_MSG = ('Are you using FloatProperty and/or GeoPtProperty? '

                    'Unfortunately loading float values from the datastore '

                    'file does not work with Python 2.5.0. '

                    'Please upgrade to a newer Python 2.5 release or use '

                    'the --clear_datastore flag.\n')

  def Read(self):

    """ Reads the datastore and history files into memory.

    The in-memory query history is cleared, but the datastore is *not*

    cleared; the entities in the files are merged into the entities in memory.

    If you want them to overwrite the in-memory datastore, call Clear() before

    calling Read().

    If the datastore file contains an entity with the same app name, kind, and

    key as an entity already in the datastore, the entity from the file

    overwrites the entity in the datastore.

    Also sets __next_id to one greater than the highest id allocated so far.

    """

    if self.__datastore_file and self.__datastore_file != '/dev/null':

      for encoded_entity in self.__ReadPickled(self.__datastore_file):

        try:

          entity = entity_pb.EntityProto(encoded_entity)

        except self.READ_PB_EXCEPTIONS, e:

          raise datastore_errors.InternalError(self.READ_ERROR_MSG %

                                               (self.__datastore_file, e))

        except struct.error, e:

          if (sys.version_info[0:3] == (2, 5, 0)

              and e.message.startswith('unpack requires a string argument')):

            raise datastore_errors.InternalError(self.READ_PY250_MSG +

                                                 self.READ_ERROR_MSG %

                                                 (self.__datastore_file, e))

          else:

            raise

        self._StoreEntity(entity)

        last_path = entity.key().path().element_list()[-1]

        if last_path.has_id() and last_path.id() >= self.__next_id:

          self.__next_id = last_path.id() + 1

  def Write(self):

    """ Writes out the datastore and history files. Be careful! If the files

    already exist, this method overwrites them!

    """

    self.__WriteDatastore()

  def __WriteDatastore(self):

    """ Writes out the datastore file. Be careful! If the file already exist,

    this method overwrites it!

    """

    if self.__datastore_file and self.__datastore_file != '/dev/null':

      encoded = []

      for kind_dict in self.__entities.values():

        for entity in kind_dict.values():

          encoded.append(entity.encoded_protobuf)

      self.__WritePickled(encoded, self.__datastore_file)

  def __ReadPickled(self, filename):

    """Reads a pickled object from the given file and returns it.

    """

    self.__file_lock.acquire()

    try:

      try:

        if filename and filename != '/dev/null' and os.path.isfile(filename):

          return pickle.load(open(filename, 'rb'))

        else:

          logging.warning('Could not read datastore data from %s', filename)

      except (AttributeError, LookupError, ImportError, NameError, TypeError,

              ValueError, struct.error, pickle.PickleError), e:

        raise datastore_errors.InternalError(

          'Could not read data from %s. Try running with the '

          '--clear_datastore flag. Cause:\n%r' % (filename, e))

    finally:

      self.__file_lock.release()

    return []

  def __WritePickled(self, obj, filename, openfile=file):

    """Pickles the object and writes it to the given file.

    """

    if not filename or filename == '/dev/null' or not obj:

      return

    tmpfile = openfile(os.tempnam(os.path.dirname(filename)), 'wb')

    pickler = pickle.Pickler(tmpfile, protocol=1)

    pickler.fast = True

    pickler.dump(obj)

    tmpfile.close()

    self.__file_lock.acquire()

    try:

      try:

        os.rename(tmpfile.name, filename)

      except OSError:

        try:

          os.remove(filename)

        except:

          pass

        os.rename(tmpfile.name, filename)

    finally:

      self.__file_lock.release()

  def MakeSyncCall(self, service, call, request, response):

    """ The main RPC entry point. service must be 'datastore_v3'.

    """

    self.assertPbIsInitialized(request)

    super(DatastoreFileStub, self).MakeSyncCall(service,

                                                call,

                                                request,

                                                response)

    self.assertPbIsInitialized(response)

  def assertPbIsInitialized(self, pb):

    """Raises an exception if the given PB is not initialized and valid."""

    explanation = []

    assert pb.IsInitialized(explanation), explanation

    pb.Encode()

  def QueryHistory(self):

    """Returns a dict that maps Query PBs to times they've been run.

    """

    return dict((pb, times) for pb, times in self.__query_history.items()

                if pb.app() == self.__app_id)

  def _Dynamic_Put(self, put_request, put_response):