Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
private boolean leasePinnedAtRequestStart;
private boolean sessionAffinity;
private boolean pooledConnectionPoisoned;
private boolean nativeJobConnection;

@Override
public boolean supportsQueryTiming() { return true; }
Expand All @@ -53,6 +54,7 @@ public final void connect(ConnectParams params) {
requestActive = false;
leasePinnedAtRequestStart = false;
pooledConnectionPoisoned = false;
nativeJobConnection = false;
loadDriver(params);
configuredDatabase = params.getDatabase();
connectParams = params;
Expand Down Expand Up @@ -289,6 +291,7 @@ public synchronized void disconnect() {
requestActive = false;
leasePinnedAtRequestStart = false;
pooledConnectionPoisoned = false;
nativeJobConnection = false;
identifierQuote = "";
});
}
Expand Down Expand Up @@ -347,6 +350,7 @@ final void beginPooledRequest() throws Exception {
String identity;
JdbcConnectionPoolRegistry registry;
synchronized (this) {
if (nativeJobConnection) return;
if (poolRegistry == null) {
return;
}
Expand Down Expand Up @@ -395,6 +399,7 @@ final synchronized void finishPooledRequest(
boolean endsSessionAffinity,
boolean preservesSchemaContext
) {
if (nativeJobConnection) return;
if (poolRegistry == null) {
return;
}
Expand Down Expand Up @@ -445,6 +450,53 @@ final synchronized void releaseIdlePooledConnection(JdbcExecutor executor) {
releasePooledConnection(!preparePooledConnectionForReturn());
}

final synchronized void requireIdleNativeJobSession(JdbcExecutor executor) throws Exception {
if (connectParams == null) throw new IllegalStateException("Not connected");
if (JdbcSessionRole.from(connectParams.getSessionRole()) != JdbcSessionRole.METADATA) {
throw new IllegalStateException("native_job_metadata_session_required");
}
if (nativeJobConnection || requestActive || pooledLease != null || sessionAffinity
|| executor.hasOpenSessions() || executor.hasActiveStatements() || hasRetainedResultResources()
|| (connection != null && !connection.getAutoCommit())) {
throw new IllegalStateException("native_job_session_busy");
}
}

final synchronized Connection openNativeJobConnection() throws Exception {
nativeJobConnection = true;
try {
ConnectParams params = connectParams;
if (connection == null) {
connection = poolRegistry == null ? openInitializedConnection(params)
: poolRegistry.openDedicatedConnection(() -> openInitializedConnection(params));
afterConnect(params, connection);
}
connection.setAutoCommit(false);
return connection;
} catch (Exception | LinkageError error) {
try {
closeNativeJobConnection();
} catch (Exception | LinkageError closeError) {
error.addSuppressed(closeError);
throw AgentRpcError.resource("close",
new JdbcConnectionPoolRegistry.PhysicalConnectionStateUnknownException(error));
}
if (!JdbcConnectionPoolRegistry.requiresRuntimeReplacement(error)) nativeJobConnection = false;
throw error;
}
}

final synchronized void closeNativeJobConnection() throws Exception {
if (!nativeJobConnection) return;
Connection owned = connection;
connection = null;
if (owned != null) owned.close();
}

final synchronized boolean hasNativeJobConnection() {
return nativeJobConnection;
}

private static String resolveIdentifierQuote(ConnectParams params, Connection connection) {
if (isPostgresCompatibleJdbc(params)) {
try (Statement statement = connection.createStatement()) {
Expand Down
36 changes: 34 additions & 2 deletions agents/common/src/main/java/com/dbx/agent/AgentProtocol.java
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,12 @@ public final class AgentProtocol {
public static final String METHOD_BEGIN_MANUAL_TRANSACTION = "begin_manual_transaction";
public static final String METHOD_COMMIT_MANUAL_TRANSACTION = "commit_manual_transaction";
public static final String METHOD_ROLLBACK_MANUAL_TRANSACTION = "rollback_manual_transaction";
public static final String METHOD_PREPARE_LEGACY_JOB = "prepare_legacy_job";
public static final String METHOD_COMMIT_LEGACY_JOB = "commit_legacy_job";
public static final String METHOD_ROLLBACK_LEGACY_JOB = "rollback_legacy_job";
public static final String METHOD_LEGACY_JOB_STATUS = "legacy_job_status";
public static final String METHOD_OPEN_NATIVE_JOB_SESSION = "open_native_job_session";
public static final String METHOD_CLOSE_NATIVE_JOB_SESSION = "close_native_job_session";
public static final String METHOD_DISCONNECT = "disconnect";
public static final String METHOD_SHUTDOWN = "shutdown";

Expand Down Expand Up @@ -144,6 +150,9 @@ public final class AgentProtocol {
public static final String CAPABILITY_MULTI_SESSION = "multi_session";
public static final String CAPABILITY_STRUCTURED_ERROR_V1 = "structured_error_v1";
public static final String CAPABILITY_BLOB_BIND_STATEMENTS_V1 = "blob_bind_statements_v1";
public static final String CAPABILITY_NATIVE_JOB_LEGACY_V1 = "native_job_legacy_v1";
public static final String CAPABILITY_NATIVE_JOB_LEGACY_OB_V1 = "native_job_legacy_ob_v1";
public static final String CAPABILITY_NATIVE_JOB_SESSION_V1 = "native_job_session_v1";
public static final String CAPABILITY_EXACT_METADATA_IDENTITY_V1 = "exact_metadata_identity_v1";

public static final List<String> CAPABILITIES = Collections.unmodifiableList(Arrays.asList(
Expand Down Expand Up @@ -183,6 +192,8 @@ public final class AgentProtocol {
CAPABILITY_MONGO_BULK_WRITE,
CAPABILITY_MONGO_RENAME_COLLECTION,
CAPABILITY_MONGO_FIND_CURSOR,
CAPABILITY_NATIVE_JOB_LEGACY_V1,
CAPABILITY_NATIVE_JOB_SESSION_V1,
CAPABILITY_EXACT_METADATA_IDENTITY_V1
));

Expand Down Expand Up @@ -236,6 +247,16 @@ public final class AgentProtocol {
));

public static final List<String> MULTI_SESSION_METHODS;
public static final List<String> NATIVE_JOB_LEGACY_METHODS = Collections.unmodifiableList(Arrays.asList(
METHOD_PREPARE_LEGACY_JOB,
METHOD_COMMIT_LEGACY_JOB,
METHOD_ROLLBACK_LEGACY_JOB,
METHOD_LEGACY_JOB_STATUS
));
public static final List<String> NATIVE_JOB_SESSION_METHODS = Collections.unmodifiableList(Arrays.asList(
METHOD_OPEN_NATIVE_JOB_SESSION,
METHOD_CLOSE_NATIVE_JOB_SESSION
));

static {
List<String> capabilities = new java.util.ArrayList<>(CAPABILITIES);
Expand Down Expand Up @@ -275,6 +296,7 @@ public final class AgentProtocol {
List<String> jdbcAllCapabilities = new java.util.ArrayList<>(MULTI_SESSION_ALL_CAPABILITIES);
jdbcAllCapabilities.add(CAPABILITY_STRUCTURED_ERROR_V1);
jdbcAllCapabilities.add(CAPABILITY_BLOB_BIND_STATEMENTS_V1);
jdbcAllCapabilities.add(CAPABILITY_NATIVE_JOB_LEGACY_OB_V1);
MULTI_SESSION_JDBC_ALL_CAPABILITIES = Collections.unmodifiableList(jdbcAllCapabilities);

List<String> methods = new java.util.ArrayList<>(COMMON_METHODS);
Expand Down Expand Up @@ -399,9 +421,19 @@ public static HandshakeResult multiSessionJdbcHandshakeResult() {
}

public static HandshakeResult multiSessionJdbcHandshakeResult(boolean blobBindStatements) {
if (!blobBindStatements) return multiSessionJdbcHandshakeResult();
return multiSessionJdbcHandshakeResult(blobBindStatements, false);
}

public static HandshakeResult multiSessionJdbcHandshakeResult(boolean blobBindStatements, boolean nativeJobSessions) {
return multiSessionJdbcHandshakeResult(blobBindStatements, nativeJobSessions, false);
}

public static HandshakeResult multiSessionJdbcHandshakeResult(boolean blobBindStatements, boolean nativeJobSessions, boolean nativeLegacyJobs) {
if (!blobBindStatements && !nativeJobSessions && !nativeLegacyJobs) return multiSessionJdbcHandshakeResult();
java.util.ArrayList<String> capabilities = new java.util.ArrayList<>(MULTI_SESSION_JDBC_CAPABILITIES);
capabilities.add(CAPABILITY_BLOB_BIND_STATEMENTS_V1);
if (blobBindStatements) capabilities.add(CAPABILITY_BLOB_BIND_STATEMENTS_V1);
if (nativeJobSessions) capabilities.add(CAPABILITY_NATIVE_JOB_SESSION_V1);
if (nativeLegacyJobs) capabilities.add(CAPABILITY_NATIVE_JOB_LEGACY_OB_V1);
return new HandshakeResult(MULTI_SESSION_PROTOCOL_VERSION, MULTI_SESSION_PROTOCOL_VERSION, capabilities);
}

Expand Down
6 changes: 6 additions & 0 deletions agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@ public interface DatabaseAgent {

default boolean supportsBlobBindStatements() { return false; }

/** Opt in to an operation-owned physical connection for native Job configuration. */
default boolean supportsNativeJobSessions() { return false; }

/** Opt in to OB's operation-owned Legacy SUBMIT/first-commit protocol. */
default boolean supportsNativeLegacyJobs() { return false; }

/** Opt in only when dictionary-derived owner/name values are preserved without fallback. */
default boolean supportsExactMetadataIdentity() { return false; }

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,29 @@ int poolCount() {
return pools.size();
}

/** Uses the existing physical budget and watchdogs, but never enters a reusable pool. */
Connection openDedicatedConnection(ConnectionFactory connectionFactory) throws SQLException {
if (closed.get()) throw new IllegalStateException("JDBC connection pool registry is closed");
SQLException failure = runtimeFailure.get();
if (failure != null) throw AgentRpcError.resource("connect", failure);
ConnectionFactoryDataSource source = new ConnectionFactoryDataSource(connectionFactory,
physicalConnectionBudget, physicalConnectionOpener, physicalConnectionCloser,
settings.connectionTimeoutMillis);
OperationDeadline deadline = new OperationDeadline(settings.connectionTimeoutMillis);
try (ConnectionFactoryDataSource.DeadlineRegistration ignored = source.registerCheckout(deadline)) {
Connection opened = source.getConnection();
// No Hikari setup follows on this path: the initialized physical connection is ready.
((HikariSetupTrackedConnection) opened).completeHikariSetup();
return opened;
} catch (SQLException error) {
if (contains(error, PhysicalConnectionLimitException.class)) {
throw AgentRpcError.backpressure("connect", error);
}
if (requiresRuntimeReplacement(error)) throw AgentRpcError.resource("connect", error);
throw error;
}
}

boolean isEnabled() {
return settings.enabled;
}
Expand Down
101 changes: 93 additions & 8 deletions agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ public final class JsonRpcServer {
.create();
private ConnectParams lastConnectParams;
private long lastConnectionValidationTimeMillis;
private NativeJobSession nativeJobSession;
private NativeLegacyJobSession nativeLegacyJobSession;

public JsonRpcServer(DatabaseAgent agent) {
this.agent = agent;
Expand Down Expand Up @@ -110,6 +112,7 @@ Object dispatchForRuntime(String method, JsonObject params) throws Exception {

private Object dispatchWithConnection(String method, JsonObject params) throws Exception {
return AgentExecutionContext.withJdbcExecutor(jdbcExecutor, () -> {
authorizeNativeJobRequest(method, params);
AbstractJdbcAgent jdbcAgent = pooledJdbcAgent();
if (AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method)
&& jdbcAgent != null
Expand Down Expand Up @@ -182,6 +185,28 @@ private void releaseIdlePooledConnection() {
}

private Object dispatch(String method, JsonObject params) throws Exception {
authorizeNativeJobRequest(method, params);
if (NativeLegacyJobSession.isControl(method)) return legacyJob(method, params);
if (AgentProtocol.METHOD_OPEN_NATIVE_JOB_SESSION.equals(method)) {
String operationId = NativeJobSession.operationId(params);
if (nativeJobSession != null) return nativeJobSession.openedFor(operationId);
if (!(agent instanceof AbstractJdbcAgent jdbcAgent) || !agent.supportsNativeJobSessions()) {
throw new UnsupportedOperationException("native_job_session_not_supported");
}
jdbcAgent.requireIdleNativeJobSession(jdbcExecutor);
nativeJobSession = new NativeJobSession(jdbcAgent, operationId);
try {
nativeJobSession.open();
return nativeJobSession.openedFor(operationId);
} catch (Exception | LinkageError error) {
if (!nativeJobSession.openOutcomeUnknown()) nativeJobSession = null;
throw error;
}
}
if (AgentProtocol.METHOD_CLOSE_NATIVE_JOB_SESSION.equals(method)) {
if (nativeJobSession == null) throw new IllegalStateException("native_job_session_unavailable");
return nativeJobSession.close(params);
}
if (agent instanceof AbstractJdbcAgent observingAgent) {
observingAgent.beforeAgentMethod(method, stringOrNull(params, "sessionId"));
}
Expand Down Expand Up @@ -424,22 +449,77 @@ private Object dispatch(String method, JsonObject params) throws Exception {
() -> agent.executeBatch(statements, stringOrNull(params, "schema")));
}
if (AgentProtocol.METHOD_DISCONNECT.equals(method)) {
jdbcExecutor.closeAllQuerySessions();
jdbcExecutor.closeAllTableReadSessions();
agent.disconnect();
lastConnectParams = null;
disconnect();
return Collections.singletonMap("ok", true);
}
if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
jdbcExecutor.closeAllQuerySessions();
jdbcExecutor.closeAllTableReadSessions();
agent.disconnect();
lastConnectParams = null;
disconnect();
return Collections.singletonMap("ok", true);
}
throw new IllegalArgumentException("Unknown method: " + method);
}

private void authorizeNativeJobRequest(String method, JsonObject params) throws Exception {
if (AgentProtocol.METHOD_HANDSHAKE.equals(method)
|| AgentProtocol.METHOD_DISCONNECT.equals(method) || AgentProtocol.METHOD_SHUTDOWN.equals(method)) return;
if (nativeLegacyJobSession != null && !NativeLegacyJobSession.isControl(method)) {
throw new IllegalStateException("native_job_session_reserved");
}
if (NativeLegacyJobSession.isControl(method)) {
if (nativeJobSession != null) throw new IllegalStateException("native_job_session_reserved");
return;
}
if (NativeJobSession.isControl(method)) return;
if (nativeJobSession != null) {
nativeJobSession.authorize(method, params);
} else if (params.has("nativeJobSessionToken") || params.has("nativeJobOperationId")) {
// A replacement logical session must not borrow a connection for an old lease token.
throw new IllegalStateException("native_job_session_unavailable");
}
}

private void disconnect() {
jdbcExecutor.closeAllQuerySessions();
jdbcExecutor.closeAllTableReadSessions();
JsonObject nativeClose = nativeJobSession == null ? null : nativeJobSession.dispose();
JsonObject legacyClose = nativeLegacyJobSession == null ? null : nativeLegacyJobSession.dispose();
agent.disconnect();
lastConnectParams = null;
if ((nativeClose != null && !"closed".equals(nativeClose.get("phase").getAsString()))
|| (legacyClose != null && (!legacyClose.has("disposalConfirmed") || !legacyClose.get("disposalConfirmed").getAsBoolean()))) {
throw AgentRpcError.resource("close", new JdbcConnectionPoolRegistry.PhysicalConnectionStateUnknownException(
new IllegalStateException("native_job_session_dispose_unknown")));
}
}

private JsonObject legacyJob(String method, JsonObject params) throws Exception {
String operationId = NativeJobSession.operationId(params);
if (AgentProtocol.METHOD_PREPARE_LEGACY_JOB.equals(method)) {
if (!(agent instanceof AbstractJdbcAgent jdbcAgent) || !agent.supportsNativeLegacyJobs()) {
throw new UnsupportedOperationException("native_job_legacy_protocol_unavailable");
}
NativeLegacyJobSession candidate;
try { candidate = new NativeLegacyJobSession(jdbcAgent, params); }
catch (IllegalArgumentException invalid) {
return NativeLegacyJobSession.failure(operationId, "not_started", invalid.getMessage());
}
if (nativeLegacyJobSession != null) {
return nativeLegacyJobSession.sameRequest(candidate) ? nativeLegacyJobSession.status(operationId)
: NativeLegacyJobSession.failure(operationId, "unknown", "native_job_operation_mismatch");
}
jdbcAgent.requireIdleNativeJobSession(jdbcExecutor);
nativeLegacyJobSession = candidate;
candidate.prepare();
return candidate.status(operationId);
}
if (nativeLegacyJobSession == null) {
return NativeLegacyJobSession.failure(operationId, "unknown", "native_job_legacy_session_unavailable");
}
if (AgentProtocol.METHOD_LEGACY_JOB_STATUS.equals(method)) return nativeLegacyJobSession.status(operationId);
return nativeLegacyJobSession.finish(operationId, stringOrNull(params, "jobId") == null ? "" : stringOrNull(params, "jobId"),
AgentProtocol.METHOD_COMMIT_LEGACY_JOB.equals(method));
}

private QueryResult executeStatementsWithTimeout(int timeoutSecs, DatabaseAgent.ThrowingSupplier<QueryResult> execute) throws Exception {
try (JdbcExecutor.NativeOperation operation = jdbcExecutor.beginNativeOperation()) {
return AgentExecutionContext.withStatementTimeout(timeoutSecs, operation,
Expand Down Expand Up @@ -504,6 +584,7 @@ private void switchCatalog(JsonObject params) throws Exception {
}

private void ensureLiveConnection(String method) {
if (agent instanceof AbstractJdbcAgent jdbcAgent && jdbcAgent.hasNativeJobConnection()) return;
if (!agent.permitsAutomaticReconnect()) return;
if (lastConnectParams == null || !shouldValidateConnection(method)) {
return;
Expand Down Expand Up @@ -547,6 +628,8 @@ private static boolean shouldValidateConnection(String method) {
return !AgentProtocol.METHOD_HANDSHAKE.equals(method)
&& !AgentProtocol.METHOD_CONNECT.equals(method)
&& !AgentProtocol.METHOD_TEST_CONNECTION.equals(method)
&& !NativeJobSession.isControl(method)
&& !NativeLegacyJobSession.isControl(method)
&& !AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method)
&& !AgentProtocol.METHOD_FETCH_QUERY_PAGE.equals(method)
&& !AgentProtocol.METHOD_CLOSE_QUERY_SESSION.equals(method)
Expand All @@ -567,6 +650,8 @@ private static boolean requiresConnectedConnection(String method) {
return !AgentProtocol.METHOD_HANDSHAKE.equals(method)
&& !AgentProtocol.METHOD_CONNECT.equals(method)
&& !AgentProtocol.METHOD_TEST_CONNECTION.equals(method)
&& !NativeJobSession.isControl(method)
&& !NativeLegacyJobSession.isControl(method)
&& !AgentProtocol.METHOD_DISCONNECT.equals(method)
&& !AgentProtocol.METHOD_SHUTDOWN.equals(method);
}
Expand Down
Loading
Loading