Merge pull request #283 from ydb-platform/find_export_methods

Added methods to find exisiting operations

Author: alex268

export/src/main/java/tech/ydb/export/ExportClient.java

@@ -14,6 +14,7 @@
 import tech.ydb.export.result.ExportToYtResult;
 import tech.ydb.export.settings.ExportToS3Settings;
 import tech.ydb.export.settings.ExportToYtSettings;
+import tech.ydb.export.settings.FindExportSettings;
 
 /**
  * @author Kirill Kurdyukov
@@ -24,6 +25,14 @@ static ExportClient newClient(@WillNotClose GrpcTransport transport) {
         return new ExportClientImpl(GrpcExportRpcImpl.useTransport(transport));
     }
 
+    CompletableFuture<Operation<Result<ExportToS3Result>>> findExportToS3(
+            String operationId, FindExportSettings settings
+    );
+
+    CompletableFuture<Operation<Result<ExportToYtResult>>> findExportToYT(
+            String operationId, FindExportSettings settings
+    );
+
     CompletableFuture<Operation<Result<ExportToS3Result>>> startExportToS3(
             String endpoint, String bucket, String accessKey, String secretKey, ExportToS3Settings settings
     );
@@ -40,7 +49,7 @@ default CompletableFuture<Result<ExportToS3Result>> exportToS3(
                 .thenCompose(operation -> OperationTray.fetchOperation(operation, updateRateSeconds));
     }
 
-    default CompletableFuture<Result<ExportToYtResult>> startExportToYt(
+    default CompletableFuture<Result<ExportToYtResult>> exportToYt(
             String host, String token, ExportToYtSettings settings, int updateRateSeconds
     ) {
         return startExportToYt(host, token, settings)

export/src/main/java/tech/ydb/export/ExportRpc.java

@@ -12,6 +12,14 @@
  */
 public interface ExportRpc {
 
+    CompletableFuture<Operation<Result<YdbExport.ExportToS3Result>>> findExportToS3(
+            String operationId, GrpcRequestSettings settings
+    );
+
+    CompletableFuture<Operation<Result<YdbExport.ExportToYtResult>>> findExportToYT(
+            String operationId, GrpcRequestSettings settings
+    );
+
     CompletableFuture<Operation<Result<YdbExport.ExportToS3Result>>> exportS3(
             YdbExport.ExportToS3Request request, GrpcRequestSettings settings
     );

export/src/main/java/tech/ydb/export/impl/ExportClientImpl.java

@@ -13,6 +13,7 @@
 import tech.ydb.export.result.ExportToYtResult;
 import tech.ydb.export.settings.ExportToS3Settings;
 import tech.ydb.export.settings.ExportToYtSettings;
+import tech.ydb.export.settings.FindExportSettings;
 import tech.ydb.proto.export.YdbExport;
 
 /**
@@ -26,11 +27,8 @@ public ExportClientImpl(ExportRpc exportRpc) {
         this.exportRpc = exportRpc;
     }
 
-    private String getTraceIdOrGenerateNew(String traceId) {
-        return traceId == null ? UUID.randomUUID().toString() : traceId;
-    }
-
-    private GrpcRequestSettings makeGrpcRequestSettings(BaseRequestSettings settings, String traceId) {
+    private GrpcRequestSettings makeGrpcRequestSettings(BaseRequestSettings settings) {
+        String traceId = settings.getTraceId() == null ? UUID.randomUUID().toString() : settings.getTraceId();
         return GrpcRequestSettings.newBuilder()
                 .withDeadline(settings.getRequestTimeout())
                 .withTraceId(traceId)
@@ -85,8 +83,7 @@ public CompletableFuture<Operation<Result<ExportToS3Result>>> startExportToS3(
                 .setOperationParams(Operation.buildParams(settings))
                 .build();
 
-        String traceId = getTraceIdOrGenerateNew(settings.getTraceId());
-        return exportRpc.exportS3(request, makeGrpcRequestSettings(settings, traceId))
+        return exportRpc.exportS3(request, makeGrpcRequestSettings(settings))
                 .thenApply(op -> op.transform(r -> r.map(ExportToS3Result::new)));
     }
 
@@ -128,8 +125,23 @@ public CompletableFuture<Operation<Result<ExportToYtResult>>> startExportToYt(
                 .setOperationParams(Operation.buildParams(settings))
                 .build();
 
-        String traceId = getTraceIdOrGenerateNew(settings.getTraceId());
-        return exportRpc.exportYt(request, makeGrpcRequestSettings(settings, traceId))
+        return exportRpc.exportYt(request, makeGrpcRequestSettings(settings))
+                .thenApply(op -> op.transform(r -> r.map(ExportToYtResult::new)));
+    }
+
+    @Override
+    public CompletableFuture<Operation<Result<ExportToS3Result>>> findExportToS3(
+            String operationId, FindExportSettings settings
+    ) {
+        return exportRpc.findExportToS3(operationId, makeGrpcRequestSettings(settings))
+                .thenApply(op -> op.transform(r -> r.map(ExportToS3Result::new)));
+    }
+
+    @Override
+    public CompletableFuture<Operation<Result<ExportToYtResult>>> findExportToYT(
+            String operationId, FindExportSettings settings
+    ) {
+        return exportRpc.findExportToYT(operationId, makeGrpcRequestSettings(settings))
                 .thenApply(op -> op.transform(r -> r.map(ExportToYtResult::new)));
     }
 }

export/src/main/java/tech/ydb/export/impl/GrpcExportRpcImpl.java

@@ -10,8 +10,10 @@
 import tech.ydb.core.operation.Operation;
 import tech.ydb.core.operation.OperationBinder;
 import tech.ydb.export.ExportRpc;
+import tech.ydb.proto.OperationProtos;
 import tech.ydb.proto.export.YdbExport;
 import tech.ydb.proto.export.v1.ExportServiceGrpc;
+import tech.ydb.proto.operation.v1.OperationServiceGrpc;
 
 /**
  * @author Kirill Kurdyukov
@@ -49,4 +51,32 @@ public CompletableFuture<Operation<Result<YdbExport.ExportToYtResult>>> exportYt
                 )
         );
     }
+
+    @Override
+    public CompletableFuture<Operation<Result<YdbExport.ExportToS3Result>>> findExportToS3(
+            String operationId, GrpcRequestSettings settings
+    ) {
+        OperationProtos.GetOperationRequest request = OperationProtos.GetOperationRequest.newBuilder()
+                        .setId(operationId)
+                        .build();
+        return transport
+                .unaryCall(OperationServiceGrpc.getGetOperationMethod(), settings, request)
+                .thenApply(OperationBinder.bindAsync(
+                        transport, OperationProtos.GetOperationResponse::getOperation, YdbExport.ExportToS3Result.class
+                ));
+    }
+
+    @Override
+    public CompletableFuture<Operation<Result<YdbExport.ExportToYtResult>>> findExportToYT(
+            String operationId, GrpcRequestSettings settings
+    ) {
+        OperationProtos.GetOperationRequest request = OperationProtos.GetOperationRequest.newBuilder()
+                        .setId(operationId)
+                        .build();
+        return transport
+                .unaryCall(OperationServiceGrpc.getGetOperationMethod(), settings, request)
+                .thenApply(OperationBinder.bindAsync(
+                        transport, OperationProtos.GetOperationResponse::getOperation, YdbExport.ExportToYtResult.class
+                ));
+    }
 }

export/src/main/java/tech/ydb/export/settings/FindExportSettings.java

@@ -0,0 +1,24 @@
+package tech.ydb.export.settings;
+
+import tech.ydb.core.settings.OperationSettings;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class FindExportSettings extends OperationSettings {
+    private FindExportSettings(Builder builder) {
+        super(builder);
+    }
+
+    public static Builder newBuilder() {
+        return new Builder().withAsyncMode(true);
+    }
+
+    public static class Builder extends OperationSettings.OperationBuilder<Builder> {
+        @Override
+        public FindExportSettings build() {
+            return new FindExportSettings(this);
+        }
+    }
+}