Skip to content

Commit

Permalink
Merge pull request #370 from oleschoenburg/os/decode-keys
Browse files Browse the repository at this point in the history 
feat: decode keys in state list command
  • Loading branch information
@Zelldon
Zelldon authored Mar 19, 2024
2 parents 8ac9722 + 81a5e01 commit 3d42cc0
Show file tree
Hide file tree
Showing 3 changed files with 240 additions and 10 deletions.
117 changes: 117 additions & 0 deletions cli/src/main/java/io/zell/zdb/state/KeyFormatters.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
/*
* Copyright © 2021 Christopher Kujawa ([email protected])
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.zell.zdb.state;

import io.camunda.zeebe.db.DbValue;
import io.camunda.zeebe.db.impl.*;
import io.camunda.zeebe.protocol.ZbColumnFamilies;
import org.agrona.concurrent.UnsafeBuffer;

import java.util.HexFormat;
import java.util.Map;

public interface KeyFormatters {
KeyFormatter HEX_FORMATTER = new KeyFormatter.HexFormatter();
Map<ZbColumnFamilies, KeyFormatter> FORMATTERS = Map.of(
ZbColumnFamilies.DEFAULT, KeyFormatter.DbValueFormatter.of("s"),
ZbColumnFamilies.KEY, KeyFormatter.DbValueFormatter.of("s"),
ZbColumnFamilies.BANNED_INSTANCE, KeyFormatter.DbValueFormatter.of("l"),
ZbColumnFamilies.MESSAGE_SUBSCRIPTION_BY_KEY, KeyFormatter.DbValueFormatter.of("ls")
);

KeyFormatter forColumnFamily(ZbColumnFamilies columnFamily);

static KeyFormatters ofDefault() {
return columnFamily -> FORMATTERS.getOrDefault(columnFamily, HEX_FORMATTER);
}

static KeyFormatters ofHex() {
return columnFamily -> HEX_FORMATTER;
}

static KeyFormatters ofFormat(String format) {
return columnFamily -> KeyFormatter.DbValueFormatter.of(format);
}

interface KeyFormatter {
String formatKey(final byte[] key);

final class HexFormatter implements KeyFormatter {
@Override
public String formatKey(final byte[] key) {
return HexFormat.ofDelimiter(" ").formatHex(key);
}
}

final class DbValueFormatter implements KeyFormatter {
final DbValue[] values;

private DbValueFormatter(DbValue[] values) {
this.values = values;
}

/**
* Takes a format string consisting of a sequence of 's', 'l', 'i', 'b', and 'B' characters
* which specify the format of the keys. 's' is a string, 'l' is a long, 'i' is an int, 'b' is
* a byte, and 'B' is a byte array.
*/
public static DbValueFormatter of(String keyFormat) {
final var components = new DbValue[keyFormat.length()];
final var chars = keyFormat.toCharArray();
for (int i = 0; i < chars.length; i++) {
components[i] =
switch (chars[i]) {
case 's' -> new DbString();
case 'l' -> new DbLong();
case 'i' -> new DbInt();
case 'b' -> new DbByte();
case 'B' -> new DbBytes();
default ->
throw new IllegalArgumentException("Unknown key format component: " + chars[i]);
};
}
return new DbValueFormatter(components);
}

@Override
public String formatKey(final byte[] key) {
final var formatted = new StringBuilder();
final var keyBuffer = new UnsafeBuffer(key);
int offset = 8;
for (final var dbValue : values) {
dbValue.wrap(keyBuffer, offset, key.length - offset);
offset += dbValue.getLength();
if (!formatted.isEmpty()) {
formatted.append(":");
}
switch (dbValue) {
case DbString dbString -> formatted.append(dbString);
case DbLong dbLong -> formatted.append(dbLong.getValue());
case DbInt dbInt -> formatted.append(dbInt.getValue());
case DbByte dbByte -> formatted.append(dbByte.getValue());
case DbBytes dbBytes -> {
final var buf = dbBytes.getDirectBuffer();
final var bytes = new byte[dbBytes.getLength()];
buf.getBytes(0, bytes);
formatted.append(HEX_FORMATTER.formatKey(bytes)); }
default -> formatted.append(dbValue);
}
}
return formatted.toString();
}
}
}
}
37 changes: 27 additions & 10 deletions cli/src/main/java/io/zell/zdb/state/StateCommand.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
import io.camunda.zeebe.protocol.ZbColumnFamilies;
import io.zell.zdb.JsonPrinter;
import java.nio.file.Path;
import java.util.HexFormat;
import java.util.concurrent.Callable;
import picocli.CommandLine;
import picocli.CommandLine.Command;
Expand Down Expand Up @@ -53,7 +52,14 @@ public int list(
names = {"-cf", "--columnFamily"},
paramLabel = "COLUMNFAMILY",
description = "The column family name to filter for")
final String columnFamilyName) {
final String columnFamilyName,
@Option(
names = {"-kf", "--keyFormat"},
paramLabel = "KEY_FORMAT",
description =
"The format of the key (default, hex, or a format string like 'silbB' for 'string, int, long, byte, byte[])")
final String keyFormat) {
final var keyFormatters = chooseKeyFormatters(keyFormat);

new JsonPrinter()
.surround(
Expand All @@ -62,13 +68,15 @@ public int list(
// we print incrementally in order to avoid to build up big state in the application
if (noColumnFamilyGiven(columnFamilyName)) {
zeebeDbReader.visitDBWithJsonValues(
((cf, key, valueJson) ->
printer.accept(
String.format(
ENTRY_FORMAT,
cf,
HexFormat.ofDelimiter(" ").formatHex(key),
valueJson))));
((cfName, key, valueJson) -> {
final var cf = ZbColumnFamilies.valueOf(cfName);
printer.accept(
String.format(
ENTRY_FORMAT,
cf,
keyFormatters.forColumnFamily(cf).formatKey(key),
valueJson));
}));
} else {
final var cf = ZbColumnFamilies.valueOf(columnFamilyName);
zeebeDbReader.visitDBWithPrefix(
Expand All @@ -78,13 +86,22 @@ public int list(
String.format(
ENTRY_FORMAT,
cf,
HexFormat.ofDelimiter(" ").formatHex(key),
keyFormatters.forColumnFamily(cf).formatKey(key),
valueJson))));
}
});
return 0;
}

private KeyFormatters chooseKeyFormatters(String keyFormat) {
return switch (keyFormat) {
case "default", "" -> KeyFormatters.ofDefault();
case null -> KeyFormatters.ofDefault();
case "hex" -> KeyFormatters.ofHex();
default -> KeyFormatters.ofFormat(keyFormat);
};
}

private static boolean noColumnFamilyGiven(String columnFamilyName) {
return columnFamilyName == null || columnFamilyName.isEmpty();
}
Expand Down
96 changes: 96 additions & 0 deletions cli/src/test/java/io/zell/zdb/KeyFormattersTest.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
package io.zell.zdb;

import io.camunda.zeebe.db.impl.*;
import io.camunda.zeebe.protocol.ZbColumnFamilies;
import io.zell.zdb.state.KeyFormatters;
import org.agrona.ExpandableArrayBuffer;
import org.agrona.collections.MutableInteger;
import org.junit.jupiter.api.Test;

import static org.assertj.core.api.Assertions.assertThat;

final class KeyFormattersTest {
@Test
void shouldFallbackToHex() {
// given
final var key = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

// when -- using a column family that is not registered with a specific formatter
final var formatter = KeyFormatters.ofDefault().forColumnFamily(ZbColumnFamilies.MIGRATIONS_STATE);

// then
assertThat(formatter.formatKey(key)).isEqualTo("01 02 03 04 05 06 07 08 09 0a");
}

@Test
void shouldDecodeWithHex() {
// given
final var key = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

// when
final var formatter = KeyFormatters.ofHex().forColumnFamily(ZbColumnFamilies.DEFAULT);

// then
assertThat(formatter.formatKey(key)).isEqualTo("01 02 03 04 05 06 07 08 09 0a");
}

@Test
void shouldDecodeWithFormat() {
// given -- key consisting of ColumnFamily, DbLong, DbInt, DbString, DbByte and DbBytes
final var cf = new DbLong();
final var dbLong = new DbLong();
final var dbInt = new DbInt();
final var dbString = new DbString();
final var dbByte = new DbByte();
final var dbBytes = new DbBytes();
cf.wrapLong(1);
dbLong.wrapLong(5);
dbInt.wrapInt(987);
dbString.wrapString("hello");
dbByte.wrapByte((byte) 123);
dbBytes.wrapBytes(new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10});

final var keyBuffer = new ExpandableArrayBuffer();
final var offset = new MutableInteger(0);
cf.write(keyBuffer, offset.getAndAdd(cf.getLength()));
dbLong.write(keyBuffer, offset.getAndAdd(dbLong.getLength()));
dbInt.write(keyBuffer, offset.getAndAdd(dbInt.getLength()));
dbString.write(keyBuffer, offset.getAndAdd(dbString.getLength()));
dbByte.write(keyBuffer, offset.getAndAdd(dbByte.getLength()));
dbBytes.write(keyBuffer, offset.getAndAdd(dbBytes.getLength()));
final var key = new byte[offset.get()];
keyBuffer.getBytes(0, key, 0, key.length);

// when
final var fullFormatter = KeyFormatters.ofFormat("lisbB").forColumnFamily(ZbColumnFamilies.DEFAULT);
final var partialFormatter = KeyFormatters.ofFormat("lis").forColumnFamily(ZbColumnFamilies.DEFAULT);

// then
assertThat(fullFormatter.formatKey(key)).isEqualTo("5:987:hello:123:01 02 03 04 05 06 07 08 09 0a");
assertThat(partialFormatter.formatKey(key)).isEqualTo("5:987:hello");
}

@Test
void shouldUseRegisteredFormatter() {
// given
final var cf = new DbLong();
final var dbString = new DbString();
cf.wrapLong(ZbColumnFamilies.DEFAULT.getValue());
dbString.wrapString("hello");

final var keyBuffer = new ExpandableArrayBuffer();
final var offset = new MutableInteger(0);
cf.write(keyBuffer, offset.getAndAdd(cf.getLength()));
dbString.write(keyBuffer, offset.getAndAdd(dbString.getLength()));

final var key = new byte[offset.get()];
keyBuffer.getBytes(0, key, 0, key.length);

// when
final var formatter = KeyFormatters.ofDefault().forColumnFamily(ZbColumnFamilies.DEFAULT);

// then
assertThat(formatter.formatKey(key)).isEqualTo("hello");
}

}

0 comments on commit 3d42cc0

Please sign in to comment.