Updates test for generated tests

jdbc/src/main/java/tech/ydb/jdbc/YdbTracer.java

@@ -9,20 +9,12 @@
  * @author Aleksandr Gorshenin
  */
 public interface YdbTracer {
-    interface Storage {
-        YdbTracer get();
-
-        default void clear() {
-
-        }
-    }
-
     static YdbTracer current() {
-        return YdbTracerImpl.ENABLED.get();
+        return YdbTracerImpl.get();
     }
 
     static void clear() {
-        YdbTracerImpl.ENABLED.clear();
+        YdbTracerImpl.clear();
     }
 
     void setId(String id);

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

@@ -27,7 +27,6 @@
 import tech.ydb.jdbc.YdbTracer;
 import tech.ydb.jdbc.common.YdbTypes;
 import tech.ydb.jdbc.exception.ExceptionFactory;
-import tech.ydb.jdbc.impl.YdbTracerImpl;
 import tech.ydb.jdbc.impl.YdbTracerNone;
 import tech.ydb.jdbc.query.QueryKey;
 import tech.ydb.jdbc.query.QueryType;
@@ -92,8 +91,6 @@ public class YdbContext implements AutoCloseable {
     private final boolean autoResizeSessionPool;
     private final AtomicInteger connectionsCount = new AtomicInteger();
 
-    private YdbTracer.Storage tracerStorage;
-
     private YdbContext(
             YdbConfig config,
             YdbOperationProperties operationProperties,
@@ -141,18 +138,12 @@ private YdbContext(
             prefixPath = transport.getDatabase();
             prefixPragma = "";
         }
-
-        this.tracerStorage = config.isTxTracedEnabled() ? YdbTracerImpl.ENABLED : YdbTracerNone.DISABLED;
     }
 
     public YdbTypes getTypes() {
         return types;
     }
 
-    public void setTracerStorage(YdbTracer.Storage storage) {
-        this.tracerStorage = storage;
-    }
-
     /**
      * Grpc Transport for other API YDB server clients
      *
@@ -163,7 +154,7 @@ public GrpcTransport getGrpcTransport() {
     }
 
     public YdbTracer getTracer() {
-        return tracerStorage.get();
+        return config.isTxTracedEnabled() ? YdbTracer.current() : YdbTracerNone.DISABLED;
     }
 
     private String joined(String path1, String path2) {
@@ -270,7 +261,7 @@ public void register() {
     }
 
     public void deregister() {
-        tracerStorage.clear();
+        YdbTracer.clear();
 
         int actual = connectionsCount.decrementAndGet();
         int maxSize = tableClient.sessionPoolStats().getMaxSize();
@@ -416,7 +407,7 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
             String tablePath = joined(getPrefixPath(), query.getYqlBatcher().getTableName());
             TableDescription description = tableDescribeCache.getIfPresent(tablePath);
             if (description == null) {
-                YdbTracer tracer = tracerStorage.get();
+                YdbTracer tracer = getTracer();
                 tracer.trace("--> describe table");
                 tracer.trace(tablePath);
 
@@ -459,7 +450,7 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
         Map<String, Type> queryTypes = queryParamsCache.getIfPresent(query.getOriginQuery());
         if (queryTypes == null) {
             String yql = prefixPragma + query.getPreparedYql();
-            YdbTracer tracer = tracerStorage.get();
+            YdbTracer tracer = getTracer();
             tracer.trace("--> prepare data query");
             tracer.trace(yql);
 

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

@@ -153,8 +153,7 @@ protected boolean updateState(YdbQueryResult result) throws SQLException {
         return state.hasResultSets();
     }
 
-    protected YdbQueryResult executeBulkUpsert(YdbQuery query, String tablePath, ListValue rows)
-            throws SQLException {
+    protected YdbQueryResult executeBulkUpsert(YdbQuery query, String tablePath, ListValue rows) throws SQLException {
         ensureOpened();
 
         if (connection.getExecutor().isInsideTransaction()) {

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

@@ -93,10 +93,12 @@ public int[] executeBatch() throws SQLException {
         try {
             if (query.getType() == QueryType.BULK_QUERY && (prepared instanceof BulkUpsertQuery)) {
                 BulkUpsertQuery bulk = (BulkUpsertQuery) prepared;
-                executeBulkUpsert(query, bulk.getTablePath(), bulk.getBatchedBulk());
+                YdbQueryResult newState = executeBulkUpsert(query, bulk.getTablePath(), bulk.getBatchedBulk());
+                updateState(newState);
             } else {
                 for (Params prm: prepared.getBatchParams()) {
-                    executeDataQuery(query, prepared.getBatchText(prm), prm);
+                    YdbQueryResult newState = executeDataQuery(query, prepared.getBatchText(prm), prm);
+                    updateState(newState);
                 }
             }
         } finally {

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

@@ -19,32 +19,6 @@ public class YdbTracerImpl implements YdbTracer {
     private static final ThreadLocal<YdbTracer> LOCAL = new ThreadLocal<>();
     private static final AtomicLong ANONYMOUS_COUNTER = new AtomicLong(0);
 
-    public static final Storage ENABLED = new Storage() {
-        @Override
-        public YdbTracer get() {
-            YdbTracer tracer = LOCAL.get();
-            if (tracer == null) {
-                tracer = new YdbTracerImpl();
-                LOCAL.set(tracer);
-            }
-
-            return tracer;
-        }
-
-        @Override
-        public void clear() {
-            LOCAL.remove();
-        }
-    };
-
-    private final Date startDate = new Date();
-    private final long startedAt = System.currentTimeMillis();
-    private final List<Record> records = new ArrayList<>();
-
-    private String txID = null;
-    private String label = null;
-    private boolean isMarked = false;
-
     private class Record {
         private final long executedAt = System.currentTimeMillis();
         private final String message;
@@ -56,70 +30,119 @@ private class Record {
         }
     }
 
+    private class Tx {
+        private final Date startDate = new Date();
+        private final long startedAt = System.currentTimeMillis();
+        private final List<Record> records = new ArrayList<>();
+
+        private String id = null;
+        private String label = null;
+        private boolean isMarked = false;
+
+        private void log(Level level) {
+            if (!LOGGER.isLoggable(level) || records.isEmpty()) {
+                return;
+            }
+
+            long finishedAt = System.currentTimeMillis();
+            long requestsTime = 0;
+
+            String idName = id != null ? id : "anonymous-" + ANONYMOUS_COUNTER.incrementAndGet();
+            String traceID = label == null ? idName : label + "-" + idName;
+            LOGGER.log(level, "Trace[{0}] started at {1}", new Object[] {traceID, startDate});
+            long last = startedAt;
+            long requestsCount = 0;
+            boolean lastIsRequest = false;
+            for (Record record: records) {
+                if (record.isRequest) {
+                    requestsCount++;
+                    lastIsRequest = true;
+                    if (record.message != null) {
+                        String clean =  record.message.replaceAll("\\s", " ");
+                        LOGGER.log(level, "Query[{0}] {1}", new Object[] {traceID, clean});
+                    }
+                } else {
+                    long ms = record.executedAt - last;
+                    if (lastIsRequest) {
+                        requestsTime += ms;
+                        lastIsRequest = false;
+                    }
+                    LOGGER.log(level, "Trace[{0}] {1} ms {2}", new Object[] {traceID, ms, record.message});
+                    last = record.executedAt;
+                }
+            }
+            LOGGER.log(level, "Trace[{0}] finished in {1} ms, {2} requests take {3} ms", new Object[] {
+                traceID, finishedAt - startedAt, requestsCount, requestsTime
+            });
+        }
+    }
+
+    private Tx tx = null;
+
+    public static <T extends YdbTracer> T use(T tracer) {
+        LOCAL.set(tracer);
+        return tracer;
+    }
+
+    public static YdbTracer get() {
+        YdbTracer tracer = LOCAL.get();
+        if (tracer == null) {
+            tracer = new YdbTracerImpl();
+            LOCAL.set(tracer);
+        }
+
+        return tracer;
+    }
+
+    public static void clear() {
+        YdbTracer tracer = LOCAL.get();
+        if (tracer != null) {
+            tracer.close();
+        }
+        LOCAL.remove();
+    }
+
+    private Tx ensureOpen() {
+        if (tx == null) {
+            tx = new Tx();
+        }
+        return tx;
+    }
+
     @Override
     public void trace(String message) {
-        records.add(new Record(message, false));
+        ensureOpen().records.add(new Record(message, false));
     }
 
     @Override
     public void query(String queryText) {
-        records.add(new Record(queryText, true));
+        ensureOpen().records.add(new Record(queryText, true));
     }
 
     @Override
     public void setId(String id) {
-        if (!Objects.equals(id, txID)) {
-            this.txID = id;
+        Tx local = ensureOpen();
+        if (!Objects.equals(id, local.id)) {
+            local.id = id;
             trace("set-id " + id);
         }
     }
 
     @Override
     public void markToPrint(String label) {
-        if (!isMarked || !Objects.equals(label, this.label)) {
-            this.isMarked = true;
-            this.label = label;
+        Tx local = ensureOpen();
+        if (!local.isMarked || !Objects.equals(label, local.label)) {
+            local.isMarked = true;
+            local.label = label;
             trace("markToPrint " + label);
         }
     }
 
     @Override
     public void close() {
-        LOCAL.remove();
-
-        final Level level = isMarked ? Level.INFO : Level.FINE;
-        if (!LOGGER.isLoggable(level) || records.isEmpty()) {
-            return;
-        }
-
-        long finishedAt = System.currentTimeMillis();
-        long requestsTime = 0;
-
-        String id = txID != null ? txID : "anonymous-" + ANONYMOUS_COUNTER.incrementAndGet();
-        String traceID = label == null ? id : label + "-" + id;
-        LOGGER.log(level, "Trace[{0}] started at {1}", new Object[] {traceID, startDate});
-        long last = startedAt;
-        long requestsCount = 0;
-        boolean lastIsRequest = false;
-        for (Record record: records) {
-            if (record.isRequest) {
-                requestsCount++;
-                lastIsRequest = true;
-                if (record.message != null) {
-                    LOGGER.log(level, "Query[{0}] {1}", new Object[] {traceID, record.message.replaceAll("\\s", " ")});
-                }
-            } else {
-                long ms = record.executedAt - last;
-                if (lastIsRequest) {
-                    requestsTime += ms;
-                    lastIsRequest = false;
-                }
-                LOGGER.log(level, "Trace[{0}] {1} ms {2}", new Object[] {traceID, ms, record.message});
-                last = record.executedAt;
-            }
+        if (tx != null) {
+            tx.log(tx.isMarked ? Level.INFO : Level.FINE);
+            tx = null;
         }
-        LOGGER.log(level, "Trace[{0}] finished in {1} ms, {2} requests take {3} ms", new Object[] {
-            traceID, finishedAt - startedAt, requestsCount, requestsTime
-        });
     }
 }

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

@@ -7,8 +7,7 @@
  * @author Aleksandr Gorshenin
  */
 public class YdbTracerNone implements YdbTracer {
-    private static final YdbTracer NONE = new YdbTracerNone();
-    public static final YdbTracer.Storage DISABLED = () -> NONE;
+    public static final YdbTracer DISABLED = new YdbTracerNone();
 
     @Override
     public void trace(String message) { }