Isolate all grpc calls from parent context (#102)

Author: alex268

jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java

@@ -53,7 +53,7 @@ public BaseYdbExecutor(YdbContext ctx) {
     }
 
     protected Session createNewTableSession(YdbValidator validator) throws SQLException {
-        return validator.call("Get session", () -> tableClient.createSession(sessionTimeout));
+        return validator.call("Get session", null, () -> tableClient.createSession(sessionTimeout));
     }
 
     protected void closeCurrentResult() throws SQLException {
@@ -163,7 +163,7 @@ public YdbQueryResult executeScanQuery(YdbStatement statement, YdbQuery query, S
             }
         }
 
-        StreamQueryResult lazy = validator.call(msg, () -> {
+        StreamQueryResult lazy = validator.call(msg, null, () -> {
             final CompletableFuture<Result<StreamQueryResult>> future = new CompletableFuture<>();
             final GrpcReadStream<ResultSetReader> stream = session.executeScanQuery(yql, params, settings);
             final StreamQueryResult result = new StreamQueryResult(msg, statement, query, stream::cancel);

jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java

@@ -68,7 +68,7 @@ public QueryServiceExecutor(YdbContext ctx, int transactionLevel, boolean autoCo
     }
 
     protected QuerySession createNewQuerySession(YdbValidator validator) throws SQLException {
-        return validator.call("Get query session", () -> queryClient.createSession(sessionTimeout));
+        return validator.call("Get query session", null, () -> queryClient.createSession(sessionTimeout));
     }
 
     @Override
@@ -255,7 +255,7 @@ public YdbQueryResult executeDataQuery(
             tracer.query(yql);
             String msg = "STREAM_QUERY >>\n" + yql;
 
-            StreamQueryResult lazy = validator.call(msg, () -> {
+            StreamQueryResult lazy = validator.call(msg, tracer, () -> {
                 final CompletableFuture<Result<StreamQueryResult>> future = new CompletableFuture<>();
                 final QueryStream stream = localTx.createQuery(yql, isAutoCommit, params, settings);
                 final StreamQueryResult result = new StreamQueryResult(msg, statement, query, stream::cancel);

jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java

@@ -248,8 +248,7 @@ public boolean isValid(YdbValidator validator, int timeout) throws SQLException
         try {
             KeepAliveSessionSettings settings = new KeepAliveSessionSettings().setTimeout(Duration.ofSeconds(timeout));
             Session.State keepAlive = validator.call(
-                    "Keep alive: " + tx.txID(),
-                    () -> session.keepAlive(settings)
+                    "Keep alive: " + tx.txID(), null, () -> session.keepAlive(settings)
             );
             return keepAlive == Session.State.READY;
         } finally {

jdbc/src/main/java/tech/ydb/jdbc/context/YdbValidator.java

@@ -10,6 +10,8 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import io.grpc.Context;
+
 import tech.ydb.core.Issue;
 import tech.ydb.core.Result;
 import tech.ydb.core.Status;
@@ -54,8 +56,18 @@ public void clearWarnings() {
         this.issues.clear();
     }
 
+    private <T> T joinFuture(Supplier<CompletableFuture<T>> supplier) {
+        Context ctx = Context.current().fork();
+        Context previous = ctx.attach();
+        try {
+            return supplier.get().join();
+        } finally {
+            ctx.detach(previous);
+        }
+    }
+
     public void execute(String msg, YdbTracer tracer, Supplier<CompletableFuture<Status>> fn) throws SQLException {
-        Status status = fn.get().join();
+        Status status = joinFuture(fn);
         addStatusIssues(status);
 
         tracer.trace("<-- " + status.toString());
@@ -67,25 +79,18 @@ public void execute(String msg, YdbTracer tracer, Supplier<CompletableFuture<Sta
         }
     }
 
-    public <R> R call(String msg, Supplier<CompletableFuture<Result<R>>> fn) throws SQLException {
-        try {
-            Result<R> result = fn.get().join();
-            addStatusIssues(result.getStatus());
-            return result.getValue();
-        } catch (UnexpectedResultException ex) {
-            LOGGER.log(Level.FINE, "call problem {0}", ex.getStatus());
-            throw ExceptionFactory.createException("Cannot call '" + msg + "' with " + ex.getStatus(), ex);
-        }
-    }
-
     public <R> R call(String msg, YdbTracer tracer, Supplier<CompletableFuture<Result<R>>> fn) throws SQLException {
         try {
-            Result<R> result = fn.get().join();
+            Result<R> result = joinFuture(fn);
             addStatusIssues(result.getStatus());
-            tracer.trace("<-- " + result.getStatus().toString());
+            if (tracer != null) {
+                tracer.trace("<-- " + result.getStatus().toString());
+            }
             return result.getValue();
         } catch (UnexpectedResultException ex) {
-            tracer.close();
+            if (tracer != null) {
+                tracer.close();
+            }
             LOGGER.log(Level.FINE, "call problem {0}", ex.getStatus());
             throw ExceptionFactory.createException("Cannot call '" + msg + "' with " + ex.getStatus(), ex);
         }

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbDatabaseMetaDataImpl.java

@@ -1343,7 +1343,7 @@ private List<String> listTables(Predicate<String> filter) throws SQLException {
 
     private List<String> tables(String databasePrefix, String path, Predicate<String> filter) throws SQLException {
         ListDirectoryResult result = validator.call(
-                "List tables from " + path, () -> executor.listDirectory(path)
+                "List tables from " + path, null, () -> executor.listDirectory(path)
         );
 
         List<String> tables = new ArrayList<>();
@@ -1376,7 +1376,7 @@ private TableDescription describeTable(String table) throws SQLException {
 
         String databaseWithSuffix = withSuffix(connection.getCtx().getPrefixPath());
 
-        return validator.call("Describe table " + table, () -> executor
+        return validator.call("Describe table " + table, null, () -> executor
                 .describeTable(databaseWithSuffix + table, settings)
                 .thenApply(result -> {
                     // ignore scheme errors like path not found

jdbc/src/test/resources/logging.properties

@@ -2,7 +2,7 @@ handlers= java.util.logging.ConsoleHandler
 .level=INFO
 java.util.logging.ConsoleHandler.level = ALL
 java.util.logging.ConsoleHandler.formatter = java.util.logging.SimpleFormatter
-java.util.logging.SimpleFormatter.format= %1$tF %1$tT [%4$s] %3$s %5$s %n
+java.util.logging.SimpleFormatter.format= %1$tF %1$tT [%4$s]\t%3$s %5$s %n
 
 tech.ydb.jdbc.level=ALL
 #tech.ydb.level=FINEST