Skip to content

Commit 6d36012

Browse files
ableegoldmanableegoldman
committed
basic implementation adding kafka partition
1 parent 2c0e55f commit 6d36012

File tree

4 files changed

+70
-36
lines changed

4 files changed

+70
-36
lines changed

‎kafka-client/src/main/java/dev/responsive/kafka/internal/db/mongo/KVDoc.java‎

Lines changed: 31 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -24,47 +24,59 @@
2424
public class KVDoc {
2525

2626
// TODO(agavra): figure out if we can use @BsonProperty to set the names explicitly
27-
public static final String ID = "_id";
27+
public static final String KEY = "_id";
2828
public static final String VALUE = "value";
29+
public static final String KAFKA_PARTITION = "partition";
2930
public static final String EPOCH = "epoch";
3031
public static final String TOMBSTONE_TS = "tombstoneTs";
3132

3233
@BsonId
33-
byte[] id;
34+
byte[] key;
3435
byte[] value;
36+
int kafkaPartition;
3537
long epoch;
3638
Date tombstoneTs;
3739

3840
public KVDoc() {
3941
}
4042

41-
public KVDoc(final byte[] key, final byte[] value, final long epoch) {
43+
public KVDoc(final byte[] key, final byte[] value, final int kafkaPartition, final long epoch) {
44+
this.key = key;
4245
this.value = value;
46+
this.kafkaPartition = kafkaPartition;
4347
this.epoch = epoch;
4448
}
4549

4650
public byte[] getKey() {
47-
return id;
51+
return key;
4852
}
4953

5054
public void setKey(final byte[] id) {
51-
this.id = id;
55+
this.key = id;
56+
}
57+
58+
public byte[] getValue() {
59+
return value;
5260
}
5361

5462
public void setValue(final byte[] value) {
5563
this.value = value;
5664
}
5765

66+
public long getEpoch() {
67+
return epoch;
68+
}
69+
5870
public void setEpoch(final long epoch) {
5971
this.epoch = epoch;
6072
}
6173

62-
public byte[] getValue() {
63-
return value;
74+
public int getKafkaPartition() {
75+
return kafkaPartition;
6476
}
6577

66-
public long getEpoch() {
67-
return epoch;
78+
public void setKafkaPartition(final int kafkaPartition) {
79+
this.kafkaPartition = kafkaPartition;
6880
}
6981

7082
public Date getTombstoneTs() {
@@ -83,25 +95,28 @@ public boolean equals(final Object o) {
8395
if (o == null || getClass() != o.getClass()) {
8496
return false;
8597
}
86-
final KVDoc kvDoc = (KVDoc) o;
87-
return epoch == kvDoc.epoch
88-
&& Arrays.equals(id, kvDoc.id)
89-
&& Arrays.equals(value, kvDoc.value)
90-
&& Objects.equals(tombstoneTs, kvDoc.tombstoneTs);
98+
final KVDoc other = (KVDoc) o;
99+
return epoch == other.epoch
100+
&& kafkaPartition == other.kafkaPartition
101+
&& Arrays.equals(key, other.key)
102+
&& Arrays.equals(value, other.value)
103+
&& Objects.equals(tombstoneTs, other.tombstoneTs);
91104
}
92105

93106
@Override
94107
public int hashCode() {
95-
int result = Objects.hash(id, epoch, tombstoneTs);
108+
int result = Objects.hash(kafkaPartition, epoch, tombstoneTs);
109+
result = 31 * result + Arrays.hashCode(key);
96110
result = 31 * result + Arrays.hashCode(value);
97111
return result;
98112
}
99113

100114
@Override
101115
public String toString() {
102116
return "KVDoc{"
103-
+ "id=" + Arrays.toString(id)
117+
+ "id=" + Arrays.toString(key)
104118
+ ", value=" + Arrays.toString(value)
119+
+ ", kafkaPartition=" + kafkaPartition
105120
+ ", epoch=" + epoch
106121
+ ", tombstoneTs=" + tombstoneTs
107122
+ '}';

‎kafka-client/src/main/java/dev/responsive/kafka/internal/db/mongo/KVMetadataDoc.java‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,23 +20,23 @@
2020

2121
public class KVMetadataDoc {
2222

23-
public static final String PARTITION = "_id";
23+
public static final String KAFKA_PARTITION = "_id";
2424
public static final String OFFSET = "offset";
2525
public static final String EPOCH = "epoch";
2626

27-
int partition;
27+
int kafkaPartition;
2828
long offset;
2929
long epoch;
3030

3131
public KVMetadataDoc() {
3232
}
3333

3434
public int partition() {
35-
return partition;
35+
return kafkaPartition;
3636
}
3737

38-
public void setPartition(final int partition) {
39-
this.partition = partition;
38+
public void setKafkaPartition(final int kafkaPartition) {
39+
this.kafkaPartition = kafkaPartition;
4040
}
4141

4242
public long offset() {
@@ -64,20 +64,20 @@ public boolean equals(final Object o) {
6464
return false;
6565
}
6666
final KVMetadataDoc that = (KVMetadataDoc) o;
67-
return partition == that.partition
67+
return kafkaPartition == that.kafkaPartition
6868
&& offset == that.offset
6969
&& epoch == that.epoch;
7070
}
7171

7272
@Override
7373
public int hashCode() {
74-
return Objects.hash(partition, epoch, offset);
74+
return Objects.hash(kafkaPartition, epoch, offset);
7575
}
7676

7777
@Override
7878
public String toString() {
7979
return "KVMetadataDoc{"
80-
+ ", partition=" + partition
80+
+ ", partition=" + kafkaPartition
8181
+ ", offset=" + offset
8282
+ ", epoch=" + epoch
8383
+ '}';

‎kafka-client/src/main/java/dev/responsive/kafka/internal/db/mongo/MongoKVTable.java‎

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -36,12 +36,14 @@
3636
import com.mongodb.client.result.UpdateResult;
3737
import dev.responsive.kafka.internal.db.MongoKVFlushManager;
3838
import dev.responsive.kafka.internal.db.RemoteKVTable;
39+
import dev.responsive.kafka.internal.utils.Iterators;
3940
import java.time.Instant;
4041
import java.util.Date;
4142
import java.util.concurrent.ConcurrentHashMap;
4243
import java.util.concurrent.ConcurrentMap;
4344
import java.util.concurrent.TimeUnit;
4445
import org.apache.kafka.common.utils.Bytes;
46+
import org.apache.kafka.streams.KeyValue;
4547
import org.apache.kafka.streams.state.KeyValueIterator;
4648
import org.bson.codecs.configuration.CodecProvider;
4749
import org.bson.codecs.configuration.CodecRegistry;
@@ -105,10 +107,10 @@ public String name() {
105107
@Override
106108
public MongoKVFlushManager init(final int kafkaPartition) {
107109
final KVMetadataDoc metaDoc = metadata.findOneAndUpdate(
108-
Filters.eq(KVMetadataDoc.PARTITION, kafkaPartition),
110+
Filters.eq(KVMetadataDoc.KAFKA_PARTITION, kafkaPartition),
109111
Updates.combine(
110-
Updates.setOnInsert(KVMetadataDoc.PARTITION, kafkaPartition),
111-
Updates.setOnInsert(KVMetadataDoc.PARTITION, kafkaPartition),
112+
Updates.setOnInsert(KVMetadataDoc.KAFKA_PARTITION, kafkaPartition),
113+
Updates.setOnInsert(KVMetadataDoc.KAFKA_PARTITION, kafkaPartition),
112114
Updates.setOnInsert(KVMetadataDoc.OFFSET, NO_COMMITTED_OFFSET),
113115
Updates.inc(KVMetadataDoc.EPOCH, 1) // will set the value to 1 if it doesn't exist
114116
),
@@ -129,7 +131,7 @@ public MongoKVFlushManager init(final int kafkaPartition) {
129131

130132
@Override
131133
public byte[] get(final int kafkaPartition, final Bytes key, final long minValidTs) {
132-
final KVDoc v = docs.find(Filters.eq(KVDoc.ID, key.get())).first();
134+
final KVDoc v = docs.find(Filters.eq(KVDoc.KEY, key.get())).first();
133135
return v == null ? null : v.getValue();
134136
}
135137

@@ -140,12 +142,20 @@ public KeyValueIterator<Bytes, byte[]> range(
140142
final Bytes to,
141143
final long minValidTs
142144
) {
143-
throw new UnsupportedOperationException();
145+
final Iterable<KVDoc> results = docs.find(Filters.and(
146+
Filters.eq(KVDoc.KAFKA_PARTITION, kafkaPartition),
147+
Filters.gte(KVDoc.KEY, from.get()),
148+
Filters.lte(KVDoc.KEY, to.get())
149+
));
150+
return Iterators.kv(results.iterator(), MongoKVTable::extractKeyValue);
144151
}
145152

146153
@Override
147154
public KeyValueIterator<Bytes, byte[]> all(final int kafkaPartition, final long minValidTs) {
148-
throw new UnsupportedOperationException();
155+
final Iterable<KVDoc> results = docs.find(
156+
Filters.eq(KVDoc.KAFKA_PARTITION, kafkaPartition)
157+
);
158+
return Iterators.kv(results.iterator(), MongoKVTable::extractKeyValue);
149159
}
150160

151161
@Override
@@ -158,7 +168,8 @@ public WriteModel<KVDoc> insert(
158168
final long epoch = kafkaPartitionToEpoch.get(kafkaPartition);
159169
return new UpdateOneModel<>(
160170
Filters.and(
161-
Filters.eq(KVDoc.ID, key.get()),
171+
Filters.eq(KVDoc.KEY, key.get()),
172+
Filters.eq(KVDoc.KAFKA_PARTITION, kafkaPartition),
162173
Filters.lte(KVDoc.EPOCH, epoch)
163174
),
164175
Updates.combine(
@@ -175,7 +186,8 @@ public WriteModel<KVDoc> delete(final int kafkaPartition, final Bytes key) {
175186
final long epoch = kafkaPartitionToEpoch.get(kafkaPartition);
176187
return new UpdateOneModel<>(
177188
Filters.and(
178-
Filters.eq(KVDoc.ID, key.get()),
189+
Filters.eq(KVDoc.KEY, key.get()),
190+
Filters.eq(KVDoc.KAFKA_PARTITION, kafkaPartition),
179191
Filters.lte(KVDoc.EPOCH, epoch)
180192
),
181193
Updates.combine(
@@ -190,7 +202,7 @@ public WriteModel<KVDoc> delete(final int kafkaPartition, final Bytes key) {
190202
@Override
191203
public long fetchOffset(final int kafkaPartition) {
192204
final KVMetadataDoc result = metadata.find(
193-
Filters.eq(KVMetadataDoc.PARTITION, kafkaPartition)
205+
Filters.eq(KVMetadataDoc.KAFKA_PARTITION, kafkaPartition)
194206
).first();
195207
if (result == null) {
196208
throw new IllegalStateException("Expected to find metadata row");
@@ -203,7 +215,7 @@ public UpdateResult setOffset(final int kafkaPartition, final long offset) {
203215

204216
return metadata.updateOne(
205217
Filters.and(
206-
Filters.eq(KVMetadataDoc.PARTITION, kafkaPartition),
218+
Filters.eq(KVMetadataDoc.KAFKA_PARTITION, kafkaPartition),
207219
Filters.lte(KVMetadataDoc.EPOCH, epoch)
208220
),
209221
Updates.combine(
@@ -219,7 +231,7 @@ public long localEpoch(final int kafkaPartition) {
219231

220232
public long fetchEpoch(final int kafkaPartition) {
221233
final KVMetadataDoc result = metadata.find(
222-
Filters.eq(KVMetadataDoc.PARTITION, kafkaPartition)
234+
Filters.eq(KVMetadataDoc.KAFKA_PARTITION, kafkaPartition)
223235
).first();
224236
if (result == null) {
225237
throw new IllegalStateException("Expected to find metadata row");
@@ -233,4 +245,11 @@ public long approximateNumEntries(final int kafkaPartition) {
233245
return 0;
234246
}
235247

248+
private static KeyValue<Bytes, byte[]> extractKeyValue(final KVDoc row) {
249+
return new KeyValue<>(
250+
Bytes.wrap(row.getKey()),
251+
row.getValue()
252+
);
253+
}
254+
236255
}

‎kafka-client/src/test/java/dev/responsive/kafka/integration/StoreQueryIntegrationTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@
7474
public class StoreQueryIntegrationTest {
7575

7676
@RegisterExtension
77-
static ResponsiveExtension EXTENSION = new ResponsiveExtension(StorageBackend.CASSANDRA);
77+
static ResponsiveExtension EXTENSION = new ResponsiveExtension(StorageBackend.MONGO_DB);
7878

7979
private static final String INPUT_TOPIC = "input";
8080
private static final String OUTPUT_TOPIC = "output";

0 commit comments

Comments
 (0)