java/src/main/com/google/appengine/api/labs/datastore/overlay/NamespacePinnedAsyncDatastoreServiceImpl.java

   1 package com.google.appengine.api.labs.datastore.overlay;
   2
   3 import static com.google.common.base.Preconditions.checkNotNull;
   4
   5 import com.google.appengine.api.NamespaceManager;
   6 import com.google.appengine.api.datastore.AsyncDatastoreService;
   7 import com.google.appengine.api.datastore.BaseDatastoreService;
   8 import com.google.appengine.api.datastore.DatastoreAttributes;
   9 import com.google.appengine.api.datastore.DatastoreFailureException;
  10 import com.google.appengine.api.datastore.Entity;
  11 import com.google.appengine.api.datastore.Index;
  12 import com.google.appengine.api.datastore.Key;
  13 import com.google.appengine.api.datastore.KeyRange;
  14 import com.google.appengine.api.datastore.Transaction;
  15 import com.google.appengine.api.datastore.TransactionOptions;
  16 import com.google.apphosting.api.ApiProxy;
  17 import com.google.apphosting.datastore.DatastoreV4;
  18 import com.google.common.collect.ImmutableList;
  19 import com.google.protobuf.InvalidProtocolBufferException;
  20
  21 import java.util.List;
  22 import java.util.Map;
  23 import java.util.concurrent.Future;
  24
  25 /**
  26  * A thin wrapper around {@link AsyncDatastoreService}, where all operations are redirected to use
  27  * an alternate namespace.
  28  *
  29  * <p>For operations which use the ambient namespace from the {@link NamespaceManager} (such as
  30  * {@code allocateIds}), {@code namespacePrefix} is applied to the current namespace
  31  * before completing the operation.
  32  *
  33  * <p>For operations whose arguments contain namespaces in some form (such as keys),
  34  * {@code namespacePrefix} is applied to the relevant namespaces before completing the
  35  * operation.
  36  *
  37  * <p>In all cases, there are no lasting effects on the ambient namespace; it is restored to the
  38  * original value after the operation is done.
  39  */
  40 final class NamespacePinnedAsyncDatastoreServiceImpl extends NamespacePinnedBaseDatastoreServiceImpl
  41     implements AsyncDatastoreService {
  42   private final AsyncDatastoreService datastore;
  43
  44   /**
  45    * Constructs a new {@link NamespacePinnedAsyncDatastoreServiceImpl}.
  46    *
  47    * @param datastore the underlying {@link AsyncDatastoreService}
  48    * @param namespacePrefix the prefix to apply to all namespaces
  49    * @param shouldReserveIds a flag that indicates whether calling {@code put()} should explicitly
  50    *        reserve (not allocate) the numeric IDs used in any complete keys before storing them
  51    */
  52   NamespacePinnedAsyncDatastoreServiceImpl(AsyncDatastoreService datastore, String namespacePrefix,
  53       boolean shouldReserveIds) {
  54     super(checkNotNull(namespacePrefix), shouldReserveIds);
  55     this.datastore = checkNotNull(datastore);
  56   }
  57
  58   @Override
  59   public Future<Entity> get(Key key) {
  60     checkNotNull(key);
  61     return getImpl(datastore, key);
  62   }
  63
  64   @Override
  65   public Future<Entity> get( Transaction txn, Key key) {
  66     checkNotNull(key);
  67     return getImpl(getTxnAsyncDatastore(txn), key);
  68   }
  69
  70   private Future<Entity> getImpl(AsyncDatastoreService datastore, Key key) {
  71     checkNotNull(datastore);
  72     checkNotNull(key);
  73     return new RethrowingFutureWrapper<Entity, Entity>(
  74         datastore.get(getAlternateNamespaceKey(key))) {
  75       @Override
  76       protected Entity wrap(Entity entity) throws Exception {
  77         return getOriginalNamespaceEntity(entity);
  78       }
  79     };
  80   }
  81
  82   @Override
  83   public Future<Map<Key, Entity>> get(Iterable<Key> keys) {
  84     checkNotNull(keys);
  85     return getImpl(datastore, keys);
  86   }
  87
  88   @Override
  89   public Future<Map<Key, Entity>> get( Transaction txn, Iterable<Key> keys) {
  90     checkNotNull(keys);
  91     return getImpl(getTxnAsyncDatastore(txn), keys);
  92   }
  93
  94   private Future<Map<Key, Entity>> getImpl(AsyncDatastoreService datastore, Iterable<Key> keys) {
  95     checkNotNull(datastore);
  96     checkNotNull(keys);
  97     return new RethrowingFutureWrapper<Map<Key, Entity>, Map<Key, Entity>>(
  98         datastore.get(getAlternateNamespaceKeys(keys))) {
  99       @Override
 100       protected Map<Key, Entity> wrap(Map<Key, Entity> entityMap) throws Exception {
 101         return getOriginalNamespaceEntityMap(entityMap);
 102       }
 103     };
 104   }
 105
 106   @Override
 107   public Future<Key> put(Entity entity) {
 108     checkNotNull(entity);
 109     return putImpl(datastore, entity);
 110   }
 111
 112   @Override
 113   public Future<Key> put( Transaction txn, Entity entity) {
 114     checkNotNull(entity);
 115     return putImpl(getTxnAsyncDatastore(txn), entity);
 116   }
 117
 118   private Future<Key> putImpl(final AsyncDatastoreService datastore, Entity entity) {
 119     checkNotNull(datastore);
 120     checkNotNull(entity);
 121     final Entity alternateEntity = getAlternateNamespaceEntity(entity);
 122     if (shouldReserveIds) {
 123       return new RethrowingFutureWrapper<Void, Key>(reserveId(alternateEntity)) {
 124         @Override
 125         protected Key wrap(Void v) throws Exception {
 126           Key key = datastore.put(alternateEntity).get();
 127           return getOriginalNamespaceKey(key);
 128         }
 129       };
 130     } else {
 131       return new RethrowingFutureWrapper<Key, Key>(datastore.put(alternateEntity)) {
 132         @Override
 133         protected Key wrap(Key key) throws Exception {
 134           return getOriginalNamespaceKey(key);
 135         }
 136       };
 137     }
 138   }
 139
 140   @Override
 141   public Future<List<Key>> put(Iterable<Entity> entities) {
 142     checkNotNull(entities);
 143     return putImpl(datastore, entities);
 144   }
 145
 146   @Override
 147   public Future<List<Key>> put( Transaction txn, Iterable<Entity> entities) {
 148     checkNotNull(entities);
 149     return putImpl(getTxnAsyncDatastore(txn), entities);
 150   }
 151
 152   private Future<List<Key>> putImpl(final AsyncDatastoreService datastore,
 153       Iterable<Entity> entities) {
 154     checkNotNull(datastore);
 155     checkNotNull(entities);
 156     final List<Entity> alternateEntities = getAlternateNamespaceEntities(entities);
 157     if (shouldReserveIds) {
 158       return new RethrowingFutureWrapper<Void, List<Key>>(reserveIds(alternateEntities)) {
 159         @Override
 160         protected List<Key> wrap(Void v) throws Exception {
 161           List<Key> keys = datastore.put(alternateEntities).get();
 162           return getOriginalNamespaceKeys(keys);
 163         }
 164       };
 165     } else {
 166       return new RethrowingFutureWrapper<List<Key>, List<Key>>(datastore.put(alternateEntities)) {
 167         @Override
 168         protected List<Key> wrap(List<Key> keys) throws Exception {
 169           return getOriginalNamespaceKeys(keys);
 170         }
 171       };
 172     }
 173   }
 174
 175   @Override
 176   public Future<Void> delete(Key... keys) {
 177     checkNotNull(keys);
 178     return delete(ImmutableList.copyOf(keys));
 179   }
 180
 181   @Override
 182   public Future<Void> delete( Transaction txn, Key... keys) {
 183     checkNotNull(keys);
 184     return delete(txn, ImmutableList.copyOf(keys));
 185   }
 186
 187   @Override
 188   public Future<Void> delete(Iterable<Key> keys) {
 189     checkNotNull(keys);
 190     return datastore.delete(getAlternateNamespaceKeys(keys));
 191   }
 192
 193   @Override
 194   public Future<Void> delete( Transaction txn, Iterable<Key> keys) {
 195     checkNotNull(keys);
 196     return datastore.delete(txn, getAlternateNamespaceKeys(keys));
 197   }
 198
 199   @Override
 200   public Future<Transaction> beginTransaction() {
 201     return datastore.beginTransaction();
 202   }
 203
 204   @Override
 205   public Future<Transaction> beginTransaction(TransactionOptions options) {
 206     checkNotNull(options);
 207     return datastore.beginTransaction(options);
 208   }
 209
 210   @Override
 211   public Future<KeyRange> allocateIds(String kind, long num) {
 212     checkNotNull(kind);
 213     String currentNamespace = NamespaceManager.get();
 214     try {
 215       NamespaceManager.set(getAlternateNamespace(currentNamespace));
 216       return new RethrowingFutureWrapper<KeyRange, KeyRange>(datastore.allocateIds(kind, num)) {
 217         @Override
 218         protected KeyRange wrap(KeyRange range) throws Exception {
 219           return getOriginalNamespaceKeyRange(range);
 220         }
 221       };
 222     } finally {
 223       NamespaceManager.set(currentNamespace);
 224     }
 225   }
 226
 227   @Override
 228   public Future<KeyRange> allocateIds(Key parent, String kind, long num) {
 229     checkNotNull(parent);
 230     checkNotNull(kind);
 231     String currentNamespace = NamespaceManager.get();
 232     try {
 233       NamespaceManager.set(getAlternateNamespace(parent.getNamespace()));
 234       return new RethrowingFutureWrapper<KeyRange, KeyRange>(
 235           datastore.allocateIds(getAlternateNamespaceKey(parent), kind, num)) {
 236         @Override
 237         protected KeyRange wrap(KeyRange range) throws Exception {
 238           return getOriginalNamespaceKeyRange(range);
 239         }
 240       };
 241     } finally {
 242       NamespaceManager.set(currentNamespace);
 243     }
 244   }
 245
 246   @Override
 247   public Future<DatastoreAttributes> getDatastoreAttributes() {
 248     return datastore.getDatastoreAttributes();
 249   }
 250
 251   @Override
 252   public Future<Map<Index, Index.IndexState>> getIndexes() {
 253     return datastore.getIndexes();
 254   }
 255
 256   @Override
 257   protected BaseDatastoreService getUnderlyingBaseDatastoreService() {
 258     return datastore;
 259   }
 260
 261   /**
 262    * Issues an RPC to reserve the key of {@code entity}. This is used when encountering
 263    * complete keys that have possibly been assigned IDs by the parent Datastore, but that the
 264    * overlay may not have seen yet.
 265    */
 266   private static Future<Void> reserveId(Entity entity) {
 267     checkNotNull(entity);
 268     return reserveIds(ImmutableList.of(entity));
 269   }
 270
 271   /**
 272    * Issues an RPC to reserve the keys of {@code entities}. This is used when encountering
 273    * complete keys that have possibly been assigned IDs by the parent Datastore, but that the
 274    * overlay may not have seen yet.
 275    */
 276   private static Future<Void> reserveIds(Iterable<Entity> entities) {
 277     checkNotNull(entities);
 278     DatastoreV4.AllocateIdsRequest req = makeReserveIdsRequest(entities);
 279     return new RethrowingFutureWrapper<byte[], Void>(
 280         ApiProxy.makeAsyncCall("datastore_v4", "AllocateIds", req.toByteArray())) {
 281       @Override
 282       protected Void wrap(byte[] respBytes) throws Exception {
 283         try {
 284           DatastoreV4.AllocateIdsResponse.parseFrom(respBytes);
 285         } catch (InvalidProtocolBufferException e) {
 286           throw new DatastoreFailureException("error reserving IDs", e);
 287         }
 288         return null;
 289       }
 290     };
 291   }
 292
 293   @Override
 294   protected BaseDatastoreService getTxnBaseDatastore( Transaction txn) {
 295     return getTxnAsyncDatastore(txn);
 296   }
 297
 298   /**
 299    * Gets a version of the Datastore that is linked to {@code txn}.
 300    */
 301   private AsyncDatastoreService getTxnAsyncDatastore( Transaction txn) {
 302     return new TransactionLinkedAsyncDatastoreServiceImpl(datastore, txn);
 303   }
 304 }