Merge remote-tracking branch 'origin/jh_experiment_new_ilp' into jh_experiment_new_ilp_udp

Author: jerrinot

core/src/main/java/io/questdb/client/Sender.java

@@ -55,7 +55,9 @@
 import java.io.Closeable;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
+import java.nio.charset.StandardCharsets;
 import java.security.PrivateKey;
+import java.util.Base64;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.concurrent.TimeUnit;
@@ -884,6 +886,8 @@ public Sender build() {
                         : TimeUnit.MILLISECONDS.toNanos(autoFlushIntervalMillis);
                 int actualInFlightWindowSize = inFlightWindowSize == PARAMETER_NOT_SET_EXPLICITLY ? DEFAULT_IN_FLIGHT_WINDOW_SIZE : inFlightWindowSize;
 
+                String wsAuthHeader = buildWebSocketAuthHeader();
+
                 if (asyncMode) {
                     return QwpWebSocketSender.connectAsync(
                             hosts.getQuick(0),
@@ -892,7 +896,8 @@ public Sender build() {
                             actualAutoFlushRows,
                             actualAutoFlushBytes,
                             actualAutoFlushIntervalNanos,
-                            actualInFlightWindowSize
+                            actualInFlightWindowSize,
+                            wsAuthHeader
                     );
                 } else {
                     return QwpWebSocketSender.connect(
@@ -901,7 +906,8 @@ public Sender build() {
                             tlsEnabled,
                             actualAutoFlushRows,
                             actualAutoFlushBytes,
-                            actualAutoFlushIntervalNanos
+                            actualAutoFlushIntervalNanos,
+                            wsAuthHeader
                     );
                 }
             }
@@ -1666,15 +1672,15 @@ private LineSenderBuilder fromConfig(CharSequence configurationString) {
                     }
                 } else if (Chars.equals("token", sink)) {
                     pos = getValue(configurationString, pos, sink, "token");
-                    if (protocol == PROTOCOL_TCP) {
+                    if (protocol == PROTOCOL_WEBSOCKET) {
+                        throw new LineSenderException("token is not supported for WebSocket protocol");
+                    } else if (protocol == PROTOCOL_TCP) {
                         tcpToken = sink.toString();
                         // will configure later, we need to know a keyId first
-                    } else if (protocol == PROTOCOL_HTTP) {
-                        httpToken(sink.toString());
                     } else if (protocol == PROTOCOL_UDP) {
                         throw new LineSenderException("token is not supported for UDP transport");
                     } else {
-                        throw new LineSenderException("token is not supported for WebSocket protocol");
+                        httpToken(sink.toString());
                     }
                 } else if (Chars.equals("retry_timeout", sink)) {
                     pos = getValue(configurationString, pos, sink, "retry_timeout");
@@ -1788,11 +1794,11 @@ private LineSenderBuilder fromConfig(CharSequence configurationString) {
             } else if (trustStorePassword != null) {
                 throw new LineSenderException("tls_roots_password was configured, but tls_roots is missing");
             }
-            if (protocol == PROTOCOL_HTTP) {
+            if (protocol == PROTOCOL_HTTP || protocol == PROTOCOL_WEBSOCKET) {
                 if (user != null) {
                     httpUsernamePassword(user, password);
                 } else if (password != null) {
-                    throw new LineSenderException("HTTP password is configured, but username is missing");
+                    throw new LineSenderException("password is configured, but username is missing");
                 }
             } else {
                 if (user != null) {
@@ -1935,9 +1941,8 @@ private void validateParameters() {
                 if (privateKey != null) {
                     throw new LineSenderException("TCP authentication is not supported for WebSocket protocol");
                 }
-                if (httpToken != null || username != null || password != null) {
-                    // TODO: WebSocket auth not yet implemented
-                    throw new LineSenderException("Authentication is not yet supported for WebSocket protocol");
+                if (httpToken != null && (username != null || password != null)) {
+                    throw new LineSenderException("cannot use both token and username/password authentication");
                 }
                 if (inFlightWindowSize != PARAMETER_NOT_SET_EXPLICITLY && !asyncMode) {
                     throw new LineSenderException("in-flight window size requires async mode");
@@ -1969,6 +1974,17 @@ private void validateParameters() {
             }
         }
 
+        private String buildWebSocketAuthHeader() {
+            if (username != null && password != null) {
+                String credentials = username + ":" + password;
+                return "Basic " + Base64.getEncoder().encodeToString(credentials.getBytes(StandardCharsets.UTF_8));
+            }
+            if (httpToken != null) {
+                return "Bearer " + httpToken;
+            }
+            return null;
+        }
+
         private void udp() {
             if (protocol != PARAMETER_NOT_SET_EXPLICITLY) {
                 throw new LineSenderException("protocol was already configured ")

core/src/main/java/io/questdb/client/cutlass/http/client/WebSocketClient.java

@@ -383,10 +383,11 @@ public boolean tryReceiveFrame(WebSocketFrameHandler handler) {
     /**
      * Performs WebSocket upgrade handshake.
      *
-     * @param path    the WebSocket endpoint path (e.g., "/ws")
-     * @param timeout timeout in milliseconds
+     * @param path                the WebSocket endpoint path (e.g., "/ws")
+     * @param timeout             timeout in milliseconds
+     * @param authorizationHeader the Authorization header value (e.g., "Basic ..."), or null
      */
-    public void upgrade(CharSequence path, int timeout) {
+    public void upgrade(CharSequence path, int timeout, CharSequence authorizationHeader) {
         if (closed) {
             throw new HttpClientException("WebSocket client is closed");
         }
@@ -422,6 +423,11 @@ public void upgrade(CharSequence path, int timeout) {
         sendBuffer.putAscii(handshakeKey);
         sendBuffer.putAscii("\r\n");
         sendBuffer.putAscii("Sec-WebSocket-Version: 13\r\n");
+        if (authorizationHeader != null) {
+            sendBuffer.putAscii("Authorization: ");
+            sendBuffer.putAscii(authorizationHeader);
+            sendBuffer.putAscii("\r\n");
+        }
         sendBuffer.putAscii("\r\n");
 
         // Send request
@@ -441,7 +447,21 @@ public void upgrade(CharSequence path, int timeout) {
      * Performs upgrade with default timeout.
      */
     public void upgrade(CharSequence path) {
-        upgrade(path, defaultTimeout);
+        upgrade(path, defaultTimeout, null);
+    }
+
+    /**
+     * Performs upgrade with default timeout and authorization header.
+     */
+    public void upgrade(CharSequence path, CharSequence authorizationHeader) {
+        upgrade(path, defaultTimeout, authorizationHeader);
+    }
+
+    /**
+     * Performs upgrade without authorization header.
+     */
+    public void upgrade(CharSequence path, int timeout) {
+        upgrade(path, timeout, null);
     }
 
     private static String computeAcceptKey(String key) {

core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpWebSocketSender.java

@@ -117,6 +117,7 @@ public class QwpWebSocketSender implements Sender {
     private static final String WRITE_PATH = "/write/v4";
     private final AckFrameHandler ackHandler = new AckFrameHandler(this);
     private final WebSocketResponse ackResponse = new WebSocketResponse();
+    private final String authorizationHeader;
     private final int autoFlushBytes;
     private final long autoFlushIntervalNanos;
     // Auto-flush configuration
@@ -174,8 +175,10 @@ private QwpWebSocketSender(
             int autoFlushRows,
             int autoFlushBytes,
             long autoFlushIntervalNanos,
-            int inFlightWindowSize
+            int inFlightWindowSize,
+            String authorizationHeader
     ) {
+        this.authorizationHeader = authorizationHeader;
         this.host = host;
         this.port = port;
         this.tlsEnabled = tlsEnabled;
@@ -224,8 +227,9 @@ public static QwpWebSocketSender connect(String host, int port) {
      * @return connected sender
      */
     public static QwpWebSocketSender connect(String host, int port, boolean tlsEnabled) {
-        return connect(host, port, tlsEnabled,
-                DEFAULT_AUTO_FLUSH_ROWS, DEFAULT_AUTO_FLUSH_BYTES, DEFAULT_AUTO_FLUSH_INTERVAL_NANOS);
+        return connect(
+                host, port, tlsEnabled, DEFAULT_AUTO_FLUSH_ROWS, DEFAULT_AUTO_FLUSH_BYTES, DEFAULT_AUTO_FLUSH_INTERVAL_NANOS
+        );
     }
 
     /**
@@ -240,13 +244,30 @@ public static QwpWebSocketSender connect(String host, int port, boolean tlsEnabl
      * @param autoFlushIntervalNanos age before flush in nanos (0 = no limit)
      * @return connected sender
      */
-    public static QwpWebSocketSender connect(String host, int port, boolean tlsEnabled,
-                                             int autoFlushRows, int autoFlushBytes,
-                                             long autoFlushIntervalNanos) {
+    public static QwpWebSocketSender connect(
+            String host,
+            int port,
+            boolean tlsEnabled,
+            int autoFlushRows,
+            int autoFlushBytes,
+            long autoFlushIntervalNanos
+    ) {
+        return connect(host, port, tlsEnabled, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos, null);
+    }
+
+    public static QwpWebSocketSender connect(
+            String host,
+            int port,
+            boolean tlsEnabled,
+            int autoFlushRows,
+            int autoFlushBytes,
+            long autoFlushIntervalNanos,
+            String authorizationHeader
+    ) {
         QwpWebSocketSender sender = new QwpWebSocketSender(
-                host, port, tlsEnabled, DEFAULT_BUFFER_SIZE,
-                autoFlushRows, autoFlushBytes, autoFlushIntervalNanos,
-                1   // window=1 for sync behavior
+                host, port, tlsEnabled, DEFAULT_BUFFER_SIZE, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos,
+                1,   // window=1 for sync behavior
+                authorizationHeader
         );
         sender.ensureConnected();
         return sender;
@@ -263,11 +284,17 @@ public static QwpWebSocketSender connect(String host, int port, boolean tlsEnabl
      * @param autoFlushIntervalNanos age before flush in nanos (0 = no limit)
      * @return connected sender
      */
-    public static QwpWebSocketSender connectAsync(String host, int port, boolean tlsEnabled,
-                                                  int autoFlushRows, int autoFlushBytes,
-                                                  long autoFlushIntervalNanos) {
-        return connectAsync(host, port, tlsEnabled, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos,
-                DEFAULT_IN_FLIGHT_WINDOW_SIZE);
+    public static QwpWebSocketSender connectAsync(
+            String host,
+            int port,
+            boolean tlsEnabled,
+            int autoFlushRows,
+            int autoFlushBytes,
+            long autoFlushIntervalNanos
+    ) {
+        return connectAsync(
+                host, port, tlsEnabled, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos, DEFAULT_IN_FLIGHT_WINDOW_SIZE
+        );
     }
 
     /**
@@ -290,11 +317,24 @@ public static QwpWebSocketSender connectAsync(
             int autoFlushBytes,
             long autoFlushIntervalNanos,
             int inFlightWindowSize
+    ) {
+        return connectAsync(
+                host, port, tlsEnabled, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos, inFlightWindowSize, null
+        );
+    }
+
+    public static QwpWebSocketSender connectAsync(
+            String host,
+            int port,
+            boolean tlsEnabled,
+            int autoFlushRows,
+            int autoFlushBytes,
+            long autoFlushIntervalNanos,
+            int inFlightWindowSize,
+            String authorizationHeader
     ) {
         QwpWebSocketSender sender = new QwpWebSocketSender(
-                host, port, tlsEnabled, DEFAULT_BUFFER_SIZE,
-                autoFlushRows, autoFlushBytes, autoFlushIntervalNanos,
-                inFlightWindowSize
+                host, port, tlsEnabled, DEFAULT_BUFFER_SIZE, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos, inFlightWindowSize, authorizationHeader
         );
         sender.ensureConnected();
         return sender;
@@ -309,8 +349,9 @@ public static QwpWebSocketSender connectAsync(
      * @return connected sender
      */
     public static QwpWebSocketSender connectAsync(String host, int port, boolean tlsEnabled) {
-        return connectAsync(host, port, tlsEnabled,
-                DEFAULT_AUTO_FLUSH_ROWS, DEFAULT_AUTO_FLUSH_BYTES, DEFAULT_AUTO_FLUSH_INTERVAL_NANOS);
+        return connectAsync(
+                host, port, tlsEnabled, DEFAULT_AUTO_FLUSH_ROWS, DEFAULT_AUTO_FLUSH_BYTES, DEFAULT_AUTO_FLUSH_INTERVAL_NANOS
+        );
     }
 
     /**
@@ -326,9 +367,7 @@ public static QwpWebSocketSender connectAsync(String host, int port, boolean tls
      */
     public static QwpWebSocketSender createForTesting(String host, int port, int inFlightWindowSize) {
         return new QwpWebSocketSender(
-                host, port, false, DEFAULT_BUFFER_SIZE,
-                DEFAULT_AUTO_FLUSH_ROWS, DEFAULT_AUTO_FLUSH_BYTES, DEFAULT_AUTO_FLUSH_INTERVAL_NANOS,
-                inFlightWindowSize
+                host, port, false, DEFAULT_BUFFER_SIZE, DEFAULT_AUTO_FLUSH_ROWS, DEFAULT_AUTO_FLUSH_BYTES, DEFAULT_AUTO_FLUSH_INTERVAL_NANOS, inFlightWindowSize, null
         );
         // Note: does NOT call ensureConnected()
     }
@@ -345,13 +384,15 @@ public static QwpWebSocketSender createForTesting(String host, int port, int inF
      * @return unconnected sender
      */
     public static QwpWebSocketSender createForTesting(
-            String host, int port,
-            int autoFlushRows, int autoFlushBytes, long autoFlushIntervalNanos,
-            int inFlightWindowSize) {
+            String host,
+            int port,
+            int autoFlushRows,
+            int autoFlushBytes,
+            long autoFlushIntervalNanos,
+            int inFlightWindowSize
+    ) {
         return new QwpWebSocketSender(
-                host, port, false, DEFAULT_BUFFER_SIZE,
-                autoFlushRows, autoFlushBytes, autoFlushIntervalNanos,
-                inFlightWindowSize
+                host, port, false, DEFAULT_BUFFER_SIZE, autoFlushRows, autoFlushBytes, autoFlushIntervalNanos, inFlightWindowSize, null
         );
         // Note: does NOT call ensureConnected()
     }
@@ -1012,7 +1053,7 @@ private void ensureConnected() {
             // Connect and upgrade to WebSocket
             try {
                 client.connect(host, port);
-                client.upgrade(WRITE_PATH);
+                client.upgrade(WRITE_PATH, authorizationHeader);
             } catch (Exception e) {
                 client.close();
                 client = null;

core/src/test/java/io/questdb/client/test/cutlass/line/LineSenderBuilderTest.java

@@ -27,7 +27,6 @@
 import io.questdb.client.Sender;
 import io.questdb.client.cutlass.line.LineSenderException;
 import io.questdb.client.test.tools.TestUtils;
-import static io.questdb.client.test.tools.TestUtils.assertMemoryLeak;
 import org.junit.Test;
 
 import static io.questdb.client.test.tools.TestUtils.assertMemoryLeak;
@@ -151,8 +150,8 @@ public void testConfStringValidation() throws Exception {
             assertConfStrError("tcp::addr=localhost;token=foo;", "TCP token is configured, but user is missing");
             assertConfStrError("http::addr=localhost;user=foo;", "password cannot be empty nor null");
             assertConfStrError("http::addr=localhost;username=foo;", "password cannot be empty nor null");
-            assertConfStrError("http::addr=localhost;pass=foo;", "HTTP password is configured, but username is missing");
-            assertConfStrError("http::addr=localhost;password=foo;", "HTTP password is configured, but username is missing");
+            assertConfStrError("http::addr=localhost;pass=foo;", "password is configured, but username is missing");
+            assertConfStrError("http::addr=localhost;password=foo;", "password is configured, but username is missing");
             assertConfStrError("tcp::addr=localhost;pass=foo;", "password is not supported for TCP protocol");
             assertConfStrError("tcp::addr=localhost;password=foo;", "password is not supported for TCP protocol");
             assertConfStrError("tcp::addr=localhost;retry_timeout=;", "retry_timeout cannot be empty");