11package com .google .gcloud .datastore ;
22
3+ import com .google .api .services .datastore .DatastoreV1 ;
4+ import com .google .api .services .datastore .client .Datastore ;
5+ import com .google .api .services .datastore .client .DatastoreException ;
6+ import com .google .common .collect .AbstractIterator ;
7+
8+ import java .util .HashMap ;
39import java .util .Iterator ;
10+ import java .util .Map ;
411
512
613final class DatastoreServiceImpl implements DatastoreService {
714
815 private final DatastoreServiceOptions options ;
16+ private final Datastore datastore ;
917
10- DatastoreServiceImpl (DatastoreServiceOptions options ) {
18+ DatastoreServiceImpl (DatastoreServiceOptions options , Datastore datastore ) {
1119 this .options = options ;
20+ this .datastore = datastore ;
1221 }
1322
1423 @ Override
@@ -35,9 +44,27 @@ public Key allocateId(PartialKey key) {
3544
3645 @ Override
3746 public Iterator <Key > allocateIds (PartialKey ... key ) {
38- // TODO Auto-generated method stub
39- // Will need to populate "force" after b/18594027 is fixed.
40- return null ;
47+ DatastoreV1 .AllocateIdsRequest .Builder requestPb = DatastoreV1 .AllocateIdsRequest .newBuilder ();
48+ for (PartialKey k : key ) {
49+ requestPb .addKey (k .toPb ());
50+ }
51+ // TODO(ozarov): will need to populate "force" after b/18594027 is fixed.
52+ try {
53+ DatastoreV1 .AllocateIdsResponse responsePb = datastore .allocateIds (requestPb .build ());
54+ final Iterator <DatastoreV1 .Key > keys = responsePb .getKeyList ().iterator ();
55+ return new AbstractIterator <Key >() {
56+
57+ @ Override
58+ protected Key computeNext () {
59+ if (keys .hasNext ()) {
60+ return Key .fromPb (keys .next ());
61+ }
62+ return endOfData ();
63+ }
64+ };
65+ } catch (DatastoreException e ) {
66+ throw DatastoreServiceException .translateAndPropagate (e );
67+ }
4168 }
4269
4370 @ Override
@@ -46,29 +73,82 @@ public Entity get(Key key) {
4673 }
4774
4875 @ Override
49- public Iterator <Entity > get (Key ... key ) {
50- // TODO Auto-generated method stub
51- return null ;
76+ public Iterator <Entity > get (final Key ... key ) {
77+ DatastoreV1 .LookupRequest .Builder requestPb = DatastoreV1 .LookupRequest .newBuilder ();
78+ for (Key k : key ) {
79+ requestPb .addKey (k .toPb ());
80+ }
81+ try {
82+ DatastoreV1 .LookupResponse responsePb = datastore .lookup (requestPb .build ());
83+ final Map <Key , Entity > result = new HashMap <>();
84+ for (DatastoreV1 .EntityResult entityResultPb : responsePb .getFoundList ()) {
85+ Entity entity = Entity .fromPb (entityResultPb .getEntity ());
86+ result .put (entity .key (), entity );
87+ }
88+ return new AbstractIterator <Entity >() {
89+ int index ;
90+
91+ @ Override
92+ protected Entity computeNext () {
93+ if (index < key .length ) {
94+ return result .get (key [index ++]);
95+ }
96+ return endOfData ();
97+ }
98+ };
99+ } catch (DatastoreException e ) {
100+ throw DatastoreServiceException .translateAndPropagate (e );
101+ }
52102 }
53103
54104 @ Override
55105 public void add (Entity ... entity ) {
56- // TODO Auto-generated method stub
106+ DatastoreV1 .Mutation .Builder mutationPb = DatastoreV1 .Mutation .newBuilder ();
107+ for (Entity e : entity ) {
108+ mutationPb .addInsert (e .toPb ());
109+ }
110+ comitMutation (mutationPb );
111+ }
112+
113+ private void comitMutation (DatastoreV1 .Mutation .Builder mutationPb ) {
114+ if (options .force ()) {
115+ mutationPb .setForce (true );
116+ }
117+ DatastoreV1 .CommitRequest .Builder requestPb = DatastoreV1 .CommitRequest .newBuilder ();
118+ requestPb .setMode (DatastoreV1 .CommitRequest .Mode .NON_TRANSACTIONAL );
119+ requestPb .setMutation (mutationPb .build ());
120+ try {
121+ datastore .commit (requestPb .build ());
122+ } catch (DatastoreException e ) {
123+ throw DatastoreServiceException .translateAndPropagate (e );
124+ }
57125 }
58126
59127 @ Override
60128 public void update (Entity ... entity ) {
61- // TODO Auto-generated method stub
129+ DatastoreV1 .Mutation .Builder mutationPb = DatastoreV1 .Mutation .newBuilder ();
130+ for (Entity e : entity ) {
131+ mutationPb .addUpdate (e .toPb ());
132+ }
133+ comitMutation (mutationPb );
62134 }
63135
64136 @ Override
65137 public void put (Entity ... entity ) {
66- // TODO Auto-generated method stub
138+ DatastoreV1 .Mutation .Builder mutationPb = DatastoreV1 .Mutation .newBuilder ();
139+ for (Entity e : entity ) {
140+ mutationPb .addUpsert (e .toPb ());
141+ }
142+ comitMutation (mutationPb );
67143 }
68144
69145 @ Override
70146 public void delete (Key ... key ) {
71- // TODO Auto-generated method stub
147+ DatastoreV1 .Mutation .Builder mutationPb = DatastoreV1 .Mutation .newBuilder ();
148+ for (Key k : key ) {
149+ mutationPb .addDelete (k .toPb ());
150+ }
151+ comitMutation (mutationPb );
72152 }
73153
74154 @ Override
@@ -81,4 +161,8 @@ public QueryResult<PartialEntity> runQuery(Query query) {
81161 // TODO Auto-generated method stub
82162 return null ;
83163 }
164+
165+ Datastore datastore () {
166+ return datastore ;
167+ }
84168}
0 commit comments