Rework beacon connections and flesh out API more

beacon/src/main/java/io/xpipe/beacon/BeaconClient.java

@@ -68,6 +68,10 @@ public class BeaconClient implements AutoCloseable {
         out = socket.getOutputStream();
     }
 
+    public boolean isClosed() {
+        return socket.isClosed();
+    }
+
     public void close() throws ConnectorException {
         try {
             socket.close();
@@ -100,13 +104,32 @@ public class BeaconClient implements AutoCloseable {
         }
     }
 
+    public void receiveBody() throws ConnectorException {
+        try {
+            var sep = in.readNBytes(BODY_SEPARATOR.length);
+            if (sep.length != 0 && !Arrays.equals(BODY_SEPARATOR, sep)) {
+                throw new ConnectorException("Invalid body separator");
+            }
+        } catch (IOException ex) {
+            throw new ConnectorException(ex);
+        }
+    }
+
+    public void startBody() throws ConnectorException {
+        try {
+            out.write(BODY_SEPARATOR);
+        } catch (IOException ex) {
+            throw new ConnectorException(ex);
+        }
+    }
+
     public <REQ extends RequestMessage, RES extends ResponseMessage> RES simpleExchange(REQ req)
         throws ServerException, ConnectorException, ClientException {
         sendRequest(req);
         return this.receiveResponse();
     }
 
-    private <T extends RequestMessage> void sendRequest(T req) throws ClientException, ConnectorException {
+    public  <T extends RequestMessage> void sendRequest(T req) throws ClientException, ConnectorException {
         ObjectNode json = JacksonHelper.newMapper().valueToTree(req);
         var prov = MessageExchanges.byRequest(req);
         if (prov.isEmpty()) {
@@ -132,7 +155,7 @@ public class BeaconClient implements AutoCloseable {
         }
     }
 
-    private <T extends ResponseMessage> T receiveResponse() throws ConnectorException, ClientException, ServerException {
+    public <T extends ResponseMessage> T receiveResponse() throws ConnectorException, ClientException, ServerException {
         JsonNode read;
         try {
             var in = socket.getInputStream();

beacon/src/main/java/io/xpipe/beacon/BeaconConnection.java

@@ -0,0 +1,163 @@
+package io.xpipe.beacon;
+
+import io.xpipe.beacon.message.RequestMessage;
+import io.xpipe.beacon.message.ResponseMessage;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+public abstract class BeaconConnection implements AutoCloseable {
+
+    protected BeaconClient socket;
+
+    protected abstract void constructSocket();
+
+    @Override
+    public void close() {
+        try {
+            if (socket != null) {
+                socket.close();
+            }
+            socket = null;
+        } catch (Exception e) {
+            socket = null;
+            throw new BeaconException("Could not close beacon connection", e);
+        }
+    }
+
+    public void closeOutput() {
+        try {
+            socket.getOutputStream().close();
+        } catch (Exception e) {
+            throw new BeaconException("Could not close beacon output stream", e);
+        }
+    }
+
+    public void withOutputStream(BeaconClient.FailableConsumer<OutputStream, IOException> ex) {
+        try {
+            ex.accept(getOutputStream());
+        } catch (IOException e) {
+            throw new BeaconException("Could not write to beacon output stream", e);
+        }
+    }
+
+    public void withInputStream(BeaconClient.FailableConsumer<InputStream, IOException> ex) {
+        try {
+            ex.accept(getInputStream());
+        } catch (IOException e) {
+            throw new BeaconException("Could not read from beacon output stream", e);
+        }
+    }
+
+    public void checkClosed() {
+        if (socket == null) {
+            throw new BeaconException("Socket is closed");
+        }
+    }
+
+    public OutputStream getOutputStream() {
+        checkClosed();
+
+        return socket.getOutputStream();
+    }
+
+    public InputStream getInputStream() {
+        checkClosed();
+
+        return socket.getInputStream();
+    }
+
+    public <REQ extends RequestMessage, RES extends ResponseMessage> void performInputExchange(
+            REQ req,
+            BeaconClient.FailableBiConsumer<RES, InputStream, IOException> responseConsumer) {
+        checkClosed();
+
+        performInputOutputExchange(req, null, responseConsumer);
+    }
+
+    public <REQ extends RequestMessage, RES extends ResponseMessage> void performInputOutputExchange(
+            REQ req,
+            BeaconClient.FailableConsumer<OutputStream, IOException> reqWriter,
+            BeaconClient.FailableBiConsumer<RES, InputStream, IOException> responseConsumer) {
+        checkClosed();
+
+        try {
+            socket.exchange(req, reqWriter, responseConsumer);
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+
+    public <REQ extends RequestMessage> void sendRequest(
+            REQ req) {
+        checkClosed();
+
+        try {
+            socket.sendRequest(req);
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+
+    public <RES extends ResponseMessage> RES receiveResponse() {
+        checkClosed();
+
+        try {
+            return socket.receiveResponse();
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+
+    public void sendBodyStart() {
+        checkClosed();
+
+        try {
+            socket.startBody();
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+
+    public void receiveBody() {
+        checkClosed();
+
+        try {
+            socket.receiveBody();
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+
+    public <REQ extends RequestMessage, RES extends ResponseMessage> RES performOutputExchange(
+            REQ req,
+            BeaconClient.FailableConsumer<OutputStream, IOException> reqWriter) {
+        checkClosed();
+
+        try {
+            socket.sendRequest(req);
+            socket.startBody();
+            reqWriter.accept(socket.getOutputStream());
+            return socket.receiveResponse();
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+
+//    public void writeLength(int bytes) throws IOException {
+//        checkClosed();
+//        socket.getOutputStream().write(ByteBuffer.allocate(4).putInt(bytes).array());
+//    }
+
+    public <REQ extends RequestMessage, RES extends ResponseMessage> RES performSimpleExchange(
+            REQ req) {
+        checkClosed();
+
+        try {
+            return socket.simpleExchange(req);
+        } catch (Exception e) {
+            throw new BeaconException("Could not communicate with beacon", e);
+        }
+    }
+}

beacon/src/main/java/io/xpipe/beacon/BeaconConnector.java

@@ -1,53 +0,0 @@
-package io.xpipe.beacon;
-
-import io.xpipe.beacon.message.RequestMessage;
-import io.xpipe.beacon.message.ResponseMessage;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.nio.ByteBuffer;
-import java.util.concurrent.atomic.AtomicReference;
-
-public abstract class BeaconConnector {
-
-    protected abstract BeaconClient constructSocket() throws ConnectorException;
-
-    protected <REQ extends RequestMessage, RES extends ResponseMessage> void performInputExchange(
-            BeaconClient socket,
-            REQ req,
-            BeaconClient.FailableBiConsumer<RES, InputStream, IOException> responseConsumer) throws ServerException, ConnectorException, ClientException {
-        performInputOutputExchange(socket, req, null, responseConsumer);
-    }
-
-    protected <REQ extends RequestMessage, RES extends ResponseMessage> void performInputOutputExchange(
-            BeaconClient socket,
-            REQ req,
-            BeaconClient.FailableConsumer<OutputStream, IOException> reqWriter,
-            BeaconClient.FailableBiConsumer<RES, InputStream, IOException> responseConsumer)
-            throws ServerException, ConnectorException, ClientException {
-        socket.exchange(req, reqWriter, responseConsumer);
-    }
-
-    protected <REQ extends RequestMessage, RES extends ResponseMessage> RES performOutputExchange(
-            BeaconClient socket,
-            REQ req,
-            BeaconClient.FailableConsumer<OutputStream, IOException> reqWriter)
-            throws ServerException, ConnectorException, ClientException {
-        AtomicReference<RES> response = new AtomicReference<>();
-        socket.exchange(req, reqWriter, (RES res, InputStream in) -> {
-            response.set(res);
-        });
-        return response.get();
-    }
-
-    protected void writeLength(BeaconClient socket, int bytes) throws IOException {
-        socket.getOutputStream().write(ByteBuffer.allocate(4).putInt(bytes).array());
-    }
-
-    protected <REQ extends RequestMessage, RES extends ResponseMessage> RES performSimpleExchange(
-            BeaconClient socket,
-            REQ req) throws ServerException, ConnectorException, ClientException {
-        return socket.simpleExchange(req);
-    }
-}

beacon/src/main/java/io/xpipe/beacon/BeaconException.java

@@ -0,0 +1,23 @@
+package io.xpipe.beacon;
+
+public class BeaconException extends RuntimeException {
+
+    public BeaconException() {
+    }
+
+    public BeaconException(String message) {
+        super(message);
+    }
+
+    public BeaconException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public BeaconException(Throwable cause) {
+        super(cause);
+    }
+
+    public BeaconException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+}

beacon/src/main/java/io/xpipe/beacon/exchange/PreStoreExchange.java

@@ -2,13 +2,11 @@ package io.xpipe.beacon.exchange;
 
 import io.xpipe.beacon.message.RequestMessage;
 import io.xpipe.beacon.message.ResponseMessage;
-import io.xpipe.core.store.LocalFileDataStore;
+import io.xpipe.core.store.StreamDataStore;
 import lombok.Builder;
 import lombok.Value;
 import lombok.extern.jackson.Jacksonized;
 
-import java.nio.file.Path;
-
 public class PreStoreExchange implements MessageExchange<PreStoreExchange.Request, PreStoreExchange.Response> {
 
     @Override
@@ -36,6 +34,6 @@ public class PreStoreExchange implements MessageExchange<PreStoreExchange.Reques
     @Builder
     @Value
     public static class Response implements ResponseMessage {
-        LocalFileDataStore store;
+        StreamDataStore store;
     }
 }

beacon/src/main/java/io/xpipe/beacon/exchange/QueryDataSourceExchange.java

@@ -45,8 +45,6 @@ public class QueryDataSourceExchange implements MessageExchange<QueryDataSourceE
         @NonNull
         DataStore store;
         @NonNull
-        DataSourceDescriptor<?> descriptor;
-        @NonNull
         DataSourceConfigInstance config;
     }
 }

beacon/src/main/java/io/xpipe/beacon/exchange/ReadExecuteExchange.java

@@ -3,7 +3,7 @@ package io.xpipe.beacon.exchange;
 import io.xpipe.beacon.message.RequestMessage;
 import io.xpipe.beacon.message.ResponseMessage;
 import io.xpipe.core.source.DataSourceConfigInstance;
-import io.xpipe.core.source.DataSourceReference;
+import io.xpipe.core.source.DataSourceId;
 import io.xpipe.core.store.DataStore;
 import lombok.Builder;
 import lombok.NonNull;
@@ -35,8 +35,8 @@ public class ReadExecuteExchange implements MessageExchange<ReadExecuteExchange.
         DataStore dataStore;
         @NonNull
         DataSourceConfigInstance config;
-        @NonNull
-        DataSourceReference target;
+
+        DataSourceId target;
     }
 
     @Jacksonized

beacon/src/main/java/io/xpipe/beacon/exchange/ReadPreparationExchange.java

@@ -30,9 +30,10 @@ public class ReadPreparationExchange implements MessageExchange<ReadPreparationE
     @Builder
     @Value
     public static class Request implements RequestMessage {
-        String providerType;
+        String provider;
+
         @NonNull
-        String dataStore;
+        StreamDataStore store;
     }
 
     @Jacksonized
@@ -40,6 +41,6 @@ public class ReadPreparationExchange implements MessageExchange<ReadPreparationE
     @Value
     public static class Response implements ResponseMessage {
         DataSourceConfigInstance config;
-        StreamDataStore dataStore;
+        StreamDataStore store;
     }
 }

beacon/src/main/java/io/xpipe/beacon/exchange/StoreEditExchange.java

@@ -2,15 +2,12 @@ package io.xpipe.beacon.exchange;
 
 import io.xpipe.beacon.message.RequestMessage;
 import io.xpipe.beacon.message.ResponseMessage;
-import io.xpipe.core.source.DataSourceConfig;
+import io.xpipe.core.source.DataSourceConfigOptions;
 import io.xpipe.core.source.DataSourceId;
-import io.xpipe.core.source.DataSourceType;
 import lombok.Builder;
 import lombok.Value;
 import lombok.extern.jackson.Jacksonized;
 
-import java.net.URL;
-
 public class StoreEditExchange implements MessageExchange<StoreEditExchange.Request, StoreEditExchange.Response> {
 
     @Override
@@ -33,7 +30,7 @@ public class StoreEditExchange implements MessageExchange<StoreEditExchange.Requ
     @Value
     public static class Request implements RequestMessage {
         DataSourceId sourceId;
-        DataSourceConfig config;
+        DataSourceConfigOptions config;
     }
 
     @Jacksonized

beacon/src/main/java/io/xpipe/beacon/exchange/StoreResourceExchange.java

@@ -2,7 +2,7 @@ package io.xpipe.beacon.exchange;
 
 import io.xpipe.beacon.message.RequestMessage;
 import io.xpipe.beacon.message.ResponseMessage;
-import io.xpipe.core.source.DataSourceConfig;
+import io.xpipe.core.source.DataSourceConfigOptions;
 import io.xpipe.core.source.DataSourceId;
 import io.xpipe.core.source.DataSourceInfo;
 import lombok.Builder;
@@ -41,7 +41,7 @@ public class StoreResourceExchange implements MessageExchange<StoreResourceExcha
     @Value
     public static class Response implements ResponseMessage {
         DataSourceId sourceId;
-        DataSourceConfig config;
+        DataSourceConfigOptions config;
         DataSourceInfo info;
     }
 }

beacon/src/main/java/io/xpipe/beacon/exchange/StoreStreamExchange.java

@@ -4,7 +4,7 @@ import io.xpipe.beacon.message.RequestMessage;
 import io.xpipe.beacon.message.ResponseMessage;
 import io.xpipe.core.source.DataSourceId;
 import io.xpipe.core.source.DataSourceType;
-import io.xpipe.core.source.DataSourceConfig;
+import io.xpipe.core.source.DataSourceConfigOptions;
 import lombok.Builder;
 import lombok.Value;
 import lombok.extern.jackson.Jacksonized;
@@ -39,7 +39,7 @@ public class StoreStreamExchange implements MessageExchange<StoreStreamExchange.
     public static class Response implements ResponseMessage {
         DataSourceId sourceId;
         DataSourceType sourceType;
-        DataSourceConfig config;
+        DataSourceConfigOptions config;
         Object data;
     }
 }

beacon/src/main/java/io/xpipe/beacon/exchange/api/QueryTableDataExchange.java

@@ -0,0 +1,45 @@
+package io.xpipe.beacon.exchange.api;
+
+import io.xpipe.beacon.exchange.MessageExchange;
+import io.xpipe.beacon.message.RequestMessage;
+import io.xpipe.beacon.message.ResponseMessage;
+import io.xpipe.core.source.DataSourceId;
+import lombok.Builder;
+import lombok.NonNull;
+import lombok.Value;
+import lombok.extern.jackson.Jacksonized;
+
+public class QueryTableDataExchange implements MessageExchange<QueryTableDataExchange.Request, QueryTableDataExchange.Response> {
+
+    @Override
+    public String getId() {
+        return "queryTableData";
+    }
+
+    @Override
+    public Class<QueryTableDataExchange.Request> getRequestClass() {
+        return QueryTableDataExchange.Request.class;
+    }
+
+    @Override
+    public Class<QueryTableDataExchange.Response> getResponseClass() {
+        return QueryTableDataExchange.Response.class;
+    }
+
+    @Jacksonized
+    @Builder
+    @Value
+    public static class Request implements RequestMessage {
+        @NonNull
+        DataSourceId id;
+
+        @Builder.Default
+        int maxRows = -1;
+    }
+
+    @Jacksonized
+    @Builder
+    @Value
+    public static class Response implements ResponseMessage {
+    }
+}

beacon/src/main/java/module-info.java

@@ -1,20 +1,22 @@
 import io.xpipe.beacon.exchange.*;
+import io.xpipe.beacon.exchange.api.QueryTableDataExchange;
 
 module io.xpipe.beacon {
     exports io.xpipe.beacon;
     exports io.xpipe.beacon.exchange;
     exports io.xpipe.beacon.message;
+    exports io.xpipe.beacon.exchange.api;
+    exports io.xpipe.beacon.exchange.data;
+
+    opens io.xpipe.beacon;
+    opens io.xpipe.beacon.exchange;
+    opens io.xpipe.beacon.exchange.api;
+    opens io.xpipe.beacon.message;
+    opens io.xpipe.beacon.exchange.data;
 
     requires com.fasterxml.jackson.core;
     requires com.fasterxml.jackson.databind;
     requires transitive io.xpipe.core;
-
-    opens io.xpipe.beacon;
-    opens io.xpipe.beacon.exchange;
-    opens io.xpipe.beacon.message;
-    exports io.xpipe.beacon.exchange.data;
-    opens io.xpipe.beacon.exchange.data;
-
     requires static lombok;
 
     uses MessageExchange;
@@ -35,5 +37,6 @@ module io.xpipe.beacon {
             PreStoreExchange,
             EditPreparationExchange,
             EditExecuteExchange,
+            QueryTableDataExchange,
             VersionExchange;
 }