🐛 fix(sqlserver): 修复可选驱动查询消息透传缺失

- 为 optional-driver-agent 的 query 和 queryMulti 响应补充 messages 字段
- 在可选驱动 DB 客户端透传 SQL Server 查询提示信息与多结果集
- 补充 agent 与数据库层回归测试并更新 driver agent revision
This commit is contained in:
Syngnat
2026-06-23 08:48:42 +08:00
parent 8f1e6cf379
commit 495a985ae1
5 changed files with 472 additions and 77 deletions

View File

@@ -101,6 +101,9 @@ type fakeAgentTimeoutDB struct {
execCalled bool
execContextCalled bool
deadlineSet bool
queryMessages []string
multiResults []connection.ResultSetData
multiMessages []string
}
func (f *fakeAgentTimeoutDB) Connect(config connection.ConnectionConfig) error { return nil }
@@ -117,6 +120,14 @@ func (f *fakeAgentTimeoutDB) QueryContext(ctx context.Context, query string) ([]
}
return []map[string]interface{}{{"ok": 1}}, []string{"ok"}, nil
}
func (f *fakeAgentTimeoutDB) QueryWithMessages(query string) ([]map[string]interface{}, []string, []string, error) {
data, fields, err := f.QueryContext(context.Background(), query)
return data, fields, append([]string(nil), f.queryMessages...), err
}
func (f *fakeAgentTimeoutDB) QueryContextWithMessages(ctx context.Context, query string) ([]map[string]interface{}, []string, []string, error) {
data, fields, err := f.QueryContext(ctx, query)
return data, fields, append([]string(nil), f.queryMessages...), err
}
func (f *fakeAgentTimeoutDB) Exec(query string) (int64, error) {
f.execCalled = true
return 0, errors.New("exec should not be called")
@@ -150,6 +161,15 @@ func (f *fakeAgentTimeoutDB) GetForeignKeys(dbName, tableName string) ([]connect
func (f *fakeAgentTimeoutDB) GetTriggers(dbName, tableName string) ([]connection.TriggerDefinition, error) {
return nil, nil
}
func (f *fakeAgentTimeoutDB) QueryMultiWithMessages(query string) ([]connection.ResultSetData, []string, error) {
return append([]connection.ResultSetData(nil), f.multiResults...), append([]string(nil), f.multiMessages...), nil
}
func (f *fakeAgentTimeoutDB) QueryMultiContextWithMessages(ctx context.Context, query string) ([]connection.ResultSetData, []string, error) {
if _, ok := ctx.Deadline(); ok {
f.deadlineSet = true
}
return f.QueryMultiWithMessages(query)
}
type fakeAgentSessionDB struct {
fakeAgentTimeoutDB
@@ -165,6 +185,7 @@ type fakeAgentStatementSession struct {
queryCalls int
execCalls int
closed bool
messages []string
}
func (f *fakeAgentStatementSession) Query(query string) ([]map[string]interface{}, []string, error) {
@@ -175,6 +196,14 @@ func (f *fakeAgentStatementSession) QueryContext(ctx context.Context, query stri
f.queryCalls++
return []map[string]interface{}{{"session_ok": 1}}, []string{"session_ok"}, nil
}
func (f *fakeAgentStatementSession) QueryWithMessages(query string) ([]map[string]interface{}, []string, []string, error) {
data, fields, err := f.QueryContext(context.Background(), query)
return data, fields, append([]string(nil), f.messages...), err
}
func (f *fakeAgentStatementSession) QueryContextWithMessages(ctx context.Context, query string) ([]map[string]interface{}, []string, []string, error) {
data, fields, err := f.QueryContext(ctx, query)
return data, fields, append([]string(nil), f.messages...), err
}
func (f *fakeAgentStatementSession) Exec(query string) (int64, error) {
return f.ExecContext(context.Background(), query)
@@ -297,6 +326,77 @@ func TestQueryWithOptionalTimeout_ClickHouseLegacyModeUsesQueryContext(t *testin
}
}
func TestHandleRequest_QueryIncludesServerMessages(t *testing.T) {
old := agentDriverType
defer func() { agentDriverType = old }()
agentDriverType = "sqlserver"
fake := &fakeAgentTimeoutDB{
queryMessages: []string{"PRINT sql line 1", "PRINT sql line 2"},
}
runtimeState := &agentRuntime{inst: fake, sessions: make(map[string]db.StatementExecer)}
resp := handleRequest(runtimeState, agentRequest{
ID: 11,
Method: agentMethodQuery,
Query: "exec dbo.p_get_select",
TimeoutMs: int64((2 * time.Second).Milliseconds()),
})
if !resp.Success {
t.Fatalf("query request failed: %s", resp.Error)
}
if len(resp.Messages) != 2 || resp.Messages[0] != "PRINT sql line 1" {
t.Fatalf("expected query messages to be preserved, got %#v", resp.Messages)
}
}
func TestHandleRequest_QueryMultiIncludesResultSetsAndMessages(t *testing.T) {
old := agentDriverType
defer func() { agentDriverType = old }()
agentDriverType = "sqlserver"
fake := &fakeAgentTimeoutDB{
multiResults: []connection.ResultSetData{
{
StatementIndex: 1,
Rows: []map[string]interface{}{{"name": "master"}},
Columns: []string{"name"},
},
{
StatementIndex: 1,
Rows: []map[string]interface{}{},
Columns: []string{},
Messages: []string{"PRINT generated sql"},
},
},
multiMessages: []string{"batch top-level message"},
}
runtimeState := &agentRuntime{inst: fake, sessions: make(map[string]db.StatementExecer)}
resp := handleRequest(runtimeState, agentRequest{
ID: 12,
Method: agentMethodQueryMulti,
Query: "exec dbo.p_get_select",
TimeoutMs: int64((2 * time.Second).Milliseconds()),
})
if !resp.Success {
t.Fatalf("queryMulti request failed: %s", resp.Error)
}
if len(resp.Messages) != 1 || resp.Messages[0] != "batch top-level message" {
t.Fatalf("expected top-level messages to be preserved, got %#v", resp.Messages)
}
resultSets, ok := resp.Data.([]connection.ResultSetData)
if !ok {
t.Fatalf("expected []connection.ResultSetData, got %T", resp.Data)
}
if len(resultSets) != 2 {
t.Fatalf("expected 2 result sets, got %#v", resultSets)
}
if len(resultSets[1].Messages) != 1 || resultSets[1].Messages[0] != "PRINT generated sql" {
t.Fatalf("expected message-only result set to be preserved, got %#v", resultSets[1])
}
}
func TestHandleRequest_UsesPinnedSessionForSessionScopedQueryAndExec(t *testing.T) {
old := agentDriverType
defer func() { agentDriverType = old }()
@@ -329,6 +429,9 @@ func TestHandleRequest_UsesPinnedSessionForSessionScopedQueryAndExec(t *testing.
if !queryResp.Success {
t.Fatalf("session query failed: %s", queryResp.Error)
}
if len(queryResp.Messages) != 0 {
t.Fatalf("expected empty default session messages, got %#v", queryResp.Messages)
}
if fake.queryCalled || fake.queryContextCalled {
t.Fatalf("expected session query to bypass database-level query path, got Query=%v QueryContext=%v", fake.queryCalled, fake.queryContextCalled)
}