final com.google.datastore.v1.Entity toPb() { com.google.datastore.v1.Entity.Builder entityPb = com.google.datastore.v1.Entity.newBuilder(); for (Map.Entry<String, Value<?>> entry : properties.entrySet()) { entityPb.putProperties(entry.getKey(), entry.getValue().toPb()); } if (key != null) { entityPb.setKey(key.toPb()); } return entityPb.build(); } }
entityBuilder.setKey(datastoreKey); for (Entry<String, ByteIterator> val : values.entrySet()) { entityBuilder.getMutableProperties()
com.google.datastore.v1.Entity getPb(final String namespace) { final Map<String, com.google.datastore.v1.Value> propertiesLocal = entity.getProperties() .entrySet() .stream() .collect(Collectors.toMap( Map.Entry::getKey, e -> Value.builder(e.getValue()).build().getPb(namespace))); return com.google.datastore.v1.Entity .newBuilder() .putAllProperties(propertiesLocal) .setKey(getKey().getPb(namespace)) .build(); } }
entities.add( EntityResult.newBuilder() .setEntity(Entity.newBuilder().setKey(makeKey("key" + i, i + 1))) .build());
/** * Convert TableRow to Entity * * @param row TableRow of bigquery * @return converted Entity * @throws ParseException */ public Entity convertTableRowToEntity(TableRow row) throws ParseException { String keyName = row.get(keyColumn).toString(); Key key = getKey(keyName); Entity.Builder builder = Entity.newBuilder().setKey(key); Set<Map.Entry<String, Object>> entries = row.entrySet(); for (Map.Entry<String, Object> entry : entries) { // Skip on the key column if (entry.getKey().equals(keyColumn)) { continue; } // Put a value in the builder String propertyName = entry.getKey(); Object value = entry.getValue(); Value v = convertToDatastoreValue(propertyName, value); if (v != null) { builder.putProperties(propertyName, v); } } return builder.build(); }
/** * Build an entity for the given ancestorKey, kind, namespace and value. * * @param largePropertySize if greater than 0, add an unindexed property of the given size. */ static Entity makeEntity( Long value, Key ancestorKey, String kind, @Nullable String namespace, int largePropertySize) { Entity.Builder entityBuilder = Entity.newBuilder(); Key.Builder keyBuilder = makeKey(ancestorKey, kind, UUID.randomUUID().toString()); // NOTE: Namespace is not inherited between keys created with DatastoreHelper.makeKey, so // we must set the namespace on keyBuilder. TODO: Once partitionId inheritance is added, // we can simplify this code. if (namespace != null) { keyBuilder.getPartitionIdBuilder().setNamespaceId(namespace); } entityBuilder.setKey(keyBuilder.build()); entityBuilder.putProperties("value", makeValue(value).build()); if (largePropertySize > 0) { entityBuilder.putProperties( "unindexed_value", makeValue(new String(new char[largePropertySize]).replace("\0", "A")) .setExcludeFromIndexes(true) .build()); } return entityBuilder.build(); }
private void datastoreWriterFnTest(int numMutations) throws Exception { // Create the requested number of mutations. List<Mutation> mutations = new ArrayList<>(numMutations); for (int i = 0; i < numMutations; ++i) { mutations.add( makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build()); } DatastoreWriterFn datastoreWriter = new DatastoreWriterFn( StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher()); DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter); doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE); doFnTester.processBundle(mutations); int start = 0; while (start < numMutations) { int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START); CommitRequest.Builder commitRequest = CommitRequest.newBuilder(); commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL); commitRequest.addAllMutations(mutations.subList(start, end)); // Verify all the batch requests were made with the expected mutations. verify(mockDatastore, times(1)).commit(commitRequest.build()); start = end; } }
/** Tests {@link DatastoreWriterFn} with a failed request which is retried. */ @Test public void testDatatoreWriterFnRetriesErrors() throws Exception { List<Mutation> mutations = new ArrayList<>(); int numRpcs = 2; for (int i = 0; i < DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START * numRpcs; ++i) { mutations.add( makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build()); } CommitResponse successfulCommit = CommitResponse.getDefaultInstance(); when(mockDatastore.commit(any(CommitRequest.class))) .thenReturn(successfulCommit) .thenThrow(new DatastoreException("commit", Code.DEADLINE_EXCEEDED, "", null)) .thenReturn(successfulCommit); DatastoreWriterFn datastoreWriter = new DatastoreWriterFn( StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher()); DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter); doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE); doFnTester.processBundle(mutations); }
@ProcessElement public void processElement(ProcessContext c) { Entity.Builder entityBuilder = Entity.newBuilder(); Key key = makeKey(makeKey(kind, ancestorKey).build(), kind, c.element().getKey()).build(); entityBuilder.setKey(key); List<Value> candidates = new ArrayList<>(); Map<String, Value> properties = new HashMap<>(); for (CompletionCandidate tag : c.element().getValue()) { Entity.Builder tagEntity = Entity.newBuilder(); properties.put("tag", makeValue(tag.value).build()); properties.put("count", makeValue(tag.count).build()); candidates.add(makeValue(tagEntity).build()); } properties.put("candidates", makeValue(candidates).build()); entityBuilder.putAllProperties(properties); c.output(entityBuilder.build()); } }
@ProcessElement public void processElement(ProcessContext c) throws InvalidProtocolBufferException { String entityJson = c.element(); Entity.Builder entityBuilder = Entity.newBuilder(); entityJsonParser.merge(entityJson, entityBuilder); // Build entity who's key has an empty project Id. // This allows DatastoreIO to handle what project Entities are loaded into Key k = entityBuilder.build().getKey(); entityBuilder.setKey(Key.newBuilder() .addAllPath(k.getPathList()) .setPartitionId(PartitionId.newBuilder() .setProjectId("") .setNamespaceId(k.getPartitionId().getNamespaceId()))); c.output(entityBuilder.build()); } }
/** * Add a greeting to the specified guestbook. */ private void addGreeting(String guestbookName, String user, String message) throws DatastoreException { Entity.Builder greeting = Entity.newBuilder(); greeting.setKey(makeKey(GUESTBOOK_KIND, guestbookName, GREETING_KIND)); greeting.getMutableProperties().put(USER_PROPERTY, makeValue(user).build()); greeting.getMutableProperties().put(MESSAGE_PROPERTY, makeValue(message).build()); greeting.getMutableProperties().put(DATE_PROPERTY, makeValue(new Date()).build()); Key greetingKey = insert(greeting.build()); System.out.println("greeting key is: " + greetingKey); }
/** Builds a response of the given timestamp. */ private static RunQueryResponse makeLatestTimestampResponse(long timestamp) { RunQueryResponse.Builder timestampResponse = RunQueryResponse.newBuilder(); Entity.Builder entity = Entity.newBuilder(); entity.setKey(makeKey("dummyKind", "dummyId")); entity.putProperties("timestamp", makeValue(new Date(timestamp * 1000)).build()); EntityResult.Builder entityResult = EntityResult.newBuilder(); entityResult.setEntity(entity); QueryResultBatch.Builder batch = QueryResultBatch.newBuilder(); batch.addEntityResults(entityResult); timestampResponse.setBatch(batch); return timestampResponse.build(); }
/** Builds a per-kind statistics response with the given entity size. */ private static RunQueryResponse makeStatKindResponse(long entitySizeInBytes) { RunQueryResponse.Builder statKindResponse = RunQueryResponse.newBuilder(); Entity.Builder entity = Entity.newBuilder(); entity.setKey(makeKey("dummyKind", "dummyId")); entity.putProperties("entity_bytes", makeValue(entitySizeInBytes).build()); EntityResult.Builder entityResult = EntityResult.newBuilder(); entityResult.setEntity(entity); QueryResultBatch.Builder batch = QueryResultBatch.newBuilder(); batch.addEntityResults(entityResult); statKindResponse.setBatch(batch); return statKindResponse.build(); }
private Builder(final Key key) { this.entity = com.google.datastore.v1.Entity.newBuilder().setKey(key.getPb()); this.properties = Maps.newHashMap(); }
/** * Set the key for this entity. * * @param key the key to set for this entity. * @return this entity builder. */ public Builder key(final Key key) { entity.setKey(key.getPb()); return this; }
/** Test that entities with incomplete keys cannot be deleted. */ @Test public void testDeleteEntitiesWithIncompleteKeys() throws Exception { Key key = makeKey("bird").build(); Entity entity = Entity.newBuilder().setKey(key).build(); DeleteEntityFn deleteEntityFn = new DeleteEntityFn(); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Entities to be deleted from the Cloud Datastore must have complete keys"); deleteEntityFn.apply(entity); }
/** Test that entities with incomplete keys cannot be updated. */ @Test public void testAddEntitiesWithIncompleteKeys() throws Exception { Key key = makeKey("bird").build(); Entity entity = Entity.newBuilder().setKey(key).build(); UpsertFn upsertFn = new UpsertFn(); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Entities to be written to the Cloud Datastore must have complete keys"); upsertFn.apply(entity); }
final com.google.datastore.v1.Entity toPb() { com.google.datastore.v1.Entity.Builder entityPb = com.google.datastore.v1.Entity.newBuilder(); for (Map.Entry<String, Value<?>> entry : properties.entrySet()) { entityPb.putProperties(entry.getKey(), entry.getValue().toPb()); } if (key != null) { entityPb.setKey(key.toPb()); } return entityPb.build(); } }
/** Test that entities with valid keys are transformed to delete mutations. */ @Test public void testDeleteEntities() throws Exception { Key key = makeKey("bird", "finch").build(); Entity entity = Entity.newBuilder().setKey(key).build(); DeleteEntityFn deleteEntityFn = new DeleteEntityFn(); Mutation exceptedMutation = makeDelete(entity.getKey()).build(); assertEquals(deleteEntityFn.apply(entity), exceptedMutation); }
@Test /** Test that entities with valid keys are transformed to upsert mutations. */ public void testAddEntities() throws Exception { Key key = makeKey("bird", "finch").build(); Entity entity = Entity.newBuilder().setKey(key).build(); UpsertFn upsertFn = new UpsertFn(); Mutation exceptedMutation = makeUpsert(entity).build(); assertEquals(upsertFn.apply(entity), exceptedMutation); }