diff --git a/driver/src/main/java/org/neo4j/driver/Config.java b/driver/src/main/java/org/neo4j/driver/Config.java index 0bfeb0501a..e3b5406f70 100644 --- a/driver/src/main/java/org/neo4j/driver/Config.java +++ b/driver/src/main/java/org/neo4j/driver/Config.java @@ -148,6 +148,13 @@ public final class Config implements Serializable { */ private final MetricsAdapter metricsAdapter; + /** + * Specify if telemetry collection is disabled. + *

+ * By default, the driver will send anonymous usage statistics to the server it connects to if the server requests those. + */ + private final boolean telemetryDisabled; + private Config(ConfigBuilder builder) { this.logging = builder.logging; this.logLeakedSessions = builder.logLeakedSessions; @@ -169,6 +176,7 @@ private Config(ConfigBuilder builder) { this.eventLoopThreads = builder.eventLoopThreads; this.metricsAdapter = builder.metricsAdapter; + this.telemetryDisabled = builder.telemetryDisabled; } /** @@ -335,6 +343,18 @@ public String userAgent() { return userAgent; } + /** + * Returns if the telemetry is disabled on the driver side. + *

+ * The telemetry is collected only when it is enabled both the server and the driver. + * + * @return {@code true} if telemetry is disabled or {@code false} otherwise + * @since 5.13 + */ + public boolean isTelemetryDisabled() { + return telemetryDisabled; + } + /** * Used to build new config instances */ @@ -357,6 +377,8 @@ public static final class ConfigBuilder { private int eventLoopThreads = 0; private NotificationConfig notificationConfig = NotificationConfig.defaultConfig(); + private boolean telemetryDisabled = false; + private ConfigBuilder() {} /** @@ -748,6 +770,31 @@ public ConfigBuilder withNotificationConfig(NotificationConfig notificationConfi return this; } + /** + * Sets if telemetry is disabled on the driver side. + *

+ * By default, the driver sends anonymous telemetry data to the server it connects to if the server has + * telemetry enabled. This can be explicitly disabled on the driver side by setting this setting to + * {@code true}. + *

+ * At present, the driver sends which API type is used, like: + *

+ * + * @param telemetryDisabled {@code true} if telemetry is disabled or {@code false} otherwise + * @return this builder + * @since 5.13 + */ + public ConfigBuilder withTelemetryDisabled(boolean telemetryDisabled) { + this.telemetryDisabled = telemetryDisabled; + return this; + } + /** * Create a config instance from this builder. * diff --git a/driver/src/main/java/org/neo4j/driver/internal/DriverFactory.java b/driver/src/main/java/org/neo4j/driver/internal/DriverFactory.java index 2b0a41e207..0a5e2f6344 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/DriverFactory.java +++ b/driver/src/main/java/org/neo4j/driver/internal/DriverFactory.java @@ -269,7 +269,8 @@ protected InternalDriver createRoutingDriver( */ protected InternalDriver createDriver( SecurityPlan securityPlan, SessionFactory sessionFactory, MetricsProvider metricsProvider, Config config) { - return new InternalDriver(securityPlan, sessionFactory, metricsProvider, config.logging()); + return new InternalDriver( + securityPlan, sessionFactory, metricsProvider, config.isTelemetryDisabled(), config.logging()); } /** diff --git a/driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java b/driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java index 846e9cafd9..824b3fa740 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java +++ b/driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java @@ -65,6 +65,8 @@ public class InternalDriver implements Driver { private final SessionFactory sessionFactory; private final Logger log; + private final boolean telemetryDisabled; + private final AtomicBoolean closed = new AtomicBoolean(false); private final MetricsProvider metricsProvider; @@ -72,11 +74,13 @@ public class InternalDriver implements Driver { SecurityPlan securityPlan, SessionFactory sessionFactory, MetricsProvider metricsProvider, + boolean telemetryDisabled, Logging logging) { this.securityPlan = securityPlan; this.sessionFactory = sessionFactory; this.metricsProvider = metricsProvider; this.log = logging.getLog(getClass()); + this.telemetryDisabled = telemetryDisabled; } @Override @@ -215,7 +219,7 @@ private static RuntimeException driverCloseException() { public NetworkSession newSession(SessionConfig config, AuthToken overrideAuthToken) { assertOpen(); - var session = sessionFactory.newInstance(config, overrideAuthToken); + var session = sessionFactory.newInstance(config, overrideAuthToken, telemetryDisabled); if (closed.get()) { // session does not immediately acquire connection, it is fine to just throw throw driverCloseException(); diff --git a/driver/src/main/java/org/neo4j/driver/internal/InternalExecutableQuery.java b/driver/src/main/java/org/neo4j/driver/internal/InternalExecutableQuery.java index 096875f14c..e0a52796bc 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/InternalExecutableQuery.java +++ b/driver/src/main/java/org/neo4j/driver/internal/InternalExecutableQuery.java @@ -32,6 +32,7 @@ import org.neo4j.driver.SessionConfig; import org.neo4j.driver.TransactionCallback; import org.neo4j.driver.TransactionConfig; +import org.neo4j.driver.internal.telemetry.TelemetryApi; public class InternalExecutableQuery implements ExecutableQuery { private final Driver driver; @@ -81,7 +82,8 @@ public T execute(Collector recordCollector, ResultFinish return resultFinisher.finish(result.keys(), finishedValue, summary); }; var accessMode = config.routing().equals(RoutingControl.WRITE) ? AccessMode.WRITE : AccessMode.READ; - return session.execute(accessMode, txCallback, TransactionConfig.empty(), false); + return session.execute( + accessMode, txCallback, TransactionConfig.empty(), TelemetryApi.EXECUTABLE_QUERY, false); } } diff --git a/driver/src/main/java/org/neo4j/driver/internal/InternalSession.java b/driver/src/main/java/org/neo4j/driver/internal/InternalSession.java index f13ec36a2c..c126e40500 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/InternalSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/InternalSession.java @@ -34,6 +34,8 @@ import org.neo4j.driver.exceptions.ClientException; import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.spi.Connection; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.Futures; public class InternalSession extends AbstractQueryRunner implements Session { @@ -93,7 +95,7 @@ public Transaction beginTransaction(TransactionConfig config) { public Transaction beginTransaction(TransactionConfig config, String txType) { var tx = Futures.blockingGet( - session.beginTransactionAsync(config, txType), + session.beginTransactionAsync(config, txType, new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION)), () -> terminateConnectionOnThreadInterrupt("Thread interrupted while starting a transaction")); return new InternalTransaction(tx); } @@ -107,12 +109,12 @@ public T readTransaction(TransactionWork work) { @Override @Deprecated public T readTransaction(TransactionWork work, TransactionConfig config) { - return transaction(AccessMode.READ, work, config, true); + return transaction(AccessMode.READ, work, config, TelemetryApi.MANAGED_TRANSACTION, true); } @Override public T executeRead(TransactionCallback callback, TransactionConfig config) { - return execute(AccessMode.READ, callback, config, true); + return execute(AccessMode.READ, callback, config, TelemetryApi.MANAGED_TRANSACTION, true); } @Override @@ -124,12 +126,12 @@ public T writeTransaction(TransactionWork work) { @Override @Deprecated public T writeTransaction(TransactionWork work, TransactionConfig config) { - return transaction(AccessMode.WRITE, work, config, true); + return transaction(AccessMode.WRITE, work, config, TelemetryApi.MANAGED_TRANSACTION, true); } @Override public T executeWrite(TransactionCallback callback, TransactionConfig config) { - return execute(AccessMode.WRITE, callback, config, true); + return execute(AccessMode.WRITE, callback, config, TelemetryApi.MANAGED_TRANSACTION, true); } @Override @@ -151,21 +153,29 @@ public void reset() { () -> terminateConnectionOnThreadInterrupt("Thread interrupted while resetting the session")); } - T execute(AccessMode accessMode, TransactionCallback callback, TransactionConfig config, boolean flush) { - return transaction(accessMode, tx -> callback.execute(new DelegatingTransactionContext(tx)), config, flush); + T execute( + AccessMode accessMode, + TransactionCallback callback, + TransactionConfig config, + TelemetryApi telemetryApi, + boolean flush) { + return transaction( + accessMode, tx -> callback.execute(new DelegatingTransactionContext(tx)), config, telemetryApi, flush); } private T transaction( AccessMode mode, @SuppressWarnings("deprecation") TransactionWork work, TransactionConfig config, + TelemetryApi telemetryApi, boolean flush) { // use different code path compared to async so that work is executed in the caller thread // caller thread will also be the one who sleeps between retries; // it is unsafe to execute retries in the event loop threads because this can cause a deadlock // event loop thread will bock and wait for itself to read some data + var apiTelemetryWork = new ApiTelemetryWork(telemetryApi); return session.retryLogic().retry(() -> { - try (var tx = beginTransaction(mode, config, flush)) { + try (var tx = beginTransaction(mode, config, apiTelemetryWork, flush)) { var result = work.execute(tx); if (result instanceof Result) { @@ -182,9 +192,10 @@ private T transaction( }); } - private Transaction beginTransaction(AccessMode mode, TransactionConfig config, boolean flush) { + private Transaction beginTransaction( + AccessMode mode, TransactionConfig config, ApiTelemetryWork apiTelemetryWork, boolean flush) { var tx = Futures.blockingGet( - session.beginTransactionAsync(mode, config, null, flush), + session.beginTransactionAsync(mode, config, null, apiTelemetryWork, flush), () -> terminateConnectionOnThreadInterrupt("Thread interrupted while starting a transaction")); return new InternalTransaction(tx); } diff --git a/driver/src/main/java/org/neo4j/driver/internal/SessionFactory.java b/driver/src/main/java/org/neo4j/driver/internal/SessionFactory.java index fb334cad99..54394f700f 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/SessionFactory.java +++ b/driver/src/main/java/org/neo4j/driver/internal/SessionFactory.java @@ -24,7 +24,7 @@ import org.neo4j.driver.internal.async.NetworkSession; public interface SessionFactory { - NetworkSession newInstance(SessionConfig sessionConfig, AuthToken overrideAuthToken); + NetworkSession newInstance(SessionConfig sessionConfig, AuthToken overrideAuthToken, boolean telemetryDisabled); CompletionStage verifyConnectivity(); diff --git a/driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java b/driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java index f6f632bebb..926e4eb887 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java +++ b/driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java @@ -53,7 +53,8 @@ public class SessionFactoryImpl implements SessionFactory { } @Override - public NetworkSession newInstance(SessionConfig sessionConfig, AuthToken overrideAuthToken) { + public NetworkSession newInstance( + SessionConfig sessionConfig, AuthToken overrideAuthToken, boolean telemetryDisabled) { return createSession( connectionProvider, retryLogic, @@ -65,7 +66,8 @@ public NetworkSession newInstance(SessionConfig sessionConfig, AuthToken overrid logging, sessionConfig.bookmarkManager().orElse(NoOpBookmarkManager.INSTANCE), sessionConfig.notificationConfig(), - overrideAuthToken); + overrideAuthToken, + telemetryDisabled); } private Set toDistinctSet(Iterable bookmarks) { @@ -142,7 +144,8 @@ private NetworkSession createSession( Logging logging, BookmarkManager bookmarkManager, NotificationConfig notificationConfig, - AuthToken authToken) { + AuthToken authToken, + boolean telemetryDisabled) { Objects.requireNonNull(bookmarks, "bookmarks may not be null"); Objects.requireNonNull(bookmarkManager, "bookmarkManager may not be null"); return leakedSessionsLoggingEnabled @@ -157,7 +160,8 @@ private NetworkSession createSession( logging, bookmarkManager, notificationConfig, - authToken) + authToken, + telemetryDisabled) : new NetworkSession( connectionProvider, retryLogic, @@ -169,6 +173,7 @@ private NetworkSession createSession( logging, bookmarkManager, notificationConfig, - authToken); + authToken, + telemetryDisabled); } } diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/InternalAsyncSession.java b/driver/src/main/java/org/neo4j/driver/internal/async/InternalAsyncSession.java index 25b32df914..cc17c4e0b9 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/InternalAsyncSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/InternalAsyncSession.java @@ -38,6 +38,8 @@ import org.neo4j.driver.async.ResultCursor; import org.neo4j.driver.exceptions.ClientException; import org.neo4j.driver.internal.InternalBookmark; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.Futures; public class InternalAsyncSession extends AsyncAbstractQueryRunner implements AsyncSession { @@ -80,7 +82,8 @@ public CompletionStage beginTransactionAsync() { @Override public CompletionStage beginTransactionAsync(TransactionConfig config) { - return session.beginTransactionAsync(config).thenApply(InternalAsyncTransaction::new); + return session.beginTransactionAsync(config, new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION)) + .thenApply(InternalAsyncTransaction::new); } @Override @@ -136,9 +139,10 @@ private CompletionStage transactionAsync( AccessMode mode, @SuppressWarnings("deprecation") AsyncTransactionWork> work, TransactionConfig config) { + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.MANAGED_TRANSACTION); return session.retryLogic().retryAsync(() -> { var resultFuture = new CompletableFuture(); - var txFuture = session.beginTransactionAsync(mode, config); + var txFuture = session.beginTransactionAsync(mode, config, apiTelemetryWork); txFuture.whenComplete((tx, completionError) -> { var error = Futures.completionExceptionCause(completionError); diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSession.java b/driver/src/main/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSession.java index 35494f17ba..b16c2cda69 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSession.java @@ -48,7 +48,8 @@ public LeakLoggingNetworkSession( Logging logging, BookmarkManager bookmarkManager, NotificationConfig notificationConfig, - AuthToken overrideAuthToken) { + AuthToken overrideAuthToken, + boolean telemetryDisabled) { super( connectionProvider, retryLogic, @@ -60,7 +61,8 @@ public LeakLoggingNetworkSession( logging, bookmarkManager, notificationConfig, - overrideAuthToken); + overrideAuthToken, + telemetryDisabled); this.stackTrace = captureStackTrace(); } diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/NetworkConnection.java b/driver/src/main/java/org/neo4j/driver/internal/async/NetworkConnection.java index 005591a362..012c36a168 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/NetworkConnection.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/NetworkConnection.java @@ -67,6 +67,7 @@ public class NetworkConnection implements Connection { private final InboundMessageDispatcher messageDispatcher; private final String serverAgent; private final BoltServerAddress serverAddress; + private final boolean telemetryEnabled; private final BoltProtocol protocol; private final ExtendedChannelPool channelPool; private final CompletableFuture releaseFuture; @@ -92,6 +93,7 @@ public NetworkConnection( this.messageDispatcher = ChannelAttributes.messageDispatcher(channel); this.serverAgent = ChannelAttributes.serverAgent(channel); this.serverAddress = ChannelAttributes.serverAddress(channel); + this.telemetryEnabled = ChannelAttributes.telemetryEnabled(channel); this.protocol = BoltProtocol.forChannel(channel); this.channelPool = channelPool; this.releaseFuture = new CompletableFuture<>(); @@ -136,6 +138,11 @@ public void writeAndFlush(Message message, ResponseHandler handler) { } } + @Override + public boolean isTelemetryEnabled() { + return telemetryEnabled; + } + @Override public CompletionStage reset(Throwable throwable) { var result = new CompletableFuture(); diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/NetworkSession.java b/driver/src/main/java/org/neo4j/driver/internal/async/NetworkSession.java index 820523c302..9eb4c1ff2a 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/NetworkSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/NetworkSession.java @@ -54,6 +54,8 @@ import org.neo4j.driver.internal.retry.RetryLogic; import org.neo4j.driver.internal.spi.Connection; import org.neo4j.driver.internal.spi.ConnectionProvider; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.Futures; public class NetworkSession { @@ -74,6 +76,7 @@ public class NetworkSession { private volatile Set lastUsedBookmarks = Collections.emptySet(); private volatile Set lastReceivedBookmarks; private final NotificationConfig notificationConfig; + private final boolean telemetryDisabled; public NetworkSession( ConnectionProvider connectionProvider, @@ -86,7 +89,8 @@ public NetworkSession( Logging logging, BookmarkManager bookmarkManager, NotificationConfig notificationConfig, - AuthToken overrideAuthToken) { + AuthToken overrideAuthToken, + boolean telemetryDisabled) { Objects.requireNonNull(bookmarks, "bookmarks may not be null"); Objects.requireNonNull(bookmarkManager, "bookmarkManager may not be null"); this.connectionProvider = connectionProvider; @@ -104,6 +108,7 @@ public NetworkSession( databaseNameFuture, determineBookmarks(false), impersonatedUser, overrideAuthToken); this.fetchSize = fetchSize; this.notificationConfig = notificationConfig; + this.telemetryDisabled = telemetryDisabled; } public CompletionStage runAsync(Query query, TransactionConfig config) { @@ -126,22 +131,31 @@ public CompletionStage runRx( return newResultCursorStage; } - public CompletionStage beginTransactionAsync(TransactionConfig config) { - return beginTransactionAsync(mode, config, null, true); + public CompletionStage beginTransactionAsync( + TransactionConfig config, ApiTelemetryWork apiTelemetryWork) { + return beginTransactionAsync(mode, config, null, apiTelemetryWork, true); } - public CompletionStage beginTransactionAsync(TransactionConfig config, String txType) { - return this.beginTransactionAsync(mode, config, txType, true); + public CompletionStage beginTransactionAsync( + TransactionConfig config, String txType, ApiTelemetryWork apiTelemetryWork) { + return this.beginTransactionAsync(mode, config, txType, apiTelemetryWork, true); } - public CompletionStage beginTransactionAsync(AccessMode mode, TransactionConfig config) { - return beginTransactionAsync(mode, config, null, true); + public CompletionStage beginTransactionAsync( + AccessMode mode, TransactionConfig config, ApiTelemetryWork apiTelemetryWork) { + return beginTransactionAsync(mode, config, null, apiTelemetryWork, true); } public CompletionStage beginTransactionAsync( - AccessMode mode, TransactionConfig config, String txType, boolean flush) { + AccessMode mode, + TransactionConfig config, + String txType, + ApiTelemetryWork apiTelemetryWork, + boolean flush) { ensureSessionIsOpen(); + apiTelemetryWork.setEnabled(!telemetryDisabled); + // create a chain that acquires connection and starts a transaction var newTransactionStage = ensureNoOpenTxBeforeStartingTx() .thenCompose(ignore -> acquireConnection(mode)) @@ -149,7 +163,12 @@ public CompletionStage beginTransactionAsync( ImpersonationUtil.ensureImpersonationSupport(connection, connection.impersonatedUser())) .thenCompose(connection -> { var tx = new UnmanagedTransaction( - connection, this::handleNewBookmark, fetchSize, notificationConfig, logging); + connection, + this::handleNewBookmark, + fetchSize, + notificationConfig, + apiTelemetryWork, + logging); return tx.beginAsync(determineBookmarks(true), config, txType, flush); }); @@ -258,6 +277,9 @@ private CompletionStage buildResultCursorFactory(Query quer ImpersonationUtil.ensureImpersonationSupport(connection, connection.impersonatedUser())) .thenCompose(connection -> { try { + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.AUTO_COMMIT_TRANSACTION); + apiTelemetryWork.setEnabled(!telemetryDisabled); + var telemetryStage = apiTelemetryWork.execute(connection, connection.protocol()); var factory = connection .protocol() .runInAutoCommitTransaction( @@ -269,7 +291,13 @@ private CompletionStage buildResultCursorFactory(Query quer fetchSize, notificationConfig, logging); - return completedFuture(factory); + var future = completedFuture(factory); + telemetryStage.whenComplete((unused, throwable) -> { + if (throwable != null) { + future.completeExceptionally(throwable); + } + }); + return future; } catch (Throwable e) { return Futures.failedFuture(e); } diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/UnmanagedTransaction.java b/driver/src/main/java/org/neo4j/driver/internal/async/UnmanagedTransaction.java index e4349676e8..2286772657 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/UnmanagedTransaction.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/UnmanagedTransaction.java @@ -51,6 +51,7 @@ import org.neo4j.driver.internal.cursor.RxResultCursor; import org.neo4j.driver.internal.messaging.BoltProtocol; import org.neo4j.driver.internal.spi.Connection; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; public class UnmanagedTransaction implements TerminationAwareStateLockingExecutor { private enum State { @@ -102,13 +103,23 @@ private enum State { private final CompletableFuture beginFuture = new CompletableFuture<>(); private final Logging logging; + private final ApiTelemetryWork apiTelemetryWork; + public UnmanagedTransaction( Connection connection, Consumer bookmarkConsumer, long fetchSize, NotificationConfig notificationConfig, + ApiTelemetryWork apiTelemetryWork, Logging logging) { - this(connection, bookmarkConsumer, fetchSize, new ResultCursorsHolder(), notificationConfig, logging); + this( + connection, + bookmarkConsumer, + fetchSize, + new ResultCursorsHolder(), + notificationConfig, + apiTelemetryWork, + logging); } protected UnmanagedTransaction( @@ -117,6 +128,7 @@ protected UnmanagedTransaction( long fetchSize, ResultCursorsHolder resultCursors, NotificationConfig notificationConfig, + ApiTelemetryWork apiTelemetryWork, Logging logging) { this.connection = connection; this.protocol = connection.protocol(); @@ -125,6 +137,7 @@ protected UnmanagedTransaction( this.fetchSize = fetchSize; this.notificationConfig = notificationConfig; this.logging = logging; + this.apiTelemetryWork = apiTelemetryWork; connection.bindTerminationAwareStateLockingExecutor(this); } @@ -132,6 +145,13 @@ protected UnmanagedTransaction( // flush = false is only supported for async mode with a single subsequent run public CompletionStage beginAsync( Set initialBookmarks, TransactionConfig config, String txType, boolean flush) { + + apiTelemetryWork.execute(connection, protocol).whenComplete((unused, throwable) -> { + if (throwable != null) { + beginFuture.completeExceptionally(throwable); + } + }); + protocol.beginTransaction(connection, initialBookmarks, config, txType, notificationConfig, logging, flush) .handle((ignore, beginError) -> { if (beginError != null) { diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtil.java b/driver/src/main/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtil.java index a8357dc282..250e85ea2b 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtil.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtil.java @@ -29,7 +29,7 @@ import org.neo4j.driver.internal.messaging.v42.BoltProtocolV42; import org.neo4j.driver.internal.messaging.v44.BoltProtocolV44; import org.neo4j.driver.internal.messaging.v5.BoltProtocolV5; -import org.neo4j.driver.internal.messaging.v53.BoltProtocolV53; +import org.neo4j.driver.internal.messaging.v54.BoltProtocolV54; public final class BoltProtocolUtil { public static final int BOLT_MAGIC_PREAMBLE = 0x6060B017; @@ -41,7 +41,7 @@ public final class BoltProtocolUtil { private static final ByteBuf HANDSHAKE_BUF = unreleasableBuffer(copyInt( BOLT_MAGIC_PREAMBLE, - BoltProtocolV53.VERSION.toIntRange(BoltProtocolV5.VERSION), + BoltProtocolV54.VERSION.toIntRange(BoltProtocolV5.VERSION), BoltProtocolV44.VERSION.toIntRange(BoltProtocolV42.VERSION), BoltProtocolV41.VERSION.toInt(), BoltProtocolV3.VERSION.toInt())) diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/connection/ChannelAttributes.java b/driver/src/main/java/org/neo4j/driver/internal/async/connection/ChannelAttributes.java index 070363726a..2c56383acb 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/connection/ChannelAttributes.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/connection/ChannelAttributes.java @@ -53,6 +53,8 @@ public final class ChannelAttributes { // configuration hints provided by the server private static final AttributeKey CONNECTION_READ_TIMEOUT = newInstance("connectionReadTimeout"); + private static final AttributeKey TELEMETRY_ENABLED = newInstance("telemetryEnabled"); + private ChannelAttributes() {} public static String connectionId(Channel channel) { @@ -174,6 +176,14 @@ public static void setAuthContext(Channel channel, AuthContext authContext) { setOnce(channel, AUTH_CONTEXT, authContext); } + public static void setTelemetryEnabled(Channel channel, Boolean telemetryEnabled) { + setOnce(channel, TELEMETRY_ENABLED, telemetryEnabled); + } + + public static Boolean telemetryEnabled(Channel channel) { + return Optional.ofNullable(get(channel, TELEMETRY_ENABLED)).orElse(false); + } + private static T get(Channel channel, AttributeKey key) { return channel.attr(key).get(); } diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/connection/DirectConnection.java b/driver/src/main/java/org/neo4j/driver/internal/async/connection/DirectConnection.java index ff3d01ff44..8b401ea9ee 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/connection/DirectConnection.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/connection/DirectConnection.java @@ -64,6 +64,11 @@ public void disableAutoRead() { delegate.disableAutoRead(); } + @Override + public boolean isTelemetryEnabled() { + return delegate.isTelemetryEnabled(); + } + @Override public void write(Message message, ResponseHandler handler) { delegate.write(message, handler); diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/connection/RoutingConnection.java b/driver/src/main/java/org/neo4j/driver/internal/async/connection/RoutingConnection.java index 77ee8d0a16..6ceab579fa 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/async/connection/RoutingConnection.java +++ b/driver/src/main/java/org/neo4j/driver/internal/async/connection/RoutingConnection.java @@ -128,6 +128,11 @@ public String impersonatedUser() { return impersonatedUser; } + @Override + public boolean isTelemetryEnabled() { + return delegate.isTelemetryEnabled(); + } + private RoutingResponseHandler newRoutingResponseHandler(ResponseHandler handler) { return new RoutingResponseHandler(handler, serverAddress(), accessMode, errorHandler); } diff --git a/driver/src/main/java/org/neo4j/driver/internal/handlers/HelloV51ResponseHandler.java b/driver/src/main/java/org/neo4j/driver/internal/handlers/HelloV51ResponseHandler.java index 22ba973503..0d167c6c79 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/handlers/HelloV51ResponseHandler.java +++ b/driver/src/main/java/org/neo4j/driver/internal/handlers/HelloV51ResponseHandler.java @@ -21,6 +21,7 @@ import static org.neo4j.driver.internal.async.connection.ChannelAttributes.setConnectionId; import static org.neo4j.driver.internal.async.connection.ChannelAttributes.setConnectionReadTimeout; import static org.neo4j.driver.internal.async.connection.ChannelAttributes.setServerAgent; +import static org.neo4j.driver.internal.async.connection.ChannelAttributes.setTelemetryEnabled; import static org.neo4j.driver.internal.util.MetadataExtractor.extractServer; import io.netty.channel.Channel; @@ -35,6 +36,7 @@ public class HelloV51ResponseHandler implements ResponseHandler { private static final String CONNECTION_ID_METADATA_KEY = "connection_id"; public static final String CONFIGURATION_HINTS_KEY = "hints"; public static final String CONNECTION_RECEIVE_TIMEOUT_SECONDS_KEY = "connection.recv_timeout_seconds"; + public static final String TELEMETRY_ENABLED_KEY = "telemetry.enabled"; private final Channel channel; private final CompletableFuture helloFuture; @@ -79,6 +81,10 @@ private void processConfigurationHints(Map metadata) { .get(CONNECTION_RECEIVE_TIMEOUT_SECONDS_KEY) .asLong()) .ifPresent(timeout -> setConnectionReadTimeout(channel, timeout)); + + getFromSupplierOrEmptyOnException( + () -> configurationHints.get(TELEMETRY_ENABLED_KEY).asBoolean(false)) + .ifPresent(telemetryEnabled -> setTelemetryEnabled(channel, telemetryEnabled)); } } diff --git a/driver/src/main/java/org/neo4j/driver/internal/handlers/TelemetryResponseHandler.java b/driver/src/main/java/org/neo4j/driver/internal/handlers/TelemetryResponseHandler.java new file mode 100644 index 0000000000..303ca09c0b --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/handlers/TelemetryResponseHandler.java @@ -0,0 +1,61 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.handlers; + +import static java.util.Objects.requireNonNull; + +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import org.neo4j.driver.Value; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; +import org.neo4j.driver.internal.spi.ResponseHandler; + +/** + * Handles {@link TelemetryMessage} responses. + * + */ +public class TelemetryResponseHandler implements ResponseHandler { + private final CompletableFuture future; + + /** + * Constructor + * + * @param future The future which will be resolved + */ + public TelemetryResponseHandler(CompletableFuture future) { + this.future = requireNonNull(future); + } + + @Override + public void onSuccess(Map metadata) { + future.complete(null); + } + + @Override + public void onFailure(Throwable error) { + throw new UnsupportedOperationException("Telemetry is not expected to receive failures.", error); + } + + @Override + public void onRecord(Value[] fields) { + throw new UnsupportedOperationException( + "Telemetry is not expected to receive records: " + Arrays.toString(fields)); + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/BoltProtocol.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/BoltProtocol.java index 5a0681c6da..e60090c64f 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/messaging/BoltProtocol.java +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/BoltProtocol.java @@ -50,6 +50,7 @@ import org.neo4j.driver.internal.messaging.v51.BoltProtocolV51; import org.neo4j.driver.internal.messaging.v52.BoltProtocolV52; import org.neo4j.driver.internal.messaging.v53.BoltProtocolV53; +import org.neo4j.driver.internal.messaging.v54.BoltProtocolV54; import org.neo4j.driver.internal.spi.Connection; public interface BoltProtocol { @@ -123,6 +124,14 @@ CompletionStage beginTransaction( */ CompletionStage rollbackTransaction(Connection connection); + /** + * Sends telemetry message to the server. + * + * @param api The api number. + * @return Promise of message be delivered + */ + CompletionStage telemetry(Connection connection, Integer api); + /** * Execute the given query in an auto-commit transaction, i.e. {@link Session#run(Query)}. * @@ -202,6 +211,8 @@ static BoltProtocol forVersion(BoltProtocolVersion version) { return BoltProtocolV52.INSTANCE; } else if (BoltProtocolV53.VERSION.equals(version)) { return BoltProtocolV53.INSTANCE; + } else if (BoltProtocolV54.VERSION.equals(version)) { + return BoltProtocolV54.INSTANCE; } throw new ClientException("Unknown protocol version: " + version); } diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/encode/TelemetryMessageEncoder.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/encode/TelemetryMessageEncoder.java new file mode 100644 index 0000000000..093b7f1949 --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/encode/TelemetryMessageEncoder.java @@ -0,0 +1,38 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.encode; + +import static org.neo4j.driver.internal.util.Preconditions.checkArgument; + +import java.io.IOException; +import org.neo4j.driver.Values; +import org.neo4j.driver.internal.messaging.Message; +import org.neo4j.driver.internal.messaging.MessageEncoder; +import org.neo4j.driver.internal.messaging.ValuePacker; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; + +public class TelemetryMessageEncoder implements MessageEncoder { + @Override + public void encode(Message message, ValuePacker packer) throws IOException { + checkArgument(message, TelemetryMessage.class); + var telemetryMessage = (TelemetryMessage) message; + packer.packStructHeader(1, TelemetryMessage.SIGNATURE); + packer.pack(Values.value(telemetryMessage.api())); + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/request/TelemetryMessage.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/request/TelemetryMessage.java new file mode 100644 index 0000000000..e1ad38324c --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/request/TelemetryMessage.java @@ -0,0 +1,41 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.request; + +import org.neo4j.driver.internal.messaging.Message; + +/** + * TELEMETRY message + * Sent by the client to inform which API is used. + * + * @param api the API identification on the protocol level + */ +public record TelemetryMessage(Integer api) implements Message { + public static final byte SIGNATURE = 0x54; + + @Override + public byte signature() { + return SIGNATURE; + } + + @Override + public String toString() { + return String.format("TELEMETRY %S", api); + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3.java index 5db50387e6..a21e815fb1 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3.java +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3.java @@ -216,6 +216,11 @@ public ResultCursorFactory runInUnmanagedTransaction( return buildResultCursorFactory(connection, query, (ignored) -> {}, tx, runMessage, fetchSize); } + @Override + public CompletionStage telemetry(Connection connection, Integer api) { + return CompletableFuture.completedStage(null); + } + protected ResultCursorFactory buildResultCursorFactory( Connection connection, Query query, diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/BoltProtocolV54.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/BoltProtocolV54.java new file mode 100644 index 0000000000..5beeb8e11b --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/BoltProtocolV54.java @@ -0,0 +1,52 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.v54; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import org.neo4j.driver.internal.handlers.TelemetryResponseHandler; +import org.neo4j.driver.internal.messaging.BoltProtocol; +import org.neo4j.driver.internal.messaging.BoltProtocolVersion; +import org.neo4j.driver.internal.messaging.MessageFormat; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; +import org.neo4j.driver.internal.messaging.v53.BoltProtocolV53; +import org.neo4j.driver.internal.spi.Connection; + +public class BoltProtocolV54 extends BoltProtocolV53 { + public static final BoltProtocolVersion VERSION = new BoltProtocolVersion(5, 4); + public static final BoltProtocol INSTANCE = new BoltProtocolV54(); + + @Override + public BoltProtocolVersion version() { + return VERSION; + } + + @Override + public CompletionStage telemetry(Connection connection, Integer api) { + var telemetry = new TelemetryMessage(api); + var future = new CompletableFuture(); + connection.write(telemetry, new TelemetryResponseHandler(future)); + return future; + } + + @Override + public MessageFormat createMessageFormat() { + return new MessageFormatV54(); + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/MessageFormatV54.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/MessageFormatV54.java new file mode 100644 index 0000000000..b191a15889 --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/MessageFormatV54.java @@ -0,0 +1,36 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.v54; + +import org.neo4j.driver.internal.messaging.MessageFormat; +import org.neo4j.driver.internal.messaging.v5.MessageReaderV5; +import org.neo4j.driver.internal.packstream.PackInput; +import org.neo4j.driver.internal.packstream.PackOutput; + +public class MessageFormatV54 implements MessageFormat { + @Override + public Writer newWriter(PackOutput output) { + return new MessageWriterV54(output); + } + + @Override + public Reader newReader(PackInput input) { + return new MessageReaderV5(input); + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/MessageWriterV54.java b/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/MessageWriterV54.java new file mode 100644 index 0000000000..45dee6bcf5 --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/messaging/v54/MessageWriterV54.java @@ -0,0 +1,80 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.v54; + +import java.util.Map; +import org.neo4j.driver.internal.messaging.AbstractMessageWriter; +import org.neo4j.driver.internal.messaging.MessageEncoder; +import org.neo4j.driver.internal.messaging.common.CommonValuePacker; +import org.neo4j.driver.internal.messaging.encode.BeginMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.CommitMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.DiscardMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.GoodbyeMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.HelloMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.LogoffMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.LogonMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.PullMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.ResetMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.RollbackMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.RouteV44MessageEncoder; +import org.neo4j.driver.internal.messaging.encode.RunWithMetadataMessageEncoder; +import org.neo4j.driver.internal.messaging.encode.TelemetryMessageEncoder; +import org.neo4j.driver.internal.messaging.request.BeginMessage; +import org.neo4j.driver.internal.messaging.request.CommitMessage; +import org.neo4j.driver.internal.messaging.request.DiscardMessage; +import org.neo4j.driver.internal.messaging.request.GoodbyeMessage; +import org.neo4j.driver.internal.messaging.request.HelloMessage; +import org.neo4j.driver.internal.messaging.request.LogoffMessage; +import org.neo4j.driver.internal.messaging.request.LogonMessage; +import org.neo4j.driver.internal.messaging.request.PullMessage; +import org.neo4j.driver.internal.messaging.request.ResetMessage; +import org.neo4j.driver.internal.messaging.request.RollbackMessage; +import org.neo4j.driver.internal.messaging.request.RouteMessage; +import org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; +import org.neo4j.driver.internal.packstream.PackOutput; +import org.neo4j.driver.internal.util.Iterables; + +public class MessageWriterV54 extends AbstractMessageWriter { + public MessageWriterV54(PackOutput output) { + super(new CommonValuePacker(output, true), buildEncoders()); + } + + @SuppressWarnings("DuplicatedCode") + private static Map buildEncoders() { + Map result = Iterables.newHashMapWithSize(9); + result.put(HelloMessage.SIGNATURE, new HelloMessageEncoder()); + result.put(LogonMessage.SIGNATURE, new LogonMessageEncoder()); + result.put(LogoffMessage.SIGNATURE, new LogoffMessageEncoder()); + result.put(GoodbyeMessage.SIGNATURE, new GoodbyeMessageEncoder()); + result.put(RunWithMetadataMessage.SIGNATURE, new RunWithMetadataMessageEncoder()); + result.put(RouteMessage.SIGNATURE, new RouteV44MessageEncoder()); + + result.put(DiscardMessage.SIGNATURE, new DiscardMessageEncoder()); + result.put(PullMessage.SIGNATURE, new PullMessageEncoder()); + + result.put(BeginMessage.SIGNATURE, new BeginMessageEncoder()); + result.put(CommitMessage.SIGNATURE, new CommitMessageEncoder()); + result.put(RollbackMessage.SIGNATURE, new RollbackMessageEncoder()); + + result.put(ResetMessage.SIGNATURE, new ResetMessageEncoder()); + result.put(TelemetryMessage.SIGNATURE, new TelemetryMessageEncoder()); + return result; + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/reactive/AbstractReactiveSession.java b/driver/src/main/java/org/neo4j/driver/internal/reactive/AbstractReactiveSession.java index 95f7342f81..710181dfd4 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/reactive/AbstractReactiveSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/reactive/AbstractReactiveSession.java @@ -36,6 +36,8 @@ import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.async.UnmanagedTransaction; import org.neo4j.driver.internal.cursor.RxResultCursor; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.Futures; import org.neo4j.driver.reactive.RxResult; import org.neo4j.driver.reactivestreams.ReactiveResult; @@ -58,22 +60,24 @@ public AbstractReactiveSession(NetworkSession session) { protected abstract Publisher closeTransaction(S transaction, boolean commit); - Publisher doBeginTransaction(TransactionConfig config) { - return doBeginTransaction(config, null); + Publisher doBeginTransaction(TransactionConfig config, ApiTelemetryWork apiTelemetryWork) { + return doBeginTransaction(config, null, apiTelemetryWork); } @SuppressWarnings("DuplicatedCode") - protected Publisher doBeginTransaction(TransactionConfig config, String txType) { + protected Publisher doBeginTransaction( + TransactionConfig config, String txType, ApiTelemetryWork apiTelemetryWork) { return createSingleItemPublisher( () -> { var txFuture = new CompletableFuture(); - session.beginTransactionAsync(config, txType).whenComplete((tx, completionError) -> { - if (tx != null) { - txFuture.complete(createTransaction(tx)); - } else { - releaseConnectionBeforeReturning(txFuture, completionError); - } - }); + session.beginTransactionAsync(config, txType, apiTelemetryWork) + .whenComplete((tx, completionError) -> { + if (tx != null) { + txFuture.complete(createTransaction(tx)); + } else { + releaseConnectionBeforeReturning(txFuture, completionError); + } + }); return txFuture; }, () -> new IllegalStateException( @@ -82,17 +86,19 @@ protected Publisher doBeginTransaction(TransactionConfig config, String txTyp } @SuppressWarnings("DuplicatedCode") - private Publisher beginTransaction(AccessMode mode, TransactionConfig config) { + private Publisher beginTransaction( + AccessMode mode, TransactionConfig config, ApiTelemetryWork apiTelemetryWork) { return createSingleItemPublisher( () -> { var txFuture = new CompletableFuture(); - session.beginTransactionAsync(mode, config).whenComplete((tx, completionError) -> { - if (tx != null) { - txFuture.complete(createTransaction(tx)); - } else { - releaseConnectionBeforeReturning(txFuture, completionError); - } - }); + session.beginTransactionAsync(mode, config, apiTelemetryWork) + .whenComplete((tx, completionError) -> { + if (tx != null) { + txFuture.complete(createTransaction(tx)); + } else { + releaseConnectionBeforeReturning(txFuture, completionError); + } + }); return txFuture; }, () -> new IllegalStateException( @@ -122,8 +128,10 @@ protected Publisher runTransaction( } sink.next(value); })); + + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.MANAGED_TRANSACTION); var repeatableWork = Flux.usingWhen( - beginTransaction(mode, config), + beginTransaction(mode, config, apiTelemetryWork), work, tx -> closeTransaction(tx, true), (tx, error) -> closeTransaction(tx, false), diff --git a/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalReactiveSession.java b/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalReactiveSession.java index 67fb3736ac..a18eb70c3c 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalReactiveSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalReactiveSession.java @@ -30,6 +30,8 @@ import org.neo4j.driver.TransactionConfig; import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.async.UnmanagedTransaction; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.reactive.ReactiveResult; import org.neo4j.driver.reactive.ReactiveSession; import org.neo4j.driver.reactive.ReactiveTransaction; @@ -53,11 +55,12 @@ protected org.reactivestreams.Publisher closeTransaction(ReactiveTransacti @Override public Publisher beginTransaction(TransactionConfig config) { - return beginTransaction(config, null); + return beginTransaction(config, null, new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION)); } - public Publisher beginTransaction(TransactionConfig config, String txType) { - return publisherToFlowPublisher(doBeginTransaction(config, txType)); + public Publisher beginTransaction( + TransactionConfig config, String txType, ApiTelemetryWork apiTelemetryWork) { + return publisherToFlowPublisher(doBeginTransaction(config, txType, apiTelemetryWork)); } @Override diff --git a/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxSession.java b/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxSession.java index 7f5225daf7..201b52a988 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxSession.java @@ -29,6 +29,8 @@ import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.async.UnmanagedTransaction; import org.neo4j.driver.internal.cursor.RxResultCursor; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.Futures; import org.neo4j.driver.reactive.RxResult; import org.neo4j.driver.reactive.RxSession; @@ -55,7 +57,7 @@ protected Publisher closeTransaction(RxTransaction transaction, boolean co @Override public Publisher beginTransaction(TransactionConfig config) { - return doBeginTransaction(config); + return doBeginTransaction(config, new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION)); } @Override diff --git a/driver/src/main/java/org/neo4j/driver/internal/reactivestreams/InternalReactiveSession.java b/driver/src/main/java/org/neo4j/driver/internal/reactivestreams/InternalReactiveSession.java index 4e8c60e51d..97e3f08017 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/reactivestreams/InternalReactiveSession.java +++ b/driver/src/main/java/org/neo4j/driver/internal/reactivestreams/InternalReactiveSession.java @@ -27,6 +27,8 @@ import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.async.UnmanagedTransaction; import org.neo4j.driver.internal.reactive.AbstractReactiveSession; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.reactivestreams.ReactiveResult; import org.neo4j.driver.reactivestreams.ReactiveSession; import org.neo4j.driver.reactivestreams.ReactiveTransaction; @@ -51,11 +53,12 @@ public Publisher closeTransaction(ReactiveTransaction transaction, boolean @Override public Publisher beginTransaction(TransactionConfig config) { - return beginTransaction(config, null); + return beginTransaction(config, null, new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION)); } - public Publisher beginTransaction(TransactionConfig config, String txType) { - return doBeginTransaction(config, txType); + public Publisher beginTransaction( + TransactionConfig config, String txType, ApiTelemetryWork apiTelemetryWork) { + return doBeginTransaction(config, txType, apiTelemetryWork); } @Override diff --git a/driver/src/main/java/org/neo4j/driver/internal/spi/Connection.java b/driver/src/main/java/org/neo4j/driver/internal/spi/Connection.java index 1d07d9b8e3..ad772c18ff 100644 --- a/driver/src/main/java/org/neo4j/driver/internal/spi/Connection.java +++ b/driver/src/main/java/org/neo4j/driver/internal/spi/Connection.java @@ -39,6 +39,8 @@ public interface Connection { void writeAndFlush(Message message, ResponseHandler handler); + boolean isTelemetryEnabled(); + CompletionStage reset(Throwable throwable); CompletionStage release(); diff --git a/driver/src/main/java/org/neo4j/driver/internal/telemetry/ApiTelemetryWork.java b/driver/src/main/java/org/neo4j/driver/internal/telemetry/ApiTelemetryWork.java new file mode 100644 index 0000000000..4f7bb8c84a --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/telemetry/ApiTelemetryWork.java @@ -0,0 +1,84 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.telemetry; + +import static org.neo4j.driver.internal.util.Futures.futureCompletingConsumer; + +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.atomic.AtomicBoolean; +import org.neo4j.driver.internal.messaging.BoltProtocol; +import org.neo4j.driver.internal.spi.Connection; + +public class ApiTelemetryWork { + private final TelemetryApi telemetryApi; + private final AtomicBoolean completedWithSuccess; + + private final AtomicBoolean enabled; + + public ApiTelemetryWork(TelemetryApi telemetryApi) { + this.telemetryApi = telemetryApi; + this.completedWithSuccess = new AtomicBoolean(false); + this.enabled = new AtomicBoolean(true); + } + + public void setEnabled(boolean enabled) { + this.enabled.set(enabled); + } + + public CompletionStage execute(Connection connection, BoltProtocol protocol) { + var future = new CompletableFuture(); + if (connection.isTelemetryEnabled() && enabled.get() && !this.completedWithSuccess.get()) { + protocol.telemetry(connection, telemetryApi.getValue()) + .thenAccept((unused) -> completedWithSuccess.set(true)) + .whenComplete(futureCompletingConsumer(future)); + } else { + future.complete(null); + } + return future; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + var that = (ApiTelemetryWork) o; + return telemetryApi == that.telemetryApi + && Objects.equals(completedWithSuccess.get(), that.completedWithSuccess.get()) + && Objects.equals(enabled.get(), that.enabled.get()); + } + + @Override + public String toString() { + return "ApiTelemetryWork{" + "telemetryApi=" + + telemetryApi + ", completedWithSuccess=" + + completedWithSuccess.get() + ", enabled=" + + enabled.get() + '}'; + } + + @Override + public int hashCode() { + return Objects.hash(telemetryApi, completedWithSuccess, enabled); + } +} diff --git a/driver/src/main/java/org/neo4j/driver/internal/telemetry/TelemetryApi.java b/driver/src/main/java/org/neo4j/driver/internal/telemetry/TelemetryApi.java new file mode 100644 index 0000000000..f7511840dc --- /dev/null +++ b/driver/src/main/java/org/neo4j/driver/internal/telemetry/TelemetryApi.java @@ -0,0 +1,39 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.telemetry; + +/** + * An enum of valid telemetry metrics. + */ +public enum TelemetryApi { + MANAGED_TRANSACTION(0), + UNMANAGED_TRANSACTION(1), + AUTO_COMMIT_TRANSACTION(2), + EXECUTABLE_QUERY(3); + + private final Integer value; + + TelemetryApi(Integer value) { + this.value = value; + } + + public Integer getValue() { + return value; + } +} diff --git a/driver/src/test/java/org/neo4j/driver/ConfigTest.java b/driver/src/test/java/org/neo4j/driver/ConfigTest.java index d563940093..57168424b1 100644 --- a/driver/src/test/java/org/neo4j/driver/ConfigTest.java +++ b/driver/src/test/java/org/neo4j/driver/ConfigTest.java @@ -467,6 +467,7 @@ void shouldSerialize() throws Exception { .disableCategories( Set.of(NotificationCategory.UNSUPPORTED, NotificationCategory.UNRECOGNIZED)), config.notificationConfig()); + assertEquals(config.isTelemetryDisabled(), verify.isTelemetryDisabled()); } @Test @@ -505,4 +506,29 @@ void shouldHaveDefaultUserAgent() { assertTrue(config.userAgent().matches("^neo4j-java/.+$")); } + + @Test + void shouldDefaultToTelemetryEnabled() { + // Given + var config = Config.defaultConfig(); + + // When + var telemetryDisabled = config.isTelemetryDisabled(); + + // Then + assertFalse(telemetryDisabled); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void shouldChangeTelemetryDisabled(boolean disabled) { + // Given + var config = Config.builder().withTelemetryDisabled(disabled).build(); + + // When + var telemetryDisabled = config.isTelemetryDisabled(); + + // Then + assertEquals(disabled, telemetryDisabled); + } } diff --git a/driver/src/test/java/org/neo4j/driver/ParametersTest.java b/driver/src/test/java/org/neo4j/driver/ParametersTest.java index 02f3cd3c6e..c739766628 100644 --- a/driver/src/test/java/org/neo4j/driver/ParametersTest.java +++ b/driver/src/test/java/org/neo4j/driver/ParametersTest.java @@ -115,7 +115,8 @@ private Session mockedSession() { DEV_NULL_LOGGING, mock(BookmarkManager.class), null, - null); + null, + true); return new InternalSession(session); } } diff --git a/driver/src/test/java/org/neo4j/driver/integration/UnmanagedTransactionIT.java b/driver/src/test/java/org/neo4j/driver/integration/UnmanagedTransactionIT.java index 80c8341494..6836a10a50 100644 --- a/driver/src/test/java/org/neo4j/driver/integration/UnmanagedTransactionIT.java +++ b/driver/src/test/java/org/neo4j/driver/integration/UnmanagedTransactionIT.java @@ -49,6 +49,8 @@ import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.async.UnmanagedTransaction; import org.neo4j.driver.internal.security.SecurityPlanImpl; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.io.ChannelTrackingDriverFactory; import org.neo4j.driver.testutil.DatabaseExtension; import org.neo4j.driver.testutil.ParallelizableIT; @@ -76,7 +78,8 @@ private UnmanagedTransaction beginTransaction() { } private UnmanagedTransaction beginTransaction(NetworkSession session) { - return await(session.beginTransactionAsync(TransactionConfig.empty())); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + return await(session.beginTransactionAsync(TransactionConfig.empty(), apiTelemetryWork)); } private ResultCursor sessionRun(NetworkSession session, Query query) { @@ -169,12 +172,13 @@ void shouldFailToRunQueryWhenTerminated() { void shouldBePossibleToRunMoreTransactionsAfterOneIsTerminated() { var tx1 = beginTransaction(); tx1.markTerminated(null); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); // commit should fail, make session forget about this transaction and release the connection to the pool var e = assertThrows(TransactionTerminatedException.class, () -> await(tx1.commitAsync())); assertThat(e.getMessage(), startsWith("Transaction can't be committed")); - await(session.beginTransactionAsync(TransactionConfig.empty()) + await(session.beginTransactionAsync(TransactionConfig.empty(), apiTelemetryWork) .thenCompose(tx -> tx.runAsync(new Query("CREATE (:Node {id: 42})")) .thenCompose(ResultCursor::consumeAsync) .thenApply(ignore -> tx)) diff --git a/driver/src/test/java/org/neo4j/driver/integration/reactive/InternalReactiveSessionIT.java b/driver/src/test/java/org/neo4j/driver/integration/reactive/InternalReactiveSessionIT.java index d4c6cef63a..b965366f5c 100644 --- a/driver/src/test/java/org/neo4j/driver/integration/reactive/InternalReactiveSessionIT.java +++ b/driver/src/test/java/org/neo4j/driver/integration/reactive/InternalReactiveSessionIT.java @@ -29,6 +29,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.neo4j.driver.TransactionConfig; import org.neo4j.driver.internal.reactive.InternalReactiveSession; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.EnabledOnNeo4jWith; import org.neo4j.driver.reactive.ReactiveSession; import org.neo4j.driver.reactive.ReactiveTransaction; @@ -58,7 +60,8 @@ void setUp() { void shouldAcceptTxTypeWhenAvailable(String txType) { // GIVEN var txConfig = TransactionConfig.empty(); - var txMono = Mono.fromDirect(flowPublisherToFlux(session.beginTransaction(txConfig, txType))); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var txMono = Mono.fromDirect(flowPublisherToFlux(session.beginTransaction(txConfig, txType, apiTelemetryWork))); Function> txUnit = tx -> Mono.fromDirect(flowPublisherToFlux(tx.run("RETURN 1"))) .flatMap(result -> Mono.fromDirect(flowPublisherToFlux(result.consume()))); diff --git a/driver/src/test/java/org/neo4j/driver/internal/DriverFactoryTest.java b/driver/src/test/java/org/neo4j/driver/internal/DriverFactoryTest.java index 66548d9536..ca4e3aa12a 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/DriverFactoryTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/DriverFactoryTest.java @@ -119,7 +119,9 @@ void usesStandardSessionFactoryWhenNothingConfigured(String uri) { createDriver(uri, factory, config); var capturedFactory = factory.capturedSessionFactory; - assertThat(capturedFactory.newInstance(SessionConfig.defaultConfig(), null), instanceOf(NetworkSession.class)); + assertThat( + capturedFactory.newInstance(SessionConfig.defaultConfig(), null, true), + instanceOf(NetworkSession.class)); } @ParameterizedTest @@ -133,7 +135,7 @@ void usesLeakLoggingSessionFactoryWhenConfigured(String uri) { var capturedFactory = factory.capturedSessionFactory; assertThat( - capturedFactory.newInstance(SessionConfig.defaultConfig(), null), + capturedFactory.newInstance(SessionConfig.defaultConfig(), null, true), instanceOf(LeakLoggingNetworkSession.class)); } diff --git a/driver/src/test/java/org/neo4j/driver/internal/InternalDriverTest.java b/driver/src/test/java/org/neo4j/driver/internal/InternalDriverTest.java index 86b613a0e7..8f810da177 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/InternalDriverTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/InternalDriverTest.java @@ -133,7 +133,7 @@ void shouldCreateExecutableQuery() { private static InternalDriver newDriver(SessionFactory sessionFactory) { return new InternalDriver( - SecurityPlanImpl.insecure(), sessionFactory, DevNullMetricsProvider.INSTANCE, DEV_NULL_LOGGING); + SecurityPlanImpl.insecure(), sessionFactory, DevNullMetricsProvider.INSTANCE, true, DEV_NULL_LOGGING); } private static SessionFactory sessionFactoryMock() { @@ -150,6 +150,6 @@ private static InternalDriver newDriver(boolean isMetricsEnabled) { } var metricsProvider = DriverFactory.getOrCreateMetricsProvider(config, Clock.systemUTC()); - return new InternalDriver(SecurityPlanImpl.insecure(), sessionFactory, metricsProvider, DEV_NULL_LOGGING); + return new InternalDriver(SecurityPlanImpl.insecure(), sessionFactory, metricsProvider, true, DEV_NULL_LOGGING); } } diff --git a/driver/src/test/java/org/neo4j/driver/internal/InternalExecutableQueryTest.java b/driver/src/test/java/org/neo4j/driver/internal/InternalExecutableQueryTest.java index 7ba92c4072..d9240d4d9d 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/InternalExecutableQueryTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/InternalExecutableQueryTest.java @@ -49,6 +49,7 @@ import org.neo4j.driver.TransactionCallback; import org.neo4j.driver.TransactionConfig; import org.neo4j.driver.TransactionContext; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.summary.ResultSummary; class InternalExecutableQueryTest { @@ -131,7 +132,12 @@ void shouldExecuteAndReturnResult(RoutingControl routingControl) { given(driver.session(any(SessionConfig.class))).willReturn(session); var txContext = mock(TransactionContext.class); var accessMode = routingControl.equals(RoutingControl.WRITE) ? AccessMode.WRITE : AccessMode.READ; - given(session.execute(eq(accessMode), any(), eq(TransactionConfig.empty()), eq(false))) + given(session.execute( + eq(accessMode), + any(), + eq(TransactionConfig.empty()), + eq(TelemetryApi.EXECUTABLE_QUERY), + eq(false))) .willAnswer(answer -> { TransactionCallback txCallback = answer.getArgument(1); return txCallback.execute(txContext); @@ -181,7 +187,14 @@ var record = mock(Record.class); .withBookmarkManager(bookmarkManager) .build(); assertEquals(expectedSessionConfig, sessionConfig); - then(session).should().execute(eq(accessMode), any(), eq(TransactionConfig.empty()), eq(false)); + then(session) + .should() + .execute( + eq(accessMode), + any(), + eq(TransactionConfig.empty()), + eq(TelemetryApi.EXECUTABLE_QUERY), + eq(false)); then(txContext).should().run(query.withParameters(params)); then(result).should(times(2)).hasNext(); then(result).should().next(); diff --git a/driver/src/test/java/org/neo4j/driver/internal/InternalSessionTest.java b/driver/src/test/java/org/neo4j/driver/internal/InternalSessionTest.java index c7aa0e9052..22dd971a43 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/InternalSessionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/InternalSessionTest.java @@ -39,6 +39,8 @@ import org.neo4j.driver.internal.async.NetworkSession; import org.neo4j.driver.internal.async.UnmanagedTransaction; import org.neo4j.driver.internal.retry.RetryLogic; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; public class InternalSessionTest { NetworkSession networkSession; @@ -92,12 +94,13 @@ void shouldDelegateBeginWithType() { var internalSession = (InternalSession) session; var config = TransactionConfig.empty(); var type = "TYPE"; - given(networkSession.beginTransactionAsync(config, type)) + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + given(networkSession.beginTransactionAsync(config, type, apiTelemetryWork)) .willReturn(completedFuture(mock(UnmanagedTransaction.class))); internalSession.beginTransaction(config, type); - then(networkSession).should().beginTransactionAsync(config, type); + then(networkSession).should().beginTransactionAsync(config, type, apiTelemetryWork); } static List executeVariations() { diff --git a/driver/src/test/java/org/neo4j/driver/internal/SessionFactoryImplTest.java b/driver/src/test/java/org/neo4j/driver/internal/SessionFactoryImplTest.java index 5c7ffe2765..9fff0a3a8c 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/SessionFactoryImplTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/SessionFactoryImplTest.java @@ -39,11 +39,11 @@ void createsNetworkSessions() { var factory = newSessionFactory(config); var readSession = factory.newInstance( - builder().withDefaultAccessMode(AccessMode.READ).build(), null); + builder().withDefaultAccessMode(AccessMode.READ).build(), null, true); assertThat(readSession, instanceOf(NetworkSession.class)); var writeSession = factory.newInstance( - builder().withDefaultAccessMode(AccessMode.WRITE).build(), null); + builder().withDefaultAccessMode(AccessMode.WRITE).build(), null, true); assertThat(writeSession, instanceOf(NetworkSession.class)); } @@ -56,11 +56,11 @@ void createsLeakLoggingNetworkSessions() { var factory = newSessionFactory(config); var readSession = factory.newInstance( - builder().withDefaultAccessMode(AccessMode.READ).build(), null); + builder().withDefaultAccessMode(AccessMode.READ).build(), null, true); assertThat(readSession, instanceOf(LeakLoggingNetworkSession.class)); var writeSession = factory.newInstance( - builder().withDefaultAccessMode(AccessMode.WRITE).build(), null); + builder().withDefaultAccessMode(AccessMode.WRITE).build(), null, true); assertThat(writeSession, instanceOf(LeakLoggingNetworkSession.class)); } diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSessionTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSessionTest.java index 27160c8668..249172f805 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSessionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSessionTest.java @@ -43,6 +43,8 @@ import org.neo4j.driver.internal.handlers.pulln.FetchSizeUtil; import org.neo4j.driver.internal.spi.Connection; import org.neo4j.driver.internal.spi.ConnectionProvider; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.FixedRetryLogic; import org.neo4j.driver.testutil.TestUtil; @@ -67,7 +69,8 @@ void logsMessageWithStacktraceDuringFinalizationIfLeaked(TestInfo testInfo) thro when(logging.getLog(any(Class.class))).thenReturn(log); var session = newSession(logging, true); // begin transaction to make session obtain a connection - session.beginTransactionAsync(TransactionConfig.empty()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + session.beginTransactionAsync(TransactionConfig.empty(), apiTelemetryWork); finalize(session); @@ -103,7 +106,8 @@ private static LeakLoggingNetworkSession newSession(Logging logging, boolean ope logging, mock(BookmarkManager.class), null, - null); + null, + true); } private static ConnectionProvider connectionProviderMock(boolean openConnection) { diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/NetworkConnectionTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/NetworkConnectionTest.java index 3e8969f818..fea24af560 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/NetworkConnectionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/NetworkConnectionTest.java @@ -57,6 +57,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.neo4j.driver.Query; import org.neo4j.driver.exceptions.Neo4jException; @@ -548,6 +549,26 @@ void shouldDispatchingQueryMessagesWhenExecutorAbsent(QueryMessage queryMessage) assertEquals(1, channel.outboundMessages().size()); } + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void shouldReturnTelemetryEnabledWhenSet(Boolean telemetryEnabled) { + var channel = newChannel(); + ChannelAttributes.setTelemetryEnabled(channel, telemetryEnabled); + + var connection = newConnection(channel); + + assertEquals(telemetryEnabled, connection.isTelemetryEnabled()); + } + + @Test + void shouldReturnTelemetryEnabledEqualsFalseWhenNotSet() { + var channel = newChannel(); + + var connection = newConnection(channel); + + assertFalse(connection.isTelemetryEnabled()); + } + static List queryMessages() { return List.of( new QueryMessage(false, mock(RunWithMetadataMessage.class)), diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/NetworkSessionTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/NetworkSessionTest.java index f2dd5889e5..8b6256ac7e 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/NetworkSessionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/NetworkSessionTest.java @@ -28,6 +28,9 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.BDDMockito.given; +import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.inOrder; @@ -56,6 +59,8 @@ import java.util.concurrent.CompletableFuture; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.neo4j.driver.AccessMode; import org.neo4j.driver.Query; @@ -67,8 +72,12 @@ import org.neo4j.driver.internal.messaging.request.PullMessage; import org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage; import org.neo4j.driver.internal.messaging.v4.BoltProtocolV4; +import org.neo4j.driver.internal.messaging.v54.BoltProtocolV54; import org.neo4j.driver.internal.spi.Connection; import org.neo4j.driver.internal.spi.ConnectionProvider; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; +import org.neo4j.driver.internal.util.FixedRetryLogic; class NetworkSessionTest { private static final String DATABASE = "neo4j"; @@ -470,12 +479,60 @@ void shouldMarkTransactionAsTerminatedAndThenResetConnectionOnReset() { verify(connection).reset(any()); } + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void shouldSendTelemetryIfEnabledOnBegin(boolean telemetryDisabled) { + // given + var session = newSession(connectionProvider, WRITE, new FixedRetryLogic(0), Set.of(), telemetryDisabled); + given(connection.isTelemetryEnabled()).willReturn(true); + var protocol = spy(BoltProtocolV54.INSTANCE); + when(connection.protocol()).thenReturn(protocol); + + // when + beginTransaction(session); + + // then + if (telemetryDisabled) { + then(protocol).should(never()).telemetry(any(), any()); + } else { + then(protocol) + .should(times(1)) + .telemetry(eq(connection), eq(TelemetryApi.UNMANAGED_TRANSACTION.getValue())); + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void shouldSendTelemetryIfEnabledOnRun(boolean telemetryDisabled) { + // given + var query = "RETURN 1"; + setupSuccessfulRunAndPull(connection, query); + var apiTxWork = mock(ApiTelemetryWork.class); + var session = newSession(connectionProvider, WRITE, new FixedRetryLogic(0), Set.of(), telemetryDisabled); + given(connection.isTelemetryEnabled()).willReturn(true); + var protocol = spy(BoltProtocolV54.INSTANCE); + when(connection.protocol()).thenReturn(protocol); + + // when + run(session, query); + + // then + if (telemetryDisabled) { + then(protocol).should(never()).telemetry(any(), any()); + } else { + then(protocol) + .should(times(1)) + .telemetry(eq(connection), eq(TelemetryApi.AUTO_COMMIT_TRANSACTION.getValue())); + } + } + private static void run(NetworkSession session, String query) { await(session.runAsync(new Query(query), TransactionConfig.empty())); } private static UnmanagedTransaction beginTransaction(NetworkSession session) { - return await(session.beginTransactionAsync(TransactionConfig.empty())); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + return await(session.beginTransactionAsync(TransactionConfig.empty(), apiTelemetryWork)); } private static void close(NetworkSession session) { diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/UnmanagedTransactionTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/UnmanagedTransactionTest.java index 137f3d32b6..fa84399f73 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/UnmanagedTransactionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/UnmanagedTransactionTest.java @@ -20,6 +20,7 @@ import static java.util.Collections.emptyMap; import static java.util.concurrent.CompletableFuture.completedFuture; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -28,6 +29,9 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.BDDMockito.given; @@ -82,8 +86,11 @@ import org.neo4j.driver.internal.messaging.BoltProtocol; import org.neo4j.driver.internal.messaging.v4.BoltProtocolV4; import org.neo4j.driver.internal.messaging.v53.BoltProtocolV53; +import org.neo4j.driver.internal.messaging.v54.BoltProtocolV54; import org.neo4j.driver.internal.spi.Connection; import org.neo4j.driver.internal.spi.ResponseHandler; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; class UnmanagedTransactionTest { @Test @@ -181,7 +188,9 @@ void shouldBeClosedWhenMarkedTerminatedAndClosed() { void shouldReleaseConnectionWhenBeginFails() { var error = new RuntimeException("Wrong bookmark!"); var connection = connectionWithBegin(handler -> handler.onFailure(error)); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var bookmarks = Collections.singleton(InternalBookmark.parse("SomeBookmark")); var txConfig = TransactionConfig.empty(); @@ -195,7 +204,9 @@ void shouldReleaseConnectionWhenBeginFails() { @Test void shouldNotReleaseConnectionWhenBeginSucceeds() { var connection = connectionWithBegin(handler -> handler.onSuccess(emptyMap())); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var bookmarks = Collections.singleton(InternalBookmark.parse("SomeBookmark")); var txConfig = TransactionConfig.empty(); @@ -209,7 +220,9 @@ void shouldNotReleaseConnectionWhenBeginSucceeds() { @SuppressWarnings("ThrowableNotThrown") void shouldReleaseConnectionWhenTerminatedAndCommitted() { var connection = connectionMock(); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); tx.markTerminated(null); @@ -224,9 +237,17 @@ void shouldReleaseConnectionWhenTerminatedAndCommitted() { void shouldNotCreateCircularExceptionWhenTerminationCauseEqualsToCursorFailure() { var connection = connectionMock(); var terminationCause = new ClientException("Custom exception"); + + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); var resultCursorsHolder = mockResultCursorWith(terminationCause); var tx = new UnmanagedTransaction( - connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, resultCursorsHolder, null, Logging.none()); + connection, + (ignored) -> {}, + UNLIMITED_FETCH_SIZE, + resultCursorsHolder, + null, + apiTelemetryWork, + Logging.none()); tx.markTerminated(terminationCause); @@ -241,8 +262,15 @@ void shouldNotCreateCircularExceptionWhenTerminationCauseDifferentFromCursorFail var connection = connectionMock(); var terminationCause = new ClientException("Custom exception"); var resultCursorsHolder = mockResultCursorWith(new ClientException("Cursor error")); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); var tx = new UnmanagedTransaction( - connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, resultCursorsHolder, null, Logging.none()); + connection, + (ignored) -> {}, + UNLIMITED_FETCH_SIZE, + resultCursorsHolder, + null, + apiTelemetryWork, + Logging.none()); tx.markTerminated(terminationCause); @@ -259,7 +287,9 @@ void shouldNotCreateCircularExceptionWhenTerminationCauseDifferentFromCursorFail void shouldNotCreateCircularExceptionWhenTerminatedWithoutFailure() { var connection = connectionMock(); var terminationCause = new ClientException("Custom exception"); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); tx.markTerminated(terminationCause); @@ -273,7 +303,9 @@ void shouldNotCreateCircularExceptionWhenTerminatedWithoutFailure() { @SuppressWarnings("ThrowableNotThrown") void shouldReleaseConnectionWhenTerminatedAndRolledBack() { var connection = connectionMock(); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); tx.markTerminated(null); await(tx.rollbackAsync()); @@ -284,7 +316,9 @@ void shouldReleaseConnectionWhenTerminatedAndRolledBack() { @Test void shouldReleaseConnectionWhenClose() { var connection = connectionMock(); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); await(tx.closeAsync()); @@ -295,7 +329,9 @@ void shouldReleaseConnectionWhenClose() { void shouldReleaseConnectionOnConnectionAuthorizationExpiredExceptionFailure() { var exception = new AuthorizationExpiredException("code", "message"); var connection = connectionWithBegin(handler -> handler.onFailure(exception)); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var bookmarks = Collections.singleton(InternalBookmark.parse("SomeBookmark")); var txConfig = TransactionConfig.empty(); @@ -310,7 +346,9 @@ void shouldReleaseConnectionOnConnectionAuthorizationExpiredExceptionFailure() { @Test void shouldReleaseConnectionOnConnectionReadTimeoutExceptionFailure() { var connection = connectionWithBegin(handler -> handler.onFailure(ConnectionReadTimeoutException.INSTANCE)); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var bookmarks = Collections.singleton(InternalBookmark.parse("SomeBookmark")); var txConfig = TransactionConfig.empty(); @@ -340,7 +378,9 @@ void shouldReturnExistingStageOnSimilarCompletingAction( given(connection.protocol()).willReturn(protocol); given(protocolCommit ? protocol.commitTransaction(connection) : protocol.rollbackTransaction(connection)) .willReturn(new CompletableFuture<>()); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var initialStage = mapTransactionAction(initialAction, tx).get(); var similarStage = mapTransactionAction(similarAction, tx).get(); @@ -380,7 +420,9 @@ void shouldReturnFailingStageOnConflictingCompletingAction( given(connection.protocol()).willReturn(protocol); given(protocolCommit ? protocol.commitTransaction(connection) : protocol.rollbackTransaction(connection)) .willReturn(protocolActionCompleted ? completedFuture(null) : new CompletableFuture<>()); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var originalActionStage = mapTransactionAction(initialAction, tx).get(); var conflictingActionStage = mapTransactionAction(conflictingAction, tx).get(); @@ -421,7 +463,9 @@ void shouldReturnCompletedWithNullStageOnClosingInactiveTransactionExceptCommitt given(connection.protocol()).willReturn(protocol); given(protocolCommit ? protocol.commitTransaction(connection) : protocol.rollbackTransaction(connection)) .willReturn(completedFuture(null)); - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); var originalActionStage = mapTransactionAction(originalAction, tx).get(); var closeStage = commitOnClose != null ? tx.closeAsync(commitOnClose) : tx.closeAsync(); @@ -507,6 +551,106 @@ void shouldThrowOnRunningNewQueriesWhenTransactionIsClosing(TransactionClosingTe assertEquals(testParams.expectedMessage(), exception.getMessage()); } + @Test + void shouldBeginAsyncTelemetryNotCompleteReturnedFuture() { + var protocol = mock(BoltProtocol.class); + given(protocol.version()).willReturn(BoltProtocolV54.VERSION); + var connection = connectionMock(protocol); + var apiTelemetryWork = mock(ApiTelemetryWork.class); + var beginFuture = new CompletableFuture<>(); + doReturn(CompletableFuture.completedFuture(null)).when(apiTelemetryWork).execute(connection, protocol); + doReturn(beginFuture) + .when(protocol) + .beginTransaction(any(), anySet(), any(), anyString(), any(), any(), anyBoolean()); + var unmanagedTransaction = new UnmanagedTransaction(connection, (bm) -> {}, 100, null, apiTelemetryWork, null); + + assertFalse(unmanagedTransaction + .beginAsync(Set.of(), TransactionConfig.empty(), "tx", true) + .toCompletableFuture() + .isDone()); + + beginFuture.complete(null); + + assertTrue(unmanagedTransaction + .beginAsync(Set.of(), TransactionConfig.empty(), "tx", true) + .toCompletableFuture() + .isDone()); + } + + @Test + void shouldBeginAsyncThrowErrorOnTelemetryIfFlushIsTrueAndBeginDontFinish() { + var protocol = mock(BoltProtocol.class); + given(protocol.version()).willReturn(BoltProtocolV54.VERSION); + var connection = connectionMock(protocol); + var apiTelemetryWork = mock(ApiTelemetryWork.class); + doReturn(CompletableFuture.failedFuture(new SecurityException("My Exception"))) + .when(apiTelemetryWork) + .execute(connection, protocol); + doReturn(new CompletableFuture<>()) + .when(protocol) + .beginTransaction(any(), anySet(), any(), anyString(), any(), any(), anyBoolean()); + var unmanagedTransaction = new UnmanagedTransaction(connection, (bm) -> {}, 100, null, apiTelemetryWork, null); + + assertThrows( + SecurityException.class, + () -> await(unmanagedTransaction.beginAsync(Set.of(), TransactionConfig.empty(), "tx", true))); + } + + @Test + void shouldBeginAsyncThrowErrorOnTelemetryIfFlushIsTrueAndBeginFailed() { + var protocol = mock(BoltProtocol.class); + given(protocol.version()).willReturn(BoltProtocolV54.VERSION); + var connection = connectionMock(protocol); + var apiTelemetryWork = mock(ApiTelemetryWork.class); + doReturn(CompletableFuture.failedFuture(new SecurityException("My Exception"))) + .when(apiTelemetryWork) + .execute(connection, protocol); + doReturn(CompletableFuture.failedFuture(new ClientException("other error"))) + .when(protocol) + .beginTransaction(any(), anySet(), any(), anyString(), any(), any(), anyBoolean()); + var unmanagedTransaction = new UnmanagedTransaction(connection, (bm) -> {}, 100, null, apiTelemetryWork, null); + + assertThrows( + SecurityException.class, + () -> await(unmanagedTransaction.beginAsync(Set.of(), TransactionConfig.empty(), "tx", true))); + } + + @Test + void shouldBeginAsyncNotThrowErrorOnTelemetryIfNotFlushIsTrueAndBeginDontFinish() { + var protocol = mock(BoltProtocol.class); + given(protocol.version()).willReturn(BoltProtocolV54.VERSION); + var connection = connectionMock(protocol); + var apiTelemetryWork = mock(ApiTelemetryWork.class); + doReturn(CompletableFuture.failedFuture(new SecurityException("My Exception"))) + .when(apiTelemetryWork) + .execute(connection, protocol); + doReturn(new CompletableFuture<>()) + .when(protocol) + .beginTransaction(any(), anySet(), any(), anyString(), any(), any(), anyBoolean()); + var unmanagedTransaction = new UnmanagedTransaction(connection, (bm) -> {}, 100, null, apiTelemetryWork, null); + + assertDoesNotThrow( + () -> await(unmanagedTransaction.beginAsync(Set.of(), TransactionConfig.empty(), "tx", false))); + } + + @Test + void shouldBeginAsyncNotThrowErrorOnTelemetryIfNotFlushIsTrueAndBeginFailed() { + var protocol = mock(BoltProtocol.class); + given(protocol.version()).willReturn(BoltProtocolV54.VERSION); + var connection = connectionMock(protocol); + var apiTelemetryWork = mock(ApiTelemetryWork.class); + doReturn(CompletableFuture.failedFuture(new SecurityException("My Exception"))) + .when(apiTelemetryWork) + .execute(connection, protocol); + doReturn(CompletableFuture.failedFuture(new ClientException("other error"))) + .when(protocol) + .beginTransaction(any(), anySet(), any(), anyString(), any(), any(), anyBoolean()); + var unmanagedTransaction = new UnmanagedTransaction(connection, (bm) -> {}, 100, null, apiTelemetryWork, null); + + assertDoesNotThrow( + () -> await(unmanagedTransaction.beginAsync(Set.of(), TransactionConfig.empty(), "tx", false))); + } + static List transactionClosingTestParams() { Function> asyncRun = tx -> tx.runAsync(new Query("query")); Function> reactiveRun = tx -> tx.runRx(new Query("query")); @@ -559,7 +703,9 @@ private static UnmanagedTransaction beginTx(Connection connection) { } private static UnmanagedTransaction beginTx(Connection connection, Set initialBookmarks) { - var tx = new UnmanagedTransaction(connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, Logging.none()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var tx = new UnmanagedTransaction( + connection, (ignored) -> {}, UNLIMITED_FETCH_SIZE, null, apiTelemetryWork, Logging.none()); return await(tx.beginAsync(initialBookmarks, TransactionConfig.empty(), null, true)); } diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtilTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtilTest.java index 2c2f199a2c..106007d12b 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtilTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/connection/BoltProtocolUtilTest.java @@ -32,7 +32,7 @@ import org.neo4j.driver.internal.messaging.v3.BoltProtocolV3; import org.neo4j.driver.internal.messaging.v41.BoltProtocolV41; import org.neo4j.driver.internal.messaging.v44.BoltProtocolV44; -import org.neo4j.driver.internal.messaging.v53.BoltProtocolV53; +import org.neo4j.driver.internal.messaging.v54.BoltProtocolV54; class BoltProtocolUtilTest { @Test @@ -40,7 +40,7 @@ void shouldReturnHandshakeBuf() { assertByteBufContains( handshakeBuf(), BOLT_MAGIC_PREAMBLE, - (3 << 16) | BoltProtocolV53.VERSION.toInt(), + (4 << 16) | BoltProtocolV54.VERSION.toInt(), (2 << 16) | BoltProtocolV44.VERSION.toInt(), BoltProtocolV41.VERSION.toInt(), BoltProtocolV3.VERSION.toInt()); @@ -48,7 +48,7 @@ void shouldReturnHandshakeBuf() { @Test void shouldReturnHandshakeString() { - assertEquals("[0x6060b017, 197381, 132100, 260, 3]", handshakeString()); + assertEquals("[0x6060b017, 263173, 132100, 260, 3]", handshakeString()); } @Test diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/connection/DirectConnectionTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/connection/DirectConnectionTest.java index 447ec6fa40..28c3e8a98e 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/connection/DirectConnectionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/connection/DirectConnectionTest.java @@ -21,11 +21,14 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.BDDMockito.given; import static org.mockito.BDDMockito.then; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.neo4j.driver.AccessMode.READ; import static org.neo4j.driver.internal.DatabaseNameUtil.defaultDatabase; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.neo4j.driver.internal.spi.Connection; public class DirectConnectionTest { @@ -44,4 +47,15 @@ void shouldReturnServerAgent() { assertEquals(agent, actualAgent); then(connection).should().serverAgent(); } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void shouldReturnTelemetryEnabledReturnNetworkValue(Boolean telemetryEnabled) { + var connection = mock(Connection.class); + doReturn(telemetryEnabled).when(connection).isTelemetryEnabled(); + + var directConnection = new DirectConnection(connection, defaultDatabase(), READ, null); + + assertEquals(telemetryEnabled, directConnection.isTelemetryEnabled()); + } } diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/connection/RoutingConnectionTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/connection/RoutingConnectionTest.java index da25b508fd..7d159507ed 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/async/connection/RoutingConnectionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/async/connection/RoutingConnectionTest.java @@ -23,6 +23,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.BDDMockito.given; import static org.mockito.BDDMockito.then; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -31,6 +32,8 @@ import static org.neo4j.driver.internal.messaging.request.PullAllMessage.PULL_ALL; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.neo4j.driver.internal.RoutingErrorHandler; import org.neo4j.driver.internal.handlers.RoutingResponseHandler; @@ -65,6 +68,18 @@ void shouldReturnServerAgent() { then(connection).should().serverAgent(); } + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void shouldReturnTelemetryEnabledReturnNetworkValue(Boolean telemetryEnabled) { + var connection = mock(Connection.class); + var errorHandler = mock(RoutingErrorHandler.class); + doReturn(telemetryEnabled).when(connection).isTelemetryEnabled(); + + var routingConnection = new RoutingConnection(connection, defaultDatabase(), READ, null, errorHandler); + + assertEquals(telemetryEnabled, routingConnection.isTelemetryEnabled()); + } + private static void testHandlersWrappingWithSingleMessage(boolean flush) { var connection = mock(Connection.class); var errorHandler = mock(RoutingErrorHandler.class); diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/encode/TelemetryMessageEncoderTest.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/encode/TelemetryMessageEncoderTest.java new file mode 100644 index 0000000000..2d13029ca6 --- /dev/null +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/encode/TelemetryMessageEncoderTest.java @@ -0,0 +1,59 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.encode; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.stream.Stream; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.neo4j.driver.Query; +import org.neo4j.driver.Values; +import org.neo4j.driver.internal.messaging.ValuePacker; +import org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; +import org.neo4j.driver.internal.telemetry.TelemetryApi; + +class TelemetryMessageEncoderTest { + private final TelemetryMessageEncoder encoder = new TelemetryMessageEncoder(); + private final ValuePacker packer = mock(ValuePacker.class); + + @ParameterizedTest + @MethodSource("validApis") + void shouldEncodeTelemetryMessage(int api) throws Exception { + encoder.encode(new TelemetryMessage(api), packer); + + verify(packer).packStructHeader(1, TelemetryMessage.SIGNATURE); + verify(packer).pack(Values.value(api)); + } + + @Test + void shouldFailToEncodeWrongMessage() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> encoder.encode(RunWithMetadataMessage.unmanagedTxRunMessage(new Query("RETURN 2")), packer)); + } + + private static Stream validApis() { + return Stream.of(TelemetryApi.values()).map(TelemetryApi::getValue); + } +} diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3Test.java index b78a938cb1..f6fb203438 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v3/BoltProtocolV3Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -61,6 +62,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -400,6 +402,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { assertThat(e.getMessage(), startsWith("Database name parameter for selecting database is not supported")); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + protected void testDatabaseNameSupport(boolean autoCommitTx) { ClientException e; if (autoCommitTx) { diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v4/BoltProtocolV4Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v4/BoltProtocolV4Test.java index 79923055fc..21a7474f64 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v4/BoltProtocolV4Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v4/BoltProtocolV4Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -389,6 +391,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + @SuppressWarnings("SameReturnValue") private BoltProtocol createProtocol() { return BoltProtocolV4.INSTANCE; diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v41/BoltProtocolV41Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v41/BoltProtocolV41Test.java index f8bef52b4b..9bc29cc2a8 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v41/BoltProtocolV41Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v41/BoltProtocolV41Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -394,6 +396,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV4.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v42/BoltProtocolV42Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v42/BoltProtocolV42Test.java index 157d80e490..4eb79834ab 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v42/BoltProtocolV42Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v42/BoltProtocolV42Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -394,6 +396,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV4.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v43/BoltProtocolV43Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v43/BoltProtocolV43Test.java index 8511797fcb..a5cf7cae56 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v43/BoltProtocolV43Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v43/BoltProtocolV43Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -393,6 +395,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV43.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v44/BoltProtocolV44Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v44/BoltProtocolV44Test.java index 78c405c6f9..c63251ff79 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v44/BoltProtocolV44Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v44/BoltProtocolV44Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -393,6 +395,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV44.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v5/BoltProtocolV5Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v5/BoltProtocolV5Test.java index d2e6aef764..ad974ed2b4 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v5/BoltProtocolV5Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v5/BoltProtocolV5Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -393,6 +395,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV5.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v51/BoltProtocolV51Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v51/BoltProtocolV51Test.java index 8a3b646089..240881fc1c 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v51/BoltProtocolV51Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v51/BoltProtocolV51Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -365,6 +367,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV51.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v52/BoltProtocolV52Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v52/BoltProtocolV52Test.java index 1de37e092d..68c3835b55 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v52/BoltProtocolV52Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v52/BoltProtocolV52Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -366,6 +368,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV51.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v53/BoltProtocolV53Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v53/BoltProtocolV53Test.java index 699bce3832..18261dbb23 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/messaging/v53/BoltProtocolV53Test.java +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v53/BoltProtocolV53Test.java @@ -37,6 +37,7 @@ import static org.mockito.BDDMockito.then; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -62,6 +63,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.neo4j.driver.AccessMode; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Bookmark; @@ -373,6 +375,16 @@ void shouldNotSupportDatabaseNameForAutoCommitTransactions() { Logging.none())); } + @Test + void shouldTelemetryReturnCompletedStageWithoutSendAnyMessage() { + var connection = connectionMock(); + + await(protocol.telemetry(connection, 1)); + + verify(connection, never()).write(Mockito.any(), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + private Class expectedMessageFormatType() { return MessageFormatV51.class; } diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/BoltProtocolV54Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/BoltProtocolV54Test.java new file mode 100644 index 0000000000..1391a065f6 --- /dev/null +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/BoltProtocolV54Test.java @@ -0,0 +1,601 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.v54; + +import static java.time.Duration.ofSeconds; +import static java.util.Collections.emptyMap; +import static java.util.Collections.singletonMap; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasSize; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.BDDMockito.then; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.neo4j.driver.AccessMode.WRITE; +import static org.neo4j.driver.Values.value; +import static org.neo4j.driver.internal.DatabaseNameUtil.database; +import static org.neo4j.driver.internal.DatabaseNameUtil.defaultDatabase; +import static org.neo4j.driver.internal.handlers.pulln.FetchSizeUtil.UNLIMITED_FETCH_SIZE; +import static org.neo4j.driver.testutil.TestUtil.await; +import static org.neo4j.driver.testutil.TestUtil.connectionMock; + +import io.netty.channel.embedded.EmbeddedChannel; +import java.time.Clock; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletionStage; +import java.util.function.Consumer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; +import org.neo4j.driver.AccessMode; +import org.neo4j.driver.AuthTokens; +import org.neo4j.driver.Bookmark; +import org.neo4j.driver.Logging; +import org.neo4j.driver.Query; +import org.neo4j.driver.TransactionConfig; +import org.neo4j.driver.Value; +import org.neo4j.driver.internal.BoltAgentUtil; +import org.neo4j.driver.internal.DatabaseBookmark; +import org.neo4j.driver.internal.DatabaseName; +import org.neo4j.driver.internal.InternalBookmark; +import org.neo4j.driver.internal.async.UnmanagedTransaction; +import org.neo4j.driver.internal.async.connection.ChannelAttributes; +import org.neo4j.driver.internal.async.inbound.InboundMessageDispatcher; +import org.neo4j.driver.internal.cluster.RoutingContext; +import org.neo4j.driver.internal.cursor.AsyncResultCursor; +import org.neo4j.driver.internal.handlers.BeginTxResponseHandler; +import org.neo4j.driver.internal.handlers.CommitTxResponseHandler; +import org.neo4j.driver.internal.handlers.PullAllResponseHandler; +import org.neo4j.driver.internal.handlers.RollbackTxResponseHandler; +import org.neo4j.driver.internal.handlers.RunResponseHandler; +import org.neo4j.driver.internal.handlers.TelemetryResponseHandler; +import org.neo4j.driver.internal.messaging.BoltProtocol; +import org.neo4j.driver.internal.messaging.MessageFormat; +import org.neo4j.driver.internal.messaging.request.BeginMessage; +import org.neo4j.driver.internal.messaging.request.CommitMessage; +import org.neo4j.driver.internal.messaging.request.GoodbyeMessage; +import org.neo4j.driver.internal.messaging.request.PullMessage; +import org.neo4j.driver.internal.messaging.request.RollbackMessage; +import org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; +import org.neo4j.driver.internal.security.InternalAuthToken; +import org.neo4j.driver.internal.spi.Connection; +import org.neo4j.driver.internal.spi.ResponseHandler; + +public class BoltProtocolV54Test { + protected static final String QUERY_TEXT = "RETURN $x"; + protected static final Map PARAMS = singletonMap("x", value(42)); + protected static final Query QUERY = new Query(QUERY_TEXT, value(PARAMS)); + + protected final BoltProtocol protocol = createProtocol(); + private final EmbeddedChannel channel = new EmbeddedChannel(); + private final InboundMessageDispatcher messageDispatcher = new InboundMessageDispatcher(channel, Logging.none()); + + private final TransactionConfig txConfig = TransactionConfig.builder() + .withTimeout(ofSeconds(12)) + .withMetadata(singletonMap("key", value(42))) + .build(); + + @SuppressWarnings("SameReturnValue") + protected BoltProtocol createProtocol() { + return BoltProtocolV54.INSTANCE; + } + + @BeforeEach + void beforeEach() { + ChannelAttributes.setMessageDispatcher(channel, messageDispatcher); + } + + @AfterEach + void afterEach() { + channel.finishAndReleaseAll(); + } + + @Test + void shouldCreateMessageFormat() { + assertThat(protocol.createMessageFormat(), instanceOf(expectedMessageFormatType())); + } + + @Test + void shouldInitializeChannel() { + var promise = channel.newPromise(); + + protocol.initializeChannel( + "MyDriver/0.0.1", + BoltAgentUtil.VALUE, + dummyAuthToken(), + RoutingContext.EMPTY, + promise, + null, + mock(Clock.class)); + + assertThat(channel.outboundMessages(), hasSize(0)); + assertEquals(1, messageDispatcher.queuedHandlersCount()); + assertTrue(promise.isDone()); + + Map metadata = new HashMap<>(); + metadata.put("server", value("Neo4j/4.4.0")); + metadata.put("connection_id", value("bolt-42")); + + messageDispatcher.handleSuccessMessage(metadata); + + channel.flush(); + assertTrue(promise.isDone()); + assertTrue(promise.isSuccess()); + } + + @Test + void shouldPrepareToCloseChannel() { + protocol.prepareToCloseChannel(channel); + + assertThat(channel.outboundMessages(), hasSize(1)); + assertThat(channel.outboundMessages().poll(), instanceOf(GoodbyeMessage.class)); + assertEquals(1, messageDispatcher.queuedHandlersCount()); + } + + @Test + void shouldBeginTransactionWithoutBookmark() { + var connection = connectionMock(protocol); + + var stage = protocol.beginTransaction( + connection, Collections.emptySet(), TransactionConfig.empty(), null, null, Logging.none(), true); + + verify(connection) + .writeAndFlush( + eq(new BeginMessage( + Collections.emptySet(), + TransactionConfig.empty(), + defaultDatabase(), + WRITE, + null, + null, + null, + Logging.none())), + any(BeginTxResponseHandler.class)); + assertNull(await(stage)); + } + + @Test + void shouldBeginTransactionWithBookmarks() { + var connection = connectionMock(protocol); + var bookmarks = Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx100")); + + var stage = protocol.beginTransaction( + connection, bookmarks, TransactionConfig.empty(), null, null, Logging.none(), true); + + verify(connection) + .writeAndFlush( + eq(new BeginMessage( + bookmarks, + TransactionConfig.empty(), + defaultDatabase(), + WRITE, + null, + null, + null, + Logging.none())), + any(BeginTxResponseHandler.class)); + assertNull(await(stage)); + } + + @Test + void shouldBeginTransactionWithConfig() { + var connection = connectionMock(protocol); + + var stage = protocol.beginTransaction( + connection, Collections.emptySet(), txConfig, null, null, Logging.none(), true); + + verify(connection) + .writeAndFlush( + eq(new BeginMessage( + Collections.emptySet(), + txConfig, + defaultDatabase(), + WRITE, + null, + null, + null, + Logging.none())), + any(BeginTxResponseHandler.class)); + assertNull(await(stage)); + } + + @Test + void shouldBeginTransactionWithBookmarksAndConfig() { + var connection = connectionMock(protocol); + var bookmarks = Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx4242")); + + var stage = protocol.beginTransaction(connection, bookmarks, txConfig, null, null, Logging.none(), true); + + verify(connection) + .writeAndFlush( + eq(new BeginMessage( + bookmarks, txConfig, defaultDatabase(), WRITE, null, null, null, Logging.none())), + any(BeginTxResponseHandler.class)); + assertNull(await(stage)); + } + + @Test + void shouldCommitTransaction() { + var bookmarkString = "neo4j:bookmark:v1:tx4242"; + + var connection = connectionMock(protocol); + when(connection.protocol()).thenReturn(protocol); + doAnswer(invocation -> { + ResponseHandler commitHandler = invocation.getArgument(1); + commitHandler.onSuccess(singletonMap("bookmark", value(bookmarkString))); + return null; + }) + .when(connection) + .writeAndFlush(eq(CommitMessage.COMMIT), any()); + + var stage = protocol.commitTransaction(connection); + + verify(connection).writeAndFlush(eq(CommitMessage.COMMIT), any(CommitTxResponseHandler.class)); + assertEquals(InternalBookmark.parse(bookmarkString), await(stage).bookmark()); + } + + @Test + void shouldRollbackTransaction() { + var connection = connectionMock(protocol); + + var stage = protocol.rollbackTransaction(connection); + + verify(connection).writeAndFlush(eq(RollbackMessage.ROLLBACK), any(RollbackTxResponseHandler.class)); + assertNull(await(stage)); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInAutoCommitTransactionAndWaitForRunResponse(AccessMode mode) throws Exception { + testRunAndWaitForRunResponse(true, TransactionConfig.empty(), mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInAutoCommitWithConfigTransactionAndWaitForRunResponse(AccessMode mode) throws Exception { + testRunAndWaitForRunResponse(true, txConfig, mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInAutoCommitTransactionAndWaitForSuccessRunResponse(AccessMode mode) throws Exception { + testSuccessfulRunInAutoCommitTxWithWaitingForResponse(Collections.emptySet(), TransactionConfig.empty(), mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInAutoCommitTransactionWithBookmarkAndConfigAndWaitForSuccessRunResponse(AccessMode mode) + throws Exception { + testSuccessfulRunInAutoCommitTxWithWaitingForResponse( + Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx65")), txConfig, mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInAutoCommitTransactionAndWaitForFailureRunResponse(AccessMode mode) { + testFailedRunInAutoCommitTxWithWaitingForResponse(Collections.emptySet(), TransactionConfig.empty(), mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInAutoCommitTransactionWithBookmarkAndConfigAndWaitForFailureRunResponse(AccessMode mode) { + testFailedRunInAutoCommitTxWithWaitingForResponse( + Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx163")), txConfig, mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInUnmanagedTransactionAndWaitForRunResponse(AccessMode mode) throws Exception { + testRunAndWaitForRunResponse(false, TransactionConfig.empty(), mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInUnmanagedTransactionAndWaitForSuccessRunResponse(AccessMode mode) throws Exception { + testRunInUnmanagedTransactionAndWaitForRunResponse(true, mode); + } + + @ParameterizedTest + @EnumSource(AccessMode.class) + void shouldRunInUnmanagedTransactionAndWaitForFailureRunResponse(AccessMode mode) throws Exception { + testRunInUnmanagedTransactionAndWaitForRunResponse(false, mode); + } + + @Test + void databaseNameInBeginTransaction() { + testDatabaseNameSupport(false); + } + + @Test + void databaseNameForAutoCommitTransactions() { + testDatabaseNameSupport(true); + } + + @Test + void shouldSupportDatabaseNameInBeginTransaction() { + var txStage = protocol.beginTransaction( + connectionMock("foo", protocol), + Collections.emptySet(), + TransactionConfig.empty(), + null, + null, + Logging.none(), + true); + + assertDoesNotThrow(() -> await(txStage)); + } + + @Test + void shouldNotSupportDatabaseNameForAutoCommitTransactions() { + assertDoesNotThrow(() -> protocol.runInAutoCommitTransaction( + connectionMock("foo", protocol), + new Query("RETURN 1"), + Collections.emptySet(), + (ignored) -> {}, + TransactionConfig.empty(), + UNLIMITED_FETCH_SIZE, + null, + Logging.none())); + } + + @Test + void shouldTelemetrySendTelemetryMessage() { + var connection = connectionMock(); + doAnswer((invocationOnMock) -> { + var handler = (TelemetryResponseHandler) invocationOnMock.getArgument(1); + handler.onSuccess(Map.of()); + return null; + }) + .when(connection) + .write(Mockito.any(), Mockito.any()); + var expectedApi = 1; + + await(protocol.telemetry(connection, expectedApi)); + + verify(connection).write(Mockito.eq(new TelemetryMessage(expectedApi)), Mockito.any()); + verify(connection, never()).writeAndFlush(Mockito.any(), Mockito.any()); + } + + private Class expectedMessageFormatType() { + return MessageFormatV54.class; + } + + private void testFailedRunInAutoCommitTxWithWaitingForResponse( + Set bookmarks, TransactionConfig config, AccessMode mode) { + // Given + var connection = connectionMock(mode, protocol); + @SuppressWarnings("unchecked") + Consumer bookmarkConsumer = mock(Consumer.class); + + var cursorFuture = protocol.runInAutoCommitTransaction( + connection, + QUERY, + bookmarks, + bookmarkConsumer, + config, + UNLIMITED_FETCH_SIZE, + null, + Logging.none()) + .asyncResult() + .toCompletableFuture(); + + var runHandler = verifySessionRunInvoked(connection, bookmarks, config, mode, defaultDatabase()); + assertFalse(cursorFuture.isDone()); + + // When I response to Run message with a failure + Throwable error = new RuntimeException(); + runHandler.onFailure(error); + + // Then + then(bookmarkConsumer).should(times(0)).accept(any()); + assertTrue(cursorFuture.isDone()); + var actual = + assertThrows(error.getClass(), () -> await(cursorFuture.get().mapSuccessfulRunCompletionAsync())); + assertSame(error, actual); + } + + private void testSuccessfulRunInAutoCommitTxWithWaitingForResponse( + Set bookmarks, TransactionConfig config, AccessMode mode) throws Exception { + // Given + var connection = connectionMock(mode, protocol); + @SuppressWarnings("unchecked") + Consumer bookmarkConsumer = mock(Consumer.class); + + var cursorFuture = protocol.runInAutoCommitTransaction( + connection, + QUERY, + bookmarks, + bookmarkConsumer, + config, + UNLIMITED_FETCH_SIZE, + null, + Logging.none()) + .asyncResult() + .toCompletableFuture(); + + var runHandler = verifySessionRunInvoked(connection, bookmarks, config, mode, defaultDatabase()); + assertFalse(cursorFuture.isDone()); + + // When I response to the run message + runHandler.onSuccess(emptyMap()); + + // Then + then(bookmarkConsumer).should(times(0)).accept(any()); + assertTrue(cursorFuture.isDone()); + assertNotNull(cursorFuture.get()); + } + + private void testRunInUnmanagedTransactionAndWaitForRunResponse(boolean success, AccessMode mode) throws Exception { + // Given + var connection = connectionMock(mode, protocol); + + var cursorFuture = protocol.runInUnmanagedTransaction( + connection, QUERY, mock(UnmanagedTransaction.class), UNLIMITED_FETCH_SIZE) + .asyncResult() + .toCompletableFuture(); + + var runHandler = verifyTxRunInvoked(connection); + assertFalse(cursorFuture.isDone()); + Throwable error = new RuntimeException(); + + if (success) { + runHandler.onSuccess(emptyMap()); + } else { + // When responded with a failure + runHandler.onFailure(error); + } + + // Then + assertTrue(cursorFuture.isDone()); + if (success) { + assertNotNull(await(cursorFuture.get().mapSuccessfulRunCompletionAsync())); + } else { + var actual = assertThrows( + error.getClass(), () -> await(cursorFuture.get().mapSuccessfulRunCompletionAsync())); + assertSame(error, actual); + } + } + + private void testRunAndWaitForRunResponse(boolean autoCommitTx, TransactionConfig config, AccessMode mode) + throws Exception { + // Given + var connection = connectionMock(mode, protocol); + var initialBookmarks = Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx987")); + + CompletionStage cursorStage; + if (autoCommitTx) { + cursorStage = protocol.runInAutoCommitTransaction( + connection, + QUERY, + initialBookmarks, + (ignored) -> {}, + config, + UNLIMITED_FETCH_SIZE, + null, + Logging.none()) + .asyncResult(); + } else { + cursorStage = protocol.runInUnmanagedTransaction( + connection, QUERY, mock(UnmanagedTransaction.class), UNLIMITED_FETCH_SIZE) + .asyncResult(); + } + + // When & Then + var cursorFuture = cursorStage.toCompletableFuture(); + assertFalse(cursorFuture.isDone()); + + var runResponseHandler = autoCommitTx + ? verifySessionRunInvoked(connection, initialBookmarks, config, mode, defaultDatabase()) + : verifyTxRunInvoked(connection); + runResponseHandler.onSuccess(emptyMap()); + + assertTrue(cursorFuture.isDone()); + assertNotNull(cursorFuture.get()); + } + + private void testDatabaseNameSupport(boolean autoCommitTx) { + var connection = connectionMock("foo", protocol); + if (autoCommitTx) { + var factory = protocol.runInAutoCommitTransaction( + connection, + QUERY, + Collections.emptySet(), + (ignored) -> {}, + TransactionConfig.empty(), + UNLIMITED_FETCH_SIZE, + null, + Logging.none()); + var resultStage = factory.asyncResult(); + var runHandler = verifySessionRunInvoked( + connection, Collections.emptySet(), TransactionConfig.empty(), AccessMode.WRITE, database("foo")); + runHandler.onSuccess(emptyMap()); + await(resultStage); + verifySessionRunInvoked( + connection, Collections.emptySet(), TransactionConfig.empty(), AccessMode.WRITE, database("foo")); + } else { + var txStage = protocol.beginTransaction( + connection, Collections.emptySet(), TransactionConfig.empty(), null, null, Logging.none(), true); + await(txStage); + verifyBeginInvoked(connection, Collections.emptySet(), TransactionConfig.empty(), database("foo")); + } + } + + private ResponseHandler verifyTxRunInvoked(Connection connection) { + return verifyRunInvoked(connection, RunWithMetadataMessage.unmanagedTxRunMessage(QUERY)); + } + + private ResponseHandler verifySessionRunInvoked( + Connection connection, + Set bookmark, + TransactionConfig config, + AccessMode mode, + DatabaseName databaseName) { + var runMessage = RunWithMetadataMessage.autoCommitTxRunMessage( + QUERY, config, databaseName, mode, bookmark, null, null, Logging.none()); + return verifyRunInvoked(connection, runMessage); + } + + private ResponseHandler verifyRunInvoked(Connection connection, RunWithMetadataMessage runMessage) { + var runHandlerCaptor = ArgumentCaptor.forClass(ResponseHandler.class); + var pullHandlerCaptor = ArgumentCaptor.forClass(ResponseHandler.class); + + verify(connection).write(eq(runMessage), runHandlerCaptor.capture()); + verify(connection).writeAndFlush(any(PullMessage.class), pullHandlerCaptor.capture()); + + assertThat(runHandlerCaptor.getValue(), instanceOf(RunResponseHandler.class)); + assertThat(pullHandlerCaptor.getValue(), instanceOf(PullAllResponseHandler.class)); + + return runHandlerCaptor.getValue(); + } + + private void verifyBeginInvoked( + Connection connection, Set bookmarks, TransactionConfig config, DatabaseName databaseName) { + var beginHandlerCaptor = ArgumentCaptor.forClass(ResponseHandler.class); + var beginMessage = + new BeginMessage(bookmarks, config, databaseName, AccessMode.WRITE, null, null, null, Logging.none()); + verify(connection).writeAndFlush(eq(beginMessage), beginHandlerCaptor.capture()); + assertThat(beginHandlerCaptor.getValue(), instanceOf(BeginTxResponseHandler.class)); + } + + private static InternalAuthToken dummyAuthToken() { + return (InternalAuthToken) AuthTokens.basic("hello", "world"); + } +} diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/MessageFormatV54Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/MessageFormatV54Test.java new file mode 100644 index 0000000000..06cb35c74d --- /dev/null +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/MessageFormatV54Test.java @@ -0,0 +1,47 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.v54; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.instanceOf; +import static org.mockito.Mockito.mock; + +import org.junit.jupiter.api.Test; +import org.neo4j.driver.internal.messaging.MessageFormat; +import org.neo4j.driver.internal.messaging.v5.MessageReaderV5; +import org.neo4j.driver.internal.packstream.PackInput; +import org.neo4j.driver.internal.packstream.PackOutput; + +class MessageFormatV54Test { + private static final MessageFormat format = BoltProtocolV54.INSTANCE.createMessageFormat(); + + @Test + void shouldCreateCorrectWriter() { + var writer = format.newWriter(mock(PackOutput.class)); + + assertThat(writer, instanceOf(MessageWriterV54.class)); + } + + @Test + void shouldCreateCorrectReader() { + var reader = format.newReader(mock(PackInput.class)); + + assertThat(reader, instanceOf(MessageReaderV5.class)); + } +} diff --git a/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/MessageWriterV54Test.java b/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/MessageWriterV54Test.java new file mode 100644 index 0000000000..bef9aca7a4 --- /dev/null +++ b/driver/src/test/java/org/neo4j/driver/internal/messaging/v54/MessageWriterV54Test.java @@ -0,0 +1,207 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.messaging.v54; + +import static java.time.Duration.ofSeconds; +import static java.util.Calendar.DECEMBER; +import static java.util.Collections.emptyMap; +import static java.util.Collections.singletonMap; +import static org.neo4j.driver.AccessMode.READ; +import static org.neo4j.driver.AccessMode.WRITE; +import static org.neo4j.driver.AuthTokens.basic; +import static org.neo4j.driver.Values.point; +import static org.neo4j.driver.Values.value; +import static org.neo4j.driver.internal.DatabaseNameUtil.database; +import static org.neo4j.driver.internal.DatabaseNameUtil.defaultDatabase; +import static org.neo4j.driver.internal.messaging.request.CommitMessage.COMMIT; +import static org.neo4j.driver.internal.messaging.request.DiscardAllMessage.DISCARD_ALL; +import static org.neo4j.driver.internal.messaging.request.GoodbyeMessage.GOODBYE; +import static org.neo4j.driver.internal.messaging.request.PullAllMessage.PULL_ALL; +import static org.neo4j.driver.internal.messaging.request.ResetMessage.RESET; +import static org.neo4j.driver.internal.messaging.request.RollbackMessage.ROLLBACK; +import static org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage.autoCommitTxRunMessage; +import static org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage.unmanagedTxRunMessage; + +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetTime; +import java.time.ZoneId; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Stream; +import org.neo4j.driver.Logging; +import org.neo4j.driver.Query; +import org.neo4j.driver.Value; +import org.neo4j.driver.Values; +import org.neo4j.driver.internal.BoltAgentUtil; +import org.neo4j.driver.internal.InternalBookmark; +import org.neo4j.driver.internal.messaging.Message; +import org.neo4j.driver.internal.messaging.MessageFormat; +import org.neo4j.driver.internal.messaging.request.BeginMessage; +import org.neo4j.driver.internal.messaging.request.DiscardMessage; +import org.neo4j.driver.internal.messaging.request.HelloMessage; +import org.neo4j.driver.internal.messaging.request.PullMessage; +import org.neo4j.driver.internal.messaging.request.RouteMessage; +import org.neo4j.driver.internal.messaging.request.TelemetryMessage; +import org.neo4j.driver.internal.packstream.PackOutput; +import org.neo4j.driver.internal.security.InternalAuthToken; +import org.neo4j.driver.internal.util.messaging.AbstractMessageWriterTestBase; + +public class MessageWriterV54Test extends AbstractMessageWriterTestBase { + @Override + protected MessageFormat.Writer newWriter(PackOutput output) { + return BoltProtocolV54.INSTANCE.createMessageFormat().newWriter(output); + } + + @Override + protected Stream supportedMessages() { + return Stream.of( + // Bolt V2 Data Types + unmanagedTxRunMessage(new Query("RETURN $point", singletonMap("point", point(42, 12.99, -180.0)))), + unmanagedTxRunMessage( + new Query("RETURN $point", singletonMap("point", point(42, 0.51, 2.99, 100.123)))), + unmanagedTxRunMessage( + new Query("RETURN $date", singletonMap("date", value(LocalDate.ofEpochDay(2147483650L))))), + unmanagedTxRunMessage(new Query( + "RETURN $time", singletonMap("time", value(OffsetTime.of(4, 16, 20, 999, ZoneOffset.MIN))))), + unmanagedTxRunMessage( + new Query("RETURN $time", singletonMap("time", value(LocalTime.of(12, 9, 18, 999_888))))), + unmanagedTxRunMessage(new Query( + "RETURN $dateTime", + singletonMap("dateTime", value(LocalDateTime.of(2049, DECEMBER, 12, 17, 25, 49, 199))))), + unmanagedTxRunMessage(new Query( + "RETURN $dateTime", + singletonMap( + "dateTime", + value(ZonedDateTime.of( + 2000, 1, 10, 12, 2, 49, 300, ZoneOffset.ofHoursMinutes(9, 30)))))), + unmanagedTxRunMessage(new Query( + "RETURN $dateTime", + singletonMap( + "dateTime", + value(ZonedDateTime.of(2000, 1, 10, 12, 2, 49, 300, ZoneId.of("Europe/Stockholm")))))), + + // New Bolt V4 messages + new PullMessage(100, 200), + new DiscardMessage(300, 400), + + // Bolt V3 messages + new HelloMessage( + "MyDriver/1.2.3", + BoltAgentUtil.VALUE, + ((InternalAuthToken) basic("neo4j", "neo4j")).toMap(), + Collections.emptyMap(), + false, + null), + GOODBYE, + new BeginMessage( + Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx123")), + ofSeconds(5), + singletonMap("key", value(42)), + READ, + defaultDatabase(), + null, + null, + null, + Logging.none()), + new BeginMessage( + Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx123")), + ofSeconds(5), + singletonMap("key", value(42)), + WRITE, + database("foo"), + null, + null, + null, + Logging.none()), + COMMIT, + ROLLBACK, + RESET, + autoCommitTxRunMessage( + new Query("RETURN 1"), + ofSeconds(5), + singletonMap("key", value(42)), + defaultDatabase(), + READ, + Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx1")), + null, + null, + Logging.none()), + autoCommitTxRunMessage( + new Query("RETURN 1"), + ofSeconds(5), + singletonMap("key", value(42)), + database("foo"), + WRITE, + Collections.singleton(InternalBookmark.parse("neo4j:bookmark:v1:tx1")), + null, + null, + Logging.none()), + unmanagedTxRunMessage(new Query("RETURN 1")), + + // Bolt V3 messages with struct values + autoCommitTxRunMessage( + new Query("RETURN $x", singletonMap("x", value(ZonedDateTime.now()))), + ofSeconds(1), + emptyMap(), + defaultDatabase(), + READ, + Collections.emptySet(), + null, + null, + Logging.none()), + autoCommitTxRunMessage( + new Query("RETURN $x", singletonMap("x", value(ZonedDateTime.now()))), + ofSeconds(1), + emptyMap(), + database("foo"), + WRITE, + Collections.emptySet(), + null, + null, + Logging.none()), + unmanagedTxRunMessage(new Query("RETURN $x", singletonMap("x", point(42, 1, 2, 3)))), + + // New 4.3 Messages + routeMessage(), + // New 5.4 message + telemetryMessage()); + } + + @Override + protected Stream unsupportedMessages() { + return Stream.of( + // Bolt V1, V2 and V3 messages + PULL_ALL, DISCARD_ALL); + } + + private RouteMessage routeMessage() { + Map routeContext = new HashMap<>(); + routeContext.put("someContext", Values.value(124)); + return new RouteMessage(routeContext, Collections.emptySet(), "dbName", null); + } + + private TelemetryMessage telemetryMessage() { + return new TelemetryMessage(1); + } +} diff --git a/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalReactiveSessionTest.java b/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalReactiveSessionTest.java index 03107ee790..989c6ae1cc 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalReactiveSessionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalReactiveSessionTest.java @@ -25,6 +25,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.BDDMockito.given; import static org.mockito.BDDMockito.then; @@ -62,6 +63,8 @@ import org.neo4j.driver.internal.cursor.RxResultCursor; import org.neo4j.driver.internal.cursor.RxResultCursorImpl; import org.neo4j.driver.internal.retry.RetryLogic; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.FixedRetryLogic; import org.neo4j.driver.internal.util.Futures; import org.neo4j.driver.internal.value.IntegerValue; @@ -142,8 +145,9 @@ void shouldDelegateBeginTx(Function Futures.getNow(txFuture)); MatcherAssert.assertThat(t.getCause(), equalTo(error)); verify(session).releaseConnectionAsync(); @@ -189,7 +194,8 @@ void shouldRetryOnError() { var tx = mock(UnmanagedTransaction.class); when(tx.closeAsync(false)).thenReturn(completedWithNull()); - when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class))) + when(session.beginTransactionAsync( + any(AccessMode.class), any(TransactionConfig.class), any(ApiTelemetryWork.class))) .thenReturn(completedFuture(tx)); when(session.retryLogic()).thenReturn(new FixedRetryLogic(retryCount)); var rxSession = new InternalRxSession(session); @@ -204,7 +210,8 @@ void shouldRetryOnError() { // Then verify(session, times(retryCount + 1)) - .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class)); + .beginTransactionAsync( + any(AccessMode.class), any(TransactionConfig.class), any(ApiTelemetryWork.class)); verify(tx, times(retryCount + 1)).closeAsync(false); } @@ -218,7 +225,8 @@ void shouldObtainResultIfRetrySucceed() { when(tx.closeAsync(false)).thenReturn(completedWithNull()); when(tx.closeAsync(true)).thenReturn(completedWithNull()); - when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class))) + when(session.beginTransactionAsync( + any(AccessMode.class), any(TransactionConfig.class), any(ApiTelemetryWork.class))) .thenReturn(completedFuture(tx)); when(session.retryLogic()).thenReturn(new FixedRetryLogic(retryCount)); var rxSession = new InternalRxSession(session); @@ -237,7 +245,8 @@ void shouldObtainResultIfRetrySucceed() { // Then verify(session, times(retryCount + 1)) - .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class)); + .beginTransactionAsync( + any(AccessMode.class), any(TransactionConfig.class), any(ApiTelemetryWork.class)); verify(tx, times(retryCount)).closeAsync(false); verify(tx).closeAsync(true); } diff --git a/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxSessionTest.java b/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxSessionTest.java index f7d5f5a1f7..04d15f6ce3 100644 --- a/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxSessionTest.java +++ b/driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxSessionTest.java @@ -25,6 +25,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -51,6 +52,8 @@ import org.neo4j.driver.internal.async.UnmanagedTransaction; import org.neo4j.driver.internal.cursor.RxResultCursor; import org.neo4j.driver.internal.cursor.RxResultCursorImpl; +import org.neo4j.driver.internal.telemetry.ApiTelemetryWork; +import org.neo4j.driver.internal.telemetry.TelemetryApi; import org.neo4j.driver.internal.util.FixedRetryLogic; import org.neo4j.driver.internal.util.Futures; import org.neo4j.driver.internal.value.IntegerValue; @@ -144,8 +147,9 @@ void shouldDelegateBeginTx(Function> beginTx // Given var session = mock(NetworkSession.class); var tx = mock(UnmanagedTransaction.class); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); - when(session.beginTransactionAsync(any(TransactionConfig.class), isNull())) + when(session.beginTransactionAsync(any(TransactionConfig.class), isNull(), eq(apiTelemetryWork))) .thenReturn(completedFuture(tx)); var rxSession = new InternalRxSession(session); @@ -154,7 +158,7 @@ void shouldDelegateBeginTx(Function> beginTx StepVerifier.create(Mono.from(rxTx)).expectNextCount(1).verifyComplete(); // Then - verify(session).beginTransactionAsync(any(TransactionConfig.class), isNull()); + verify(session).beginTransactionAsync(any(TransactionConfig.class), isNull(), eq(apiTelemetryWork)); } @ParameterizedTest @@ -163,9 +167,10 @@ void shouldReleaseConnectionIfFailedToBeginTx(Function Futures.getNow(txFuture)); assertThat(t.getCause(), equalTo(error)); verify(session).releaseConnectionAsync(); @@ -188,9 +193,10 @@ void shouldDelegateRunTx(Function> runTx) { // Given var session = mock(NetworkSession.class); var tx = mock(UnmanagedTransaction.class); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.MANAGED_TRANSACTION); when(tx.closeAsync(true)).thenReturn(completedWithNull()); - when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class))) + when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class), eq(apiTelemetryWork))) .thenReturn(completedFuture(tx)); when(session.retryLogic()).thenReturn(new FixedRetryLogic(1)); var rxSession = new InternalRxSession(session); @@ -200,7 +206,8 @@ void shouldDelegateRunTx(Function> runTx) { StepVerifier.create(Flux.from(strings)).expectNext("a").verifyComplete(); // Then - verify(session).beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class)); + verify(session) + .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class), eq(apiTelemetryWork)); verify(tx).closeAsync(true); } @@ -211,8 +218,9 @@ void shouldRetryOnError() { var session = mock(NetworkSession.class); var tx = mock(UnmanagedTransaction.class); when(tx.closeAsync(false)).thenReturn(completedWithNull()); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.MANAGED_TRANSACTION); - when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class))) + when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class), eq(apiTelemetryWork))) .thenReturn(completedFuture(tx)); when(session.retryLogic()).thenReturn(new FixedRetryLogic(retryCount)); var rxSession = new InternalRxSession(session); @@ -227,7 +235,7 @@ void shouldRetryOnError() { // Then verify(session, times(retryCount + 1)) - .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class)); + .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class), eq(apiTelemetryWork)); verify(tx, times(retryCount + 1)).closeAsync(false); } @@ -237,10 +245,11 @@ void shouldObtainResultIfRetrySucceed() { var retryCount = 2; var session = mock(NetworkSession.class); var tx = mock(UnmanagedTransaction.class); + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.MANAGED_TRANSACTION); when(tx.closeAsync(false)).thenReturn(completedWithNull()); when(tx.closeAsync(true)).thenReturn(completedWithNull()); - when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class))) + when(session.beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class), eq(apiTelemetryWork))) .thenReturn(completedFuture(tx)); when(session.retryLogic()).thenReturn(new FixedRetryLogic(retryCount)); var rxSession = new InternalRxSession(session); @@ -259,7 +268,7 @@ void shouldObtainResultIfRetrySucceed() { // Then verify(session, times(retryCount + 1)) - .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class)); + .beginTransactionAsync(any(AccessMode.class), any(TransactionConfig.class), eq(apiTelemetryWork)); verify(tx, times(retryCount)).closeAsync(false); verify(tx).closeAsync(true); } diff --git a/driver/src/test/java/org/neo4j/driver/internal/telemetry/ApiTelemetryWorkTest.java b/driver/src/test/java/org/neo4j/driver/internal/telemetry/ApiTelemetryWorkTest.java new file mode 100644 index 0000000000..358ba32427 --- /dev/null +++ b/driver/src/test/java/org/neo4j/driver/internal/telemetry/ApiTelemetryWorkTest.java @@ -0,0 +1,151 @@ +/* + * Copyright (c) "Neo4j" + * Neo4j Sweden AB [http://neo4j.com] + * + * This file is part of Neo4j. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.neo4j.driver.internal.telemetry; + +import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; +import java.util.stream.Stream; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mockito; +import org.neo4j.driver.internal.messaging.BoltProtocol; +import org.neo4j.driver.internal.spi.Connection; +import org.neo4j.driver.testutil.TestUtil; + +class ApiTelemetryWorkTest { + + @ParameterizedTest + @MethodSource("shouldNotSendTelemetrySource") + public void shouldNotCallTelemetryAndCompleteStage( + boolean telemetryEnabled, Consumer transformWorker) { + var apiTelemetryWork = new ApiTelemetryWork(TelemetryApi.UNMANAGED_TRANSACTION); + var protocol = Mockito.mock(BoltProtocol.class); + var connection = Mockito.mock(Connection.class); + Mockito.doReturn(telemetryEnabled).when(connection).isTelemetryEnabled(); + transformWorker.accept(apiTelemetryWork); + + TestUtil.await(apiTelemetryWork.execute(connection, protocol)); + + Mockito.verify(protocol, Mockito.never()).telemetry(Mockito.any(), Mockito.any()); + } + + @ParameterizedTest + @MethodSource("shouldCallTelemetry") + public void shouldCallTelemetryWithCorrectValuesAndResolveFuture( + TelemetryApi telemetryApi, boolean telemetryEnabled, Consumer transformWorker) { + var apiTelemetryWork = new ApiTelemetryWork(telemetryApi); + var protocol = Mockito.mock(BoltProtocol.class); + var connection = Mockito.mock(Connection.class); + Mockito.doReturn(telemetryEnabled).when(connection).isTelemetryEnabled(); + Mockito.doReturn(CompletableFuture.completedFuture(null)) + .when(protocol) + .telemetry(Mockito.any(), Mockito.any()); + transformWorker.accept(apiTelemetryWork); + + TestUtil.await(apiTelemetryWork.execute(connection, protocol)); + + Mockito.verify(protocol, Mockito.only()).telemetry(connection, telemetryApi.getValue()); + } + + @ParameterizedTest + @MethodSource("shouldCallTelemetry") + public void shouldCallTelemetryWithCorrectValuesAndFailedFuture( + TelemetryApi telemetryApi, boolean telemetryEnabled, Consumer transformWorker) { + var apiTelemetryWork = new ApiTelemetryWork(telemetryApi); + var protocol = Mockito.mock(BoltProtocol.class); + var connection = Mockito.mock(Connection.class); + var exception = new RuntimeException("something wrong"); + Mockito.doReturn(telemetryEnabled).when(connection).isTelemetryEnabled(); + Mockito.doReturn(CompletableFuture.failedFuture(exception)) + .when(protocol) + .telemetry(Mockito.any(), Mockito.any()); + transformWorker.accept(apiTelemetryWork); + + Assertions.assertThrows( + RuntimeException.class, () -> TestUtil.await(apiTelemetryWork.execute(connection, protocol))); + + Mockito.verify(protocol, Mockito.only()).telemetry(connection, telemetryApi.getValue()); + } + + public static Stream shouldNotSendTelemetrySource() { + return Stream.of( + Arguments.of(false, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkSetEnabledWithFalse), + Arguments.of(false, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkSetEnabledWithTrue), + Arguments.of(false, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkExecuteWithSuccess), + Arguments.of( + false, (Consumer) ApiTelemetryWorkTest::callApiTelemetryWorkExecuteWithError), + Arguments.of(false, (Consumer) ApiTelemetryWorkTest::noop), + Arguments.of(true, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkSetEnabledWithFalse), + Arguments.of(true, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkExecuteWithSuccess)); + } + + private static Stream shouldCallTelemetry() { + return Stream.of(TelemetryApi.values()) + .flatMap(telemetryApi -> Stream.of( + Arguments.of(telemetryApi, true, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkSetEnabledWithTrue), + Arguments.of(telemetryApi, true, (Consumer) + ApiTelemetryWorkTest::callApiTelemetryWorkExecuteWithError), + Arguments.of(telemetryApi, true, (Consumer) ApiTelemetryWorkTest::noop))); + } + + private static void callApiTelemetryWorkSetEnabledWithTrue(ApiTelemetryWork apiTelemetryWork) { + apiTelemetryWork.setEnabled(true); + } + + private static void callApiTelemetryWorkSetEnabledWithFalse(ApiTelemetryWork apiTelemetryWork) { + apiTelemetryWork.setEnabled(false); + } + + @SuppressWarnings("EmptyMethod") + private static void noop(ApiTelemetryWork apiTelemetryWork) {} + + private static void callApiTelemetryWorkExecuteWithSuccess(ApiTelemetryWork apiTelemetryWork) { + var protocol = Mockito.mock(BoltProtocol.class); + var connection = Mockito.mock(Connection.class); + Mockito.doReturn(CompletableFuture.completedFuture(null)) + .when(protocol) + .telemetry(Mockito.any(), Mockito.any()); + Mockito.doReturn(true).when(connection).isTelemetryEnabled(); + + TestUtil.await(apiTelemetryWork.execute(connection, protocol)); + } + + private static void callApiTelemetryWorkExecuteWithError(ApiTelemetryWork apiTelemetryWork) { + var protocol = Mockito.mock(BoltProtocol.class); + var connection = Mockito.mock(Connection.class); + Mockito.doReturn(CompletableFuture.failedFuture(new RuntimeException("WRONG"))) + .when(protocol) + .telemetry(Mockito.any(), Mockito.any()); + Mockito.doReturn(true).when(connection).isTelemetryEnabled(); + + try { + TestUtil.await(apiTelemetryWork.execute(connection, protocol)); + } catch (Exception ex) { + // ignore since the error is expected. + } + } +} diff --git a/driver/src/test/java/org/neo4j/driver/testutil/TestUtil.java b/driver/src/test/java/org/neo4j/driver/testutil/TestUtil.java index 2fb1691bd9..7900a2cc76 100644 --- a/driver/src/test/java/org/neo4j/driver/testutil/TestUtil.java +++ b/driver/src/test/java/org/neo4j/driver/testutil/TestUtil.java @@ -98,6 +98,7 @@ import org.neo4j.driver.internal.messaging.v51.BoltProtocolV51; import org.neo4j.driver.internal.messaging.v52.BoltProtocolV52; import org.neo4j.driver.internal.messaging.v53.BoltProtocolV53; +import org.neo4j.driver.internal.messaging.v54.BoltProtocolV54; import org.neo4j.driver.internal.retry.RetryLogic; import org.neo4j.driver.internal.spi.Connection; import org.neo4j.driver.internal.spi.ConnectionProvider; @@ -254,6 +255,15 @@ public static NetworkSession newSession(ConnectionProvider connectionProvider) { public static NetworkSession newSession( ConnectionProvider connectionProvider, AccessMode mode, RetryLogic retryLogic, Set bookmarks) { + return newSession(connectionProvider, mode, retryLogic, bookmarks, true); + } + + public static NetworkSession newSession( + ConnectionProvider connectionProvider, + AccessMode mode, + RetryLogic retryLogic, + Set bookmarks, + boolean telemetryDisabled) { return new NetworkSession( connectionProvider, retryLogic, @@ -265,7 +275,8 @@ public static NetworkSession newSession( DEV_NULL_LOGGING, NoOpBookmarkManager.INSTANCE, null, - null); + null, + telemetryDisabled); } public static void verifyRunRx(Connection connection, String query) { @@ -455,7 +466,8 @@ public static Connection connectionMock(String databaseName, AccessMode mode, Bo BoltProtocolV5.VERSION, BoltProtocolV51.VERSION, BoltProtocolV52.VERSION, - BoltProtocolV53.VERSION) + BoltProtocolV53.VERSION, + BoltProtocolV54.VERSION) .contains(version)) { setupSuccessResponse(connection, CommitMessage.class); setupSuccessResponse(connection, RollbackMessage.class); diff --git a/pom.xml b/pom.xml index ac52c851f8..eac979aaee 100644 --- a/pom.xml +++ b/pom.xml @@ -158,6 +158,12 @@ testng ${testng.version} test + + + org.webjars + jquery + + org.rauschig diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java index cc14742025..220e7bdaf7 100644 --- a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java +++ b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java @@ -42,6 +42,7 @@ public class GetFeatures implements TestkitRequest { "Feature:Bolt:5.1", "Feature:Bolt:5.2", "Feature:Bolt:5.3", + "Feature:Bolt:5.4", "AuthorizationExpiredTreatment", "ConfHint:connection.recv_timeout_seconds", "Feature:Auth:Bearer", diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/NewDriver.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/NewDriver.java index 6dd412664f..5e549eac9a 100644 --- a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/NewDriver.java +++ b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/NewDriver.java @@ -99,6 +99,7 @@ public TestkitResponse process(TestkitState testkitState) { Optional.ofNullable(data.maxConnectionPoolSize).ifPresent(configBuilder::withMaxConnectionPoolSize); Optional.ofNullable(data.connectionAcquisitionTimeoutMs) .ifPresent(timeout -> configBuilder.withConnectionAcquisitionTimeout(timeout, TimeUnit.MILLISECONDS)); + Optional.ofNullable(data.telemetryDisabled).ifPresent(configBuilder::withTelemetryDisabled); configBuilder.withNotificationConfig( toNotificationConfig(data.notificationsMinSeverity, data.notificationsDisabledCategories)); configBuilder.withDriverMetrics(); @@ -293,6 +294,7 @@ public static class NewDriverBody { private Long connectionAcquisitionTimeoutMs; private boolean encrypted; private List trustedCertificates; + private Boolean telemetryDisabled; } @RequiredArgsConstructor