feat(client): add Hrana 3 Protobuf over WebSocket transport

libsql-client/src/main/java/com/github/magnusp/libsql/client/LibsqlClientConfig.java

@@ -6,12 +6,27 @@ public record LibsqlClientConfig(
     String url,
     String authToken,
     Duration connectTimeout,
-    Duration busyTimeout
+    Duration busyTimeout,
+    Encoding encoding
 ) {
     public static final Duration DEFAULT_CONNECT_TIMEOUT = Duration.ofSeconds(10);
     public static final Duration DEFAULT_BUSY_TIMEOUT = Duration.ofSeconds(5);
 
+    /** The wire encoding used for Hrana 3 over WebSocket. HTTP always speaks JSON. */
+    public enum Encoding {
+        JSON,
+        PROTOBUF
+    }
+
+    public LibsqlClientConfig {
+        encoding = encoding == null ? Encoding.JSON : encoding;
+    }
+
+    public LibsqlClientConfig(String url, String authToken, Duration connectTimeout, Duration busyTimeout) {
+        this(url, authToken, connectTimeout, busyTimeout, Encoding.JSON);
+    }
+
     public LibsqlClientConfig(String url, String authToken) {
-        this(url, authToken, DEFAULT_CONNECT_TIMEOUT, DEFAULT_BUSY_TIMEOUT);
+        this(url, authToken, DEFAULT_CONNECT_TIMEOUT, DEFAULT_BUSY_TIMEOUT, Encoding.JSON);
     }
 }

libsql-client/src/main/java/com/github/magnusp/libsql/client/LibsqlWebSocketClient.java

@@ -40,7 +40,10 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
- * Hrana 3 client over a WebSocket, speaking the JSON encoding ({@code hrana3} subprotocol).
+ * Hrana 3 client over a WebSocket.
+ *
+ * <p>The supplied {@link HranaWsCodec} selects both the wire encoding and the subprotocol: {@code
+ * hrana3} for JSON text frames and {@code hrana3-protobuf} for Protobuf binary frames.
  *
  * <p>One WebSocket connection multiplexes any number of streams. Requests are matched to responses
  * by {@code request_id}, so the client may send a whole one-shot sequence (open, execute, close)
@@ -58,7 +61,8 @@
  */
 public class LibsqlWebSocketClient implements LibsqlClient {
 
-    private static final String SUBPROTOCOL = "hrana3";
+    private static final String SUBPROTOCOL_JSON = "hrana3";
+    private static final String SUBPROTOCOL_PROTOBUF = "hrana3-protobuf";
     private static final String WEBSOCKET_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 
     private static final int OP_CONTINUATION = 0x0;
@@ -73,6 +77,7 @@ public class LibsqlWebSocketClient implements LibsqlClient {
 
     private final LibsqlClientConfig config;
     private final HranaWsCodec codec;
+    private final String subprotocol;
     private final Duration timeout;
     private final Transport transport;
 
@@ -97,7 +102,8 @@ public LibsqlWebSocketClient(LibsqlClientConfig config, HranaWsCodec codec) {
         this.timeout = config.connectTimeout() != null
                 ? config.connectTimeout()
                 : LibsqlClientConfig.DEFAULT_CONNECT_TIMEOUT;
-        this.transport = new Transport(URI.create(config.url()), timeout);
+        this.subprotocol = codec.isBinary() ? SUBPROTOCOL_PROTOBUF : SUBPROTOCOL_JSON;
+        this.transport = new Transport(URI.create(config.url()), timeout, subprotocol);
 
         try {
             send(new WsClientMsg.Hello(authToken()));
@@ -359,14 +365,14 @@ private final class Transport {
 
         private volatile boolean open = true;
 
-        private Transport(URI uri, Duration timeout) {
+        private Transport(URI uri, Duration timeout, String subprotocol) {
             Socket connected = null;
             try {
                 connected = connect(uri, timeout);
                 this.socket = connected;
                 this.in = connected.getInputStream();
                 this.out = connected.getOutputStream();
-                handshake(uri, in, out, timeout);
+                handshake(uri, in, out, timeout, subprotocol);
             } catch (IOException e) {
                 closeQuietly(connected);
                 throw new LibsqlException(
@@ -583,8 +589,8 @@ private static Socket connect(URI uri, Duration timeout) throws IOException {
         }
     }
 
-    private static void handshake(URI uri, InputStream in, OutputStream out, Duration timeout)
-            throws IOException {
+    private static void handshake(URI uri, InputStream in, OutputStream out, Duration timeout,
+            String subprotocol) throws IOException {
         String key = createKey();
         String expectedAccept = acceptKey(key);
         String path = uri.getRawPath() == null || uri.getRawPath().isEmpty() ? "/" : uri.getRawPath();
@@ -602,7 +608,7 @@ private static void handshake(URI uri, InputStream in, OutputStream out, Duratio
                 .append("Connection: Upgrade\r\n")
                 .append("Sec-WebSocket-Key: ").append(key).append("\r\n")
                 .append("Sec-WebSocket-Version: 13\r\n")
-                .append("Sec-WebSocket-Protocol: ").append(SUBPROTOCOL).append("\r\n")
+                .append("Sec-WebSocket-Protocol: ").append(subprotocol).append("\r\n")
                 .append("\r\n");
         out.write(request.toString().getBytes(StandardCharsets.US_ASCII));
         out.flush();
@@ -627,8 +633,8 @@ private static void handshake(URI uri, InputStream in, OutputStream out, Duratio
             throw new IOException("WebSocket upgrade response has an invalid Sec-WebSocket-Accept");
         }
         String negotiated = headers.get("sec-websocket-protocol");
-        if (!SUBPROTOCOL.equals(negotiated)) {
-            throw new IOException("Server did not negotiate subprotocol " + SUBPROTOCOL
+        if (!subprotocol.equals(negotiated)) {
+            throw new IOException("Server did not negotiate subprotocol " + subprotocol
                     + " (got " + negotiated + ")");
         }
     }

libsql-hrana-codec/src/main/java/com/github/magnusp/libsql/hrana/codec/HranaWsProtobufCodec.java

@@ -0,0 +1,237 @@
+package com.github.magnusp.libsql.hrana.codec;
+
+import com.github.magnusp.libsql.hrana.proto.HranaProto;
+import com.github.magnusp.libsql.hrana.proto.HranaWsProto;
+import com.google.protobuf.ByteString;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Hrana 3 WebSocket message codec for the Protobuf encoding ({@code hrana3-protobuf}
+ * subprotocol).
+ *
+ * <p>Every client or server message is a complete {@code hrana.ws.ClientMsg} or {@code
+ * hrana.ws.ServerMsg} serialized as one binary WebSocket frame.
+ */
+public final class HranaWsProtobufCodec implements HranaWsCodec {
+
+    @Override
+    public boolean isBinary() {
+        return true;
+    }
+
+    @Override
+    public byte[] serialize(WsClientMsg msg) throws IOException {
+        HranaWsProto.ClientMsg.Builder clientMsg = HranaWsProto.ClientMsg.newBuilder();
+        switch (msg) {
+            case WsClientMsg.Hello hello -> {
+                HranaWsProto.HelloMsg.Builder helloMsg = HranaWsProto.HelloMsg.newBuilder();
+                if (hello.jwt() != null) {
+                    helloMsg.setJwt(hello.jwt());
+                }
+                clientMsg.setHello(helloMsg);
+            }
+            case WsClientMsg.Request request -> {
+                HranaWsProto.RequestMsg.Builder requestMsg = HranaWsProto.RequestMsg.newBuilder()
+                        .setRequestId(request.requestId());
+                switch (request.request()) {
+                    case WsRequest.OpenStream open -> requestMsg.setOpenStream(
+                            HranaWsProto.OpenStreamReq.newBuilder()
+                                    .setStreamId(open.streamId()));
+                    case WsRequest.CloseStream close -> requestMsg.setCloseStream(
+                            HranaWsProto.CloseStreamReq.newBuilder()
+                                    .setStreamId(close.streamId()));
+                    case WsRequest.Execute execute -> requestMsg.setExecute(
+                            HranaWsProto.ExecuteReq.newBuilder()
+                                    .setStreamId(execute.streamId())
+                                    .setStmt(toProto(execute.stmt())));
+                    case WsRequest.Batch batch -> requestMsg.setBatch(
+                            HranaWsProto.BatchReq.newBuilder()
+                                    .setStreamId(batch.streamId())
+                                    .setBatch(toProto(batch.batch())));
+                }
+                clientMsg.setRequest(requestMsg);
+            }
+        }
+        return clientMsg.build().toByteArray();
+    }
+
+    @Override
+    public WsServerMsg deserialize(byte[] payload) throws IOException {
+        HranaWsProto.ServerMsg serverMsg = HranaWsProto.ServerMsg.parseFrom(payload);
+        return switch (serverMsg.getMsgCase()) {
+            case HELLO_OK -> WsServerMsg.HelloOk.INSTANCE;
+            case HELLO_ERROR -> new WsServerMsg.HelloError(
+                    fromProto(serverMsg.getHelloError().getError()));
+            case RESPONSE_OK -> new WsServerMsg.ResponseOk(
+                    serverMsg.getResponseOk().getRequestId(),
+                    fromProto(serverMsg.getResponseOk()));
+            case RESPONSE_ERROR -> new WsServerMsg.ResponseError(
+                    serverMsg.getResponseError().getRequestId(),
+                    fromProto(serverMsg.getResponseError().getError()));
+            case MSG_NOT_SET -> throw new IOException(
+                    "Hrana protobuf server message has no msg case");
+        };
+    }
+
+    // -------------------------------------------------------------
+    // Client structures
+    // -------------------------------------------------------------
+
+    private static HranaProto.Stmt toProto(Stmt stmt) {
+        HranaProto.Stmt.Builder out = HranaProto.Stmt.newBuilder();
+        if (stmt.sql() != null) {
+            out.setSql(stmt.sql());
+        }
+        if (stmt.sqlId() != null) {
+            out.setSqlId(stmt.sqlId());
+        }
+        if (stmt.args() != null) {
+            for (Value arg : stmt.args()) {
+                out.addArgs(toProto(arg));
+            }
+        }
+        if (stmt.namedArgs() != null) {
+            for (NamedArg namedArg : stmt.namedArgs()) {
+                out.addNamedArgs(HranaProto.NamedArg.newBuilder()
+                        .setName(namedArg.name())
+                        .setValue(toProto(namedArg.value())));
+            }
+        }
+        if (stmt.wantRows() != null) {
+            out.setWantRows(stmt.wantRows());
+        }
+        return out.build();
+    }
+
+    private static HranaProto.Value toProto(Value value) {
+        HranaProto.Value.Builder out = HranaProto.Value.newBuilder();
+        switch (value) {
+            case Value.Null ignored -> out.setNull(HranaProto.Value.Null.getDefaultInstance());
+            case Value.Integer integer -> out.setInteger(integer.value());
+            case Value.Float floating -> out.setFloat(floating.value());
+            case Value.Text text -> out.setText(text.value());
+            case Value.Blob blob -> out.setBlob(ByteString.copyFrom(blob.bytes()));
+        }
+        return out.build();
+    }
+
+    private static HranaProto.Batch toProto(Batch batch) {
+        HranaProto.Batch.Builder out = HranaProto.Batch.newBuilder();
+        for (BatchStep step : batch.steps()) {
+            HranaProto.BatchStep.Builder stepMsg = HranaProto.BatchStep.newBuilder()
+                    .setStmt(toProto(step.stmt()));
+            if (step.condition() != null) {
+                stepMsg.setCondition(toProto(step.condition()));
+            }
+            out.addSteps(stepMsg);
+        }
+        return out.build();
+    }
+
+    private static HranaProto.BatchCond toProto(BatchCond condition) {
+        HranaProto.BatchCond.Builder out = HranaProto.BatchCond.newBuilder();
+        switch (condition) {
+            case BatchCond.StepOk stepOk -> out.setStepOk(stepOk.step());
+            case BatchCond.StepError stepError -> out.setStepError(stepError.step());
+            case BatchCond.Not not -> out.setNot(toProto(not.cond()));
+            case BatchCond.And and -> out.setAnd(toProtoCondList(and.conds()));
+            case BatchCond.Or or -> out.setOr(toProtoCondList(or.conds()));
+            case BatchCond.IsAutocommit ignored -> out.setIsAutocommit(
+                    HranaProto.BatchCond.IsAutocommit.getDefaultInstance());
+        }
+        return out.build();
+    }
+
+    private static HranaProto.BatchCond.CondList toProtoCondList(List<BatchCond> conditions) {
+        HranaProto.BatchCond.CondList.Builder out = HranaProto.BatchCond.CondList.newBuilder();
+        for (BatchCond condition : conditions) {
+            out.addConds(toProto(condition));
+        }
+        return out.build();
+    }
+
+    // -------------------------------------------------------------
+    // Server structures
+    // -------------------------------------------------------------
+
+    private static WsResponse fromProto(HranaWsProto.ResponseOkMsg responseMsg) throws IOException {
+        return switch (responseMsg.getResponseCase()) {
+            case OPEN_STREAM -> WsResponse.OpenStream.INSTANCE;
+            case CLOSE_STREAM -> WsResponse.CloseStream.INSTANCE;
+            case EXECUTE -> new WsResponse.Execute(
+                    fromProto(responseMsg.getExecute().getResult()));
+            case BATCH -> new WsResponse.Batch(
+                    fromProto(responseMsg.getBatch().getResult()));
+            case RESPONSE_NOT_SET -> throw new IOException(
+                    "Hrana protobuf response has no response case");
+            default -> throw new IOException("Unsupported Hrana protobuf response: "
+                    + responseMsg.getResponseCase());
+        };
+    }
+
+    private static StmtResult fromProto(HranaProto.StmtResult result) throws IOException {
+        List<Col> cols = new ArrayList<>(result.getColsCount());
+        for (HranaProto.Col col : result.getColsList()) {
+            cols.add(fromProto(col));
+        }
+        List<Row> rows = new ArrayList<>(result.getRowsCount());
+        for (HranaProto.Row row : result.getRowsList()) {
+            rows.add(fromProto(row));
+        }
+        Long lastInsertRowid = result.hasLastInsertRowid()
+                ? result.getLastInsertRowid()
+                : null;
+        return new StmtResult(cols, rows, result.getAffectedRowCount(), lastInsertRowid);
+    }
+
+    private static Col fromProto(HranaProto.Col col) {
+        return new Col(
+                col.hasName() ? col.getName() : null,
+                col.hasDecltype() ? col.getDecltype() : null);
+    }
+
+    private static Row fromProto(HranaProto.Row row) throws IOException {
+        List<Value> values = new ArrayList<>(row.getValuesCount());
+        for (HranaProto.Value value : row.getValuesList()) {
+            values.add(fromProto(value));
+        }
+        return new Row(values);
+    }
+
+    private static BatchResult fromProto(HranaProto.BatchResult result) throws IOException {
+        Map<Integer, StmtResult> stepResults = new LinkedHashMap<>();
+        for (Map.Entry<Integer, HranaProto.StmtResult> entry
+                : result.getStepResultsMap().entrySet()) {
+            stepResults.put(entry.getKey(), fromProto(entry.getValue()));
+        }
+        Map<Integer, HranaError> stepErrors = new LinkedHashMap<>();
+        for (Map.Entry<Integer, HranaProto.Error> entry
+                : result.getStepErrorsMap().entrySet()) {
+            stepErrors.put(entry.getKey(), fromProto(entry.getValue()));
+        }
+        return new BatchResult(stepResults, stepErrors);
+    }
+
+    private static HranaError fromProto(HranaProto.Error error) {
+        return new HranaError(
+                error.getMessage(),
+                error.hasCode() ? error.getCode() : null);
+    }
+
+    private static Value fromProto(HranaProto.Value value) throws IOException {
+        return switch (value.getValueCase()) {
+            case NULL -> Value.Null.INSTANCE;
+            case INTEGER -> new Value.Integer(value.getInteger());
+            case FLOAT -> new Value.Float(value.getFloat());
+            case TEXT -> new Value.Text(value.getText());
+            case BLOB -> new Value.Blob(value.getBlob().toByteArray());
+            case VALUE_NOT_SET -> throw new IOException(
+                    "Hrana protobuf Value has no value case");
+        };
+    }
+}