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.checkArgument;
   4 import static com.google.common.base.Preconditions.checkNotNull;
   5
   6 import com.google.appengine.api.NamespaceManager;
   7 import com.google.appengine.api.datastore.AsyncDatastoreService;
   8 import com.google.appengine.api.datastore.BaseDatastoreService;
   9 import com.google.appengine.api.datastore.DatastoreAttributes;
  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.PreparedQuery;
  15 import com.google.appengine.api.datastore.Projection;
  16 import com.google.appengine.api.datastore.Query;
  17 import com.google.appengine.api.datastore.Query.FilterPredicate;
  18 import com.google.appengine.api.datastore.Query.SortPredicate;
  19 import com.google.appengine.api.datastore.Transaction;
  20 import com.google.appengine.api.datastore.TransactionOptions;
  21 import com.google.common.annotations.VisibleForTesting;
  22 import com.google.common.base.Strings;
  23 import com.google.common.collect.ImmutableList;
  24 import com.google.common.collect.Lists;
  25 import com.google.common.collect.Maps;
  26
  27 import java.util.Collection;
  28 import java.util.List;
  29 import java.util.Map;
  30 import java.util.concurrent.Future;
  31
  32 /**
  33  * A thin wrapper around {@link AsyncDatastoreService}, where all operations are redirected to use
  34  * an alternate namespace.
  35  *
  36  * <p>For operations which use the ambient namespace from the {@link NamespaceManager} (such as
  37  * {@code allocateIds}), {@code namespacePrefix} is applied to the current namespace
  38  * before completing the operation.
  39  *
  40  * <p>For operations whose arguments contain namespaces in some form (such as keys),
  41  * {@code namespacePrefix} is applied to the relevant namespaces before completing the
  42  * operation.
  43  *
  44  * <p>In all cases, there are no lasting effects on the ambient namespace; it is restored to the
  45  * original value after the operation is done.
  46  */
  47 final class NamespacePinnedAsyncDatastoreServiceImpl implements AsyncDatastoreService {
  48   private final AsyncDatastoreService datastore;
  49   private final String namespacePrefix;
  50   private final boolean shouldReserveIds;
  51
  52   /**
  53    * Constructs a new {@link NamespacePinnedAsyncDatastoreServiceImpl}.
  54    *
  55    * @param datastore the underlying {@link AsyncDatastoreService}
  56    * @param namespacePrefix the prefix to apply to all namespaces
  57    * @param shouldReserveIds a flag that indicates whether calling {@code put()} should explicitly
  58    *        reserve (not allocate) the numeric IDs used in any complete keys before storing them
  59    */
  60   NamespacePinnedAsyncDatastoreServiceImpl(AsyncDatastoreService datastore, String namespacePrefix,
  61       boolean shouldReserveIds) {
  62     this.datastore = checkNotNull(datastore);
  63     this.namespacePrefix = checkNotNull(namespacePrefix);
  64     this.shouldReserveIds = shouldReserveIds;
  65   }
  66
  67   @Override
  68   public Future<Entity> get(Key key) {
  69     checkNotNull(key);
  70     return getImpl(datastore, key);
  71   }
  72
  73   @Override
  74   public Future<Entity> get( Transaction txn, Key key) {
  75     checkNotNull(key);
  76     return getImpl(getTxnAsyncDatastore(txn), key);
  77   }
  78
  79   private Future<Entity> getImpl(AsyncDatastoreService datastore, Key key) {
  80     checkNotNull(datastore);
  81     checkNotNull(key);
  82     return new RethrowingFutureWrapper<Entity, Entity>(
  83         datastore.get(getAlternateNamespaceKey(key))) {
  84       @Override
  85       protected Entity wrap(Entity entity) throws Exception {
  86         return getOriginalNamespaceEntity(entity);
  87       }
  88     };
  89   }
  90
  91   @Override
  92   public Future<Map<Key, Entity>> get(Iterable<Key> keys) {
  93     checkNotNull(keys);
  94     return getImpl(datastore, keys);
  95   }
  96
  97   @Override
  98   public Future<Map<Key, Entity>> get( Transaction txn, Iterable<Key> keys) {
  99     checkNotNull(keys);
 100     return getImpl(getTxnAsyncDatastore(txn), keys);
 101   }
 102
 103   private Future<Map<Key, Entity>> getImpl(AsyncDatastoreService datastore, Iterable<Key> keys) {
 104     checkNotNull(datastore);
 105     checkNotNull(keys);
 106     return new RethrowingFutureWrapper<Map<Key, Entity>, Map<Key, Entity>>(
 107         datastore.get(getAlternateNamespaceKeys(keys))) {
 108       @Override
 109       protected Map<Key, Entity> wrap(Map<Key, Entity> entityMap) throws Exception {
 110         return getOriginalNamespaceEntityMap(entityMap);
 111       }
 112     };
 113   }
 114
 115   @Override
 116   public Future<Key> put(Entity entity) {
 117     checkNotNull(entity);
 118     return putImpl(datastore, entity);
 119   }
 120
 121   @Override
 122   public Future<Key> put( Transaction txn, Entity entity) {
 123     checkNotNull(entity);
 124     return putImpl(getTxnAsyncDatastore(txn), entity);
 125   }
 126
 127   private Future<Key> putImpl(final AsyncDatastoreService datastore, Entity entity) {
 128     checkNotNull(datastore);
 129     checkNotNull(entity);
 130     final Entity alternateEntity = getAlternateNamespaceEntity(entity);
 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   @Override
 140   public Future<List<Key>> put(Iterable<Entity> entities) {
 141     checkNotNull(entities);
 142     return putImpl(datastore, entities);
 143   }
 144
 145   @Override
 146   public Future<List<Key>> put( Transaction txn, Iterable<Entity> entities) {
 147     checkNotNull(entities);
 148     return putImpl(getTxnAsyncDatastore(txn), entities);
 149   }
 150
 151   private Future<List<Key>> putImpl(final AsyncDatastoreService datastore,
 152       Iterable<Entity> entities) {
 153     checkNotNull(datastore);
 154     checkNotNull(entities);
 155     final List<Entity> alternateEntities = getAlternateNamespaceEntities(entities);
 156     return new RethrowingFutureWrapper<List<Key>, List<Key>>(datastore.put(alternateEntities)) {
 157       @Override
 158       protected List<Key> wrap(List<Key> keys) throws Exception {
 159         return getOriginalNamespaceKeys(keys);
 160       }
 161     };
 162   }
 163
 164   @Override
 165   public Future<Void> delete(Key... keys) {
 166     checkNotNull(keys);
 167     return delete(ImmutableList.copyOf(keys));
 168   }
 169
 170   @Override
 171   public Future<Void> delete( Transaction txn, Key... keys) {
 172     checkNotNull(keys);
 173     return delete(txn, ImmutableList.copyOf(keys));
 174   }
 175
 176   @Override
 177   public Future<Void> delete(Iterable<Key> keys) {
 178     checkNotNull(keys);
 179     return datastore.delete(getAlternateNamespaceKeys(keys));
 180   }
 181
 182   @Override
 183   public Future<Void> delete( Transaction txn, Iterable<Key> keys) {
 184     checkNotNull(keys);
 185     return datastore.delete(txn, getAlternateNamespaceKeys(keys));
 186   }
 187
 188   @Override
 189   public Future<Transaction> beginTransaction() {
 190     return datastore.beginTransaction();
 191   }
 192
 193   @Override
 194   public Future<Transaction> beginTransaction(TransactionOptions options) {
 195     checkNotNull(options);
 196     return datastore.beginTransaction(options);
 197   }
 198
 199   @Override
 200   public Future<KeyRange> allocateIds(String kind, long num) {
 201     checkNotNull(kind);
 202     String currentNamespace = NamespaceManager.get();
 203     try {
 204       NamespaceManager.set(getAlternateNamespace(currentNamespace));
 205       return new RethrowingFutureWrapper<KeyRange, KeyRange>(datastore.allocateIds(kind, num)) {
 206         @Override
 207         protected KeyRange wrap(KeyRange range) throws Exception {
 208           return getOriginalNamespaceKeyRange(range);
 209         }
 210       };
 211     } finally {
 212       NamespaceManager.set(currentNamespace);
 213     }
 214   }
 215
 216   @Override
 217   public Future<KeyRange> allocateIds(Key parent, String kind, long num) {
 218     checkNotNull(parent);
 219     checkNotNull(kind);
 220     String currentNamespace = NamespaceManager.get();
 221     try {
 222       NamespaceManager.set(getAlternateNamespace(parent.getNamespace()));
 223       return new RethrowingFutureWrapper<KeyRange, KeyRange>(
 224           datastore.allocateIds(getAlternateNamespaceKey(parent), kind, num)) {
 225         @Override
 226         protected KeyRange wrap(KeyRange range) throws Exception {
 227           return getOriginalNamespaceKeyRange(range);
 228         }
 229       };
 230     } finally {
 231       NamespaceManager.set(currentNamespace);
 232     }
 233   }
 234
 235   @Override
 236   public Future<DatastoreAttributes> getDatastoreAttributes() {
 237     return datastore.getDatastoreAttributes();
 238   }
 239
 240   @Override
 241   public Future<Map<Index, Index.IndexState>> getIndexes() {
 242     return datastore.getIndexes();
 243   }
 244
 245   @Override
 246   public Collection<Transaction> getActiveTransactions() {
 247     return datastore.getActiveTransactions();
 248   }
 249
 250   @Override
 251   public Transaction getCurrentTransaction() {
 252     return datastore.getCurrentTransaction();
 253   }
 254
 255   @Override
 256   public Transaction getCurrentTransaction(Transaction returnedIfNoTxn) {
 257     return datastore.getCurrentTransaction(returnedIfNoTxn);
 258   }
 259
 260   @Override
 261   public PreparedQuery prepare(Query query) {
 262     checkNotNull(query);
 263     return prepareImpl(datastore, query);
 264   }
 265
 266   @Override
 267   public PreparedQuery prepare(Transaction txn, Query query) {
 268     checkNotNull(query);
 269     return prepareImpl(getTxnAsyncDatastore(txn), query);
 270   }
 271
 272   private PreparedQuery prepareImpl(BaseDatastoreService datastore, Query query) {
 273     checkNotNull(datastore);
 274     checkNotNull(query);
 275     Query alternateQuery = getAlternateNamespaceQuery(query);
 276     return new NamespacePinnedPreparedQueryImpl(this, datastore.prepare(alternateQuery));
 277   }
 278
 279   /**
 280    * Returns {@code namespace} combined with the namespace prefix for this instance.
 281    */
 282   @VisibleForTesting
 283   private String getAlternateNamespace(String namespace) {
 284     return namespacePrefix + Strings.nullToEmpty(namespace);
 285   }
 286
 287   /**
 288    * Returns {@code namespace} with the prefix removed.
 289    */
 290   @VisibleForTesting
 291   private String getOriginalNamespace(String namespace) {
 292     checkNotNull(namespace);
 293     checkArgument(namespace.startsWith(namespacePrefix),
 294         "%s must start with %s", namespace, namespacePrefix);
 295     return namespace.substring(namespacePrefix.length());
 296   }
 297
 298   /**
 299    * Creates a copy of {@code entities}, where each entity is replaced with a copy whose key is in
 300    * the alternate namespace.
 301    */
 302   private List<Entity> getAlternateNamespaceEntities(Iterable<Entity> entities) {
 303     checkNotNull(entities);
 304     List<Entity> newEntities = Lists.newArrayList();
 305     for (Entity entity : entities) {
 306       newEntities.add(getAlternateNamespaceEntity(entity));
 307     }
 308     return newEntities;
 309   }
 310
 311   /**
 312    * Creates a copy of {@code entity}, where the key is in the alternate namespace.
 313    */
 314   private Entity getAlternateNamespaceEntity(Entity entity) {
 315     checkNotNull(entity);
 316     return cloneEntityWithNewKey(entity, getAlternateNamespaceKey(entity.getKey()));
 317   }
 318
 319   /**
 320    * Creates a copy of {@code entity}, where the key is in the original namespace.
 321    */
 322   Entity getOriginalNamespaceEntity(Entity entity) {
 323     checkNotNull(entity);
 324     return cloneEntityWithNewKey(entity, getOriginalNamespaceKey(entity.getKey()));
 325   }
 326
 327   /**
 328    * Creates a copy of {@code entity}, where the key is replaced with {@code newKey}.
 329    */
 330   private Entity cloneEntityWithNewKey(Entity entity, Key newKey) {
 331     Entity newEntity = new Entity(newKey);
 332     newEntity.setPropertiesFrom(entity);
 333     return newEntity;
 334   }
 335
 336   /**
 337    * Creates a copy of {@code entityMap}, where all the keys are replaced with the equivalent keys
 338    * in the original namespace.
 339    */
 340   private Map<Key, Entity> getOriginalNamespaceEntityMap(Map<Key, Entity> entityMap) {
 341     checkNotNull(entityMap);
 342     Map<Key, Entity> newMap = Maps.newHashMapWithExpectedSize(entityMap.size());
 343     for (Map.Entry<Key, Entity> entry : entityMap.entrySet()) {
 344       newMap.put(getOriginalNamespaceKey(entry.getKey()),
 345           getOriginalNamespaceEntity(entry.getValue()));
 346     }
 347     return newMap;
 348   }
 349
 350   /**
 351    * Creates a copy of {@code key}, but in the alternate namespace.
 352    */
 353   private Key getAlternateNamespaceKey(Key key) {
 354     checkNotNull(key);
 355     String currentNamespace = NamespaceManager.get();
 356     try {
 357       NamespaceManager.set(getAlternateNamespace(key.getNamespace()));
 358       return cloneKey(key);
 359     } finally {
 360       NamespaceManager.set(currentNamespace);
 361     }
 362   }
 363
 364   /**
 365    * Creates a copy of {@code key} in the original namespace. {@code key} must be in the alternate
 366    * namespace to begin with.
 367    */
 368   private Key getOriginalNamespaceKey(Key key) {
 369     checkNotNull(key);
 370     String currentNamespace = NamespaceManager.get();
 371     try {
 372       NamespaceManager.set(getOriginalNamespace(key.getNamespace()));
 373       return cloneKey(key);
 374     } finally {
 375       NamespaceManager.set(currentNamespace);
 376     }
 377   }
 378
 379   /**
 380    * Creates a copy of {@code keys}, where each key is replaced with a copy that is in the alternate
 381    * namespace.
 382    */
 383   private List<Key> getAlternateNamespaceKeys(Iterable<Key> keys) {
 384     checkNotNull(keys);
 385     List<Key> newKeys = Lists.newArrayList();
 386     for (Key key : keys) {
 387       newKeys.add(getAlternateNamespaceKey(key));
 388     }
 389     return newKeys;
 390   }
 391
 392   /**
 393    * Creates a copy of {@code keys}, where each key is replaced with a copy that is in the original
 394    * namespace. The elements of {@code keys} must be in the alternate namespace to begin with.
 395    */
 396   private List<Key> getOriginalNamespaceKeys(Iterable<Key> keys) {
 397     checkNotNull(keys);
 398     List<Key> newKeys = Lists.newArrayList();
 399     for (Key key : keys) {
 400       newKeys.add(getOriginalNamespaceKey(key));
 401     }
 402     return newKeys;
 403   }
 404
 405   /**
 406    * Creates a copy of {@code query}, where each component key is replaced with the equivalent key
 407    * in the alternate namespace.
 408    */
 409   private Query getAlternateNamespaceQuery(Query query) {
 410     checkNotNull(query);
 411     Query alternateQuery;
 412     String currentNamespace = NamespaceManager.get();
 413     try {
 414       NamespaceManager.set(getAlternateNamespace(query.getNamespace()));
 415       alternateQuery = new Query(query.getKind());
 416     } finally {
 417       NamespaceManager.set(currentNamespace);
 418     }
 419
 420     Key ancestor = query.getAncestor();
 421     if (ancestor != null) {
 422       alternateQuery.setAncestor(getAlternateNamespaceKey(ancestor));
 423     }
 424     for (SortPredicate sp : query.getSortPredicates()) {
 425       alternateQuery.addSort(sp.getPropertyName(), sp.getDirection());
 426     }
 427     alternateQuery.setFilter(query.getFilter());
 428     for (FilterPredicate fp : query.getFilterPredicates()) {
 429       alternateQuery.addFilter(fp.getPropertyName(), fp.getOperator(), fp.getValue());
 430     }
 431     if (query.isKeysOnly()) {
 432       alternateQuery.setKeysOnly();
 433     }
 434     for (Projection p : query.getProjections()) {
 435       alternateQuery.addProjection(p);
 436     }
 437     alternateQuery.setDistinct(query.getDistinct());
 438
 439     return alternateQuery;
 440   }
 441 /**
 442    * Creates a copy of {@code keyRange}, where the returned keys are in the alternate namespace.
 443    */
 444
 445   /**
 446    * Creates a copy of {@code keyRange}, where the returned keys are in the original namespace,
 447    * rather than the alternate namespace.
 448    */
 449   private KeyRange getOriginalNamespaceKeyRange(KeyRange keyRange) {
 450     Key alternateParent = keyRange.getStart().getParent();
 451     Key originalParent = alternateParent == null ? null : getOriginalNamespaceKey(alternateParent);
 452     String currentNamespace = NamespaceManager.get();
 453     try {
 454       NamespaceManager.set(getOriginalNamespace(keyRange.getStart().getNamespace()));
 455       return new KeyRange(originalParent, keyRange.getStart().getKind(),
 456           keyRange.getStart().getId(), keyRange.getEnd().getId());
 457     } finally {
 458       NamespaceManager.set(currentNamespace);
 459     }
 460   }
 461
 462   /**
 463    * Gets a version of the Datastore that is linked to {@code txn}.
 464    */
 465   private AsyncDatastoreService getTxnAsyncDatastore( Transaction txn) {
 466     return new TransactionLinkedAsyncDatastoreServiceImpl(datastore, txn);
 467   }
 468
 469   /**
 470    * Creates a copy of {@code key}, where the copy (and each of its components) takes its namespace
 471    * from the {@link NamespaceManager}.
 472    */
 473   private static Key cloneKey(Key key) {
 474     checkNotNull(key);
 475     Key parentKey = key.getParent();
 476     String name = key.getName();
 477     long id = key.getId();
 478     if (parentKey == null) {
 479       if (name != null) {
 480         return new Entity(key.getKind(), name).getKey();
 481       } else if (id != 0L) {
 482         return new Entity(key.getKind(), key.getId()).getKey();
 483       } else {
 484         return new Entity(key.getKind()).getKey();
 485       }
 486     } else {
 487       if (name != null) {
 488         return new Entity(key.getKind(), name, cloneKey(parentKey)).getKey();
 489       } else if (id != 0L) {
 490         return new Entity(key.getKind(), key.getId(), cloneKey(parentKey)).getKey();
 491       } else {
 492         return new Entity(key.getKind(), cloneKey(parentKey)).getKey();
 493       }
 494     }
 495   }
 496 }