github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(sqlserver): keep table variable scope in agent batches

Browse Source
This commit is contained in:
zipg authored and GitHub committed 2026-09-03 13:43:45 +08:00
1 parent 4b7b72dc65
commit 76c51b5fa6
6 files changed
+334 -11
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java
+9 -2
View File
@@ -196,7 +196,9 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
options.getMaxRows(),
options.getFetchSize(),
options.getTimeoutSecs(),
resultValueReader()
resultValueReader(),
JdbcExecutor.StatementMessageReader.NONE,
advancePastUpdateCounts()
);
}
@@ -211,7 +213,8 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
this::setSchemaSQL,
this::resetSchemaSQL,
options,
resultValueReader()
resultValueReader(),
advancePastUpdateCounts()
);
}
@@ -529,6 +532,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
protected void beforeQueryExecution(Connection connection, int timeoutSecs) throws Exception {
}
protected boolean advancePastUpdateCounts() {
return false;
}
protected void beforePooledConnectionReturn(Connection connection) throws Exception {
}
agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java
+110 -4
View File
@@ -102,6 +102,34 @@ public final class JdbcExecutor {
int timeoutSecs,
ResultValueReader valueReader,
StatementMessageReader statementMessageReader
) {
return execute(
conn,
sql,
schema,
setSchemaSql,
resetSchemaSql,
maxRows,
fetchSize,
timeoutSecs,
valueReader,
statementMessageReader,
false
);
}
public QueryResult execute(
Connection conn,
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
int maxRows,
Integer fetchSize,
int timeoutSecs,
ResultValueReader valueReader,
StatementMessageReader statementMessageReader,
boolean advancePastUpdateCounts
) {
return unchecked(() -> {
String trimmedSql = trimSql(sql);
@@ -122,17 +150,31 @@ public final class JdbcExecutor {
// Do not translate them to Connection.commit(), which requires autoCommit=false.
boolean hasResultSet = stmt.execute(trimmedSql);
long elapsed = System.currentTimeMillis() - start;
long affectedRows;
if (advancePastUpdateCounts) {
affectedRows = 0L;
while (!hasResultSet) {
int updateCount = stmt.getUpdateCount();
if (updateCount < 0) {
break;
}
affectedRows += updateCount;
hasResultSet = stmt.getMoreResults();
}
} else {
int updateCount = hasResultSet ? -1 : stmt.getUpdateCount();
affectedRows = updateCount >= 0 ? updateCount : 0L;
}
QueryResult result;
if (hasResultSet) {
try (ResultSet rs = stmt.getResultSet()) {
result = readResultSet(rs, elapsed, effectiveMaxRows, valueReader);
}
} else {
int updateCount = stmt.getUpdateCount();
result = new QueryResult(
Collections.emptyList(),
Collections.emptyList(),
updateCount >= 0 ? updateCount : 0,
affectedRows,
elapsed,
false
);
@@ -240,6 +282,30 @@ public final class JdbcExecutor {
return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, StatementMessageReader.NONE, sessions);
}
public QueryPageResult executePage(
Connection conn,
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
QueryPageOptions options,
ResultValueReader valueReader,
boolean advancePastUpdateCounts
) {
return executePage(
conn,
sql,
schema,
setSchemaSql,
resetSchemaSql,
options,
valueReader,
StatementMessageReader.NONE,
sessions,
advancePastUpdateCounts
);
}
public QueryPageResult startTableRead(
Connection conn,
String sql,
@@ -273,6 +339,32 @@ public final class JdbcExecutor {
ResultValueReader valueReader,
StatementMessageReader statementMessageReader,
ConcurrentHashMap<String, QuerySession> targetSessions
) {
return executePage(
conn,
sql,
schema,
setSchemaSql,
resetSchemaSql,
options,
valueReader,
statementMessageReader,
targetSessions,
false
);
}
private QueryPageResult executePage(
Connection conn,
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
QueryPageOptions options,
ResultValueReader valueReader,
StatementMessageReader statementMessageReader,
ConcurrentHashMap<String, QuerySession> targetSessions,
boolean advancePastUpdateCounts
) {
return unchecked(() -> {
expireIdleSessions(targetSessions, System.currentTimeMillis(), QUERY_SESSION_IDLE_TIMEOUT_MILLIS);
@@ -293,12 +385,26 @@ public final class JdbcExecutor {
// JDBC transaction APIs are reserved for executeTransaction.
boolean hasResultSet = stmt.execute(trimmedSql);
long elapsed = System.currentTimeMillis() - start;
long affectedRows;
if (advancePastUpdateCounts) {
affectedRows = 0L;
while (!hasResultSet) {
int updateCount = stmt.getUpdateCount();
if (updateCount < 0) {
break;
}
affectedRows += updateCount;
hasResultSet = stmt.getMoreResults();
}
} else {
int updateCount = hasResultSet ? -1 : stmt.getUpdateCount();
affectedRows = updateCount >= 0 ? updateCount : 0L;
}
if (!hasResultSet) {
int updateCount = stmt.getUpdateCount();
QueryResult result = new QueryResult(
Collections.emptyList(),
Collections.emptyList(),
updateCount >= 0 ? updateCount : 0,
affectedRows,
elapsed,
false
);