mirror of
https://github.com/Syngnat/GoNavi.git
synced 2026-08-10 16:53:35 +08:00
🐛 fix(sqlserver): 修复可选驱动查询消息透传缺失
- 为 optional-driver-agent 的 query 和 queryMulti 响应补充 messages 字段 - 在可选驱动 DB 客户端透传 SQL Server 查询提示信息与多结果集 - 补充 agent 与数据库层回归测试并更新 driver agent revision
This commit is contained in:
@@ -36,6 +36,7 @@ type agentResponse struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
Data interface{} `json:"data,omitempty"`
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
Messages []string `json:"messages,omitempty"`
|
||||
ChunkType string `json:"chunkType,omitempty"`
|
||||
RowsAffected int64 `json:"rowsAffected,omitempty"`
|
||||
}
|
||||
@@ -48,6 +49,7 @@ const (
|
||||
agentMethodOpenSession = "openSession"
|
||||
agentMethodCloseSession = "closeSession"
|
||||
agentMethodQuery = "query"
|
||||
agentMethodQueryMulti = "queryMulti"
|
||||
agentMethodStreamQuery = "streamQuery"
|
||||
agentMethodExec = "exec"
|
||||
agentMethodGetDatabases = "getDatabases"
|
||||
@@ -64,9 +66,9 @@ const (
|
||||
const legacyClickHouseDefaultTimeout = 2 * time.Hour
|
||||
|
||||
const (
|
||||
agentChunkColumns = "columns"
|
||||
agentChunkRows = "rows"
|
||||
agentChunkDone = "done"
|
||||
agentChunkColumns = "columns"
|
||||
agentChunkRows = "rows"
|
||||
agentChunkDone = "done"
|
||||
// agentStreamBatchSize 控制 driver-agent 向主进程发送 row chunk 的批次大小。
|
||||
// 调小到 64:单批 JSON 编码 + 主进程解码的瞬时内存峰值降为原来的 1/4,
|
||||
// 代价是 IPC 次数变为 4 倍,但每批仅一次 stdin/stdout 行读写,整体影响可忽略。
|
||||
@@ -236,12 +238,23 @@ func handleRequest(runtimeState *agentRuntime, req agentRequest) agentResponse {
|
||||
} else if ok {
|
||||
switch method {
|
||||
case agentMethodQuery:
|
||||
data, fields, err := queryStatementWithOptionalTimeout(session, req.Query, req.TimeoutMs)
|
||||
data, fields, messages, err := queryStatementWithMessagesOptionalTimeout(session, req.Query, req.TimeoutMs)
|
||||
if err != nil {
|
||||
return fail(resp, err.Error())
|
||||
}
|
||||
resp.Data = data
|
||||
resp.Fields = fields
|
||||
resp.Messages = messages
|
||||
case agentMethodQueryMulti:
|
||||
data, messages, supported, err := queryMultiStatementWithMessagesOptionalTimeout(session, req.Query, req.TimeoutMs)
|
||||
if err != nil {
|
||||
return fail(resp, err.Error())
|
||||
}
|
||||
if !supported {
|
||||
return fail(resp, "当前事务会话不支持多结果集查询")
|
||||
}
|
||||
resp.Data = data
|
||||
resp.Messages = messages
|
||||
case agentMethodExec:
|
||||
affected, err := execStatementWithOptionalTimeout(session, req.Query, req.TimeoutMs)
|
||||
if err != nil {
|
||||
@@ -260,12 +273,23 @@ func handleRequest(runtimeState *agentRuntime, req agentRequest) agentResponse {
|
||||
return fail(resp, err.Error())
|
||||
}
|
||||
case agentMethodQuery:
|
||||
data, fields, err := queryWithOptionalTimeout(runtimeState.inst, req.Query, req.TimeoutMs)
|
||||
data, fields, messages, err := queryWithMessagesOptionalTimeout(runtimeState.inst, req.Query, req.TimeoutMs)
|
||||
if err != nil {
|
||||
return fail(resp, err.Error())
|
||||
}
|
||||
resp.Data = data
|
||||
resp.Fields = fields
|
||||
resp.Messages = messages
|
||||
case agentMethodQueryMulti:
|
||||
data, messages, supported, err := queryMultiWithMessagesOptionalTimeout(runtimeState.inst, req.Query, req.TimeoutMs)
|
||||
if err != nil {
|
||||
return fail(resp, err.Error())
|
||||
}
|
||||
if !supported {
|
||||
return fail(resp, "当前驱动不支持原生多结果集查询")
|
||||
}
|
||||
resp.Data = data
|
||||
resp.Messages = messages
|
||||
case agentMethodExec:
|
||||
affected, err := execWithOptionalTimeout(runtimeState.inst, req.Query, req.TimeoutMs)
|
||||
if err != nil {
|
||||
@@ -581,6 +605,30 @@ type agentQueryContextRunner interface {
|
||||
QueryContext(context.Context, string) ([]map[string]interface{}, []string, error)
|
||||
}
|
||||
|
||||
type agentQueryMessageRunner interface {
|
||||
QueryWithMessages(query string) ([]map[string]interface{}, []string, []string, error)
|
||||
}
|
||||
|
||||
type agentQueryMessageContextRunner interface {
|
||||
QueryContextWithMessages(context.Context, string) ([]map[string]interface{}, []string, []string, error)
|
||||
}
|
||||
|
||||
type agentMultiResultMessageRunner interface {
|
||||
QueryMultiWithMessages(query string) ([]connection.ResultSetData, []string, error)
|
||||
}
|
||||
|
||||
type agentMultiResultMessageContextRunner interface {
|
||||
QueryMultiContextWithMessages(context.Context, string) ([]connection.ResultSetData, []string, error)
|
||||
}
|
||||
|
||||
type agentMultiResultRunner interface {
|
||||
QueryMulti(query string) ([]connection.ResultSetData, error)
|
||||
}
|
||||
|
||||
type agentMultiResultContextRunner interface {
|
||||
QueryMultiContext(context.Context, string) ([]connection.ResultSetData, error)
|
||||
}
|
||||
|
||||
type agentExecRunner interface {
|
||||
Exec(string) (int64, error)
|
||||
}
|
||||
@@ -589,20 +637,39 @@ type agentExecContextRunner interface {
|
||||
ExecContext(context.Context, string) (int64, error)
|
||||
}
|
||||
|
||||
func queryWithOptionalTimeout(inst agentQueryRunner, query string, timeoutMs int64) ([]map[string]interface{}, []string, error) {
|
||||
func queryWithMessagesOptionalTimeout(inst agentQueryRunner, query string, timeoutMs int64) ([]map[string]interface{}, []string, []string, error) {
|
||||
effectiveTimeoutMs := timeoutMs
|
||||
if effectiveTimeoutMs <= 0 && strings.EqualFold(strings.TrimSpace(agentDriverType), "clickhouse") {
|
||||
effectiveTimeoutMs = int64(legacyClickHouseDefaultTimeout / time.Millisecond)
|
||||
}
|
||||
if effectiveTimeoutMs <= 0 {
|
||||
return inst.Query(query)
|
||||
if q, ok := inst.(agentQueryMessageRunner); ok {
|
||||
return q.QueryWithMessages(query)
|
||||
}
|
||||
data, fields, err := inst.Query(query)
|
||||
return data, fields, nil, err
|
||||
}
|
||||
if q, ok := inst.(agentQueryMessageContextRunner); ok {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(effectiveTimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
return q.QueryContextWithMessages(ctx, query)
|
||||
}
|
||||
if q, ok := inst.(agentQueryContextRunner); ok {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(effectiveTimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
return q.QueryContext(ctx, query)
|
||||
data, fields, err := q.QueryContext(ctx, query)
|
||||
return data, fields, nil, err
|
||||
}
|
||||
return inst.Query(query)
|
||||
if q, ok := inst.(agentQueryMessageRunner); ok {
|
||||
return q.QueryWithMessages(query)
|
||||
}
|
||||
data, fields, err := inst.Query(query)
|
||||
return data, fields, nil, err
|
||||
}
|
||||
|
||||
func queryWithOptionalTimeout(inst agentQueryRunner, query string, timeoutMs int64) ([]map[string]interface{}, []string, error) {
|
||||
data, fields, _, err := queryWithMessagesOptionalTimeout(inst, query, timeoutMs)
|
||||
return data, fields, err
|
||||
}
|
||||
|
||||
func queryStatementWithOptionalTimeout(inst db.StatementExecer, query string, timeoutMs int64) ([]map[string]interface{}, []string, error) {
|
||||
@@ -613,6 +680,74 @@ func queryStatementWithOptionalTimeout(inst db.StatementExecer, query string, ti
|
||||
return queryWithOptionalTimeout(queryRunner, query, timeoutMs)
|
||||
}
|
||||
|
||||
func queryStatementWithMessagesOptionalTimeout(inst db.StatementExecer, query string, timeoutMs int64) ([]map[string]interface{}, []string, []string, error) {
|
||||
queryRunner, ok := inst.(agentQueryRunner)
|
||||
if !ok {
|
||||
return nil, nil, nil, fmt.Errorf("当前事务会话不支持查询语句")
|
||||
}
|
||||
return queryWithMessagesOptionalTimeout(queryRunner, query, timeoutMs)
|
||||
}
|
||||
|
||||
func queryMultiWithMessagesOptionalTimeout(inst db.Database, query string, timeoutMs int64) ([]connection.ResultSetData, []string, bool, error) {
|
||||
effectiveTimeoutMs := timeoutMs
|
||||
if effectiveTimeoutMs <= 0 && strings.EqualFold(strings.TrimSpace(agentDriverType), "clickhouse") {
|
||||
effectiveTimeoutMs = int64(legacyClickHouseDefaultTimeout / time.Millisecond)
|
||||
}
|
||||
if effectiveTimeoutMs > 0 {
|
||||
if q, ok := inst.(agentMultiResultMessageContextRunner); ok {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(effectiveTimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
data, messages, err := q.QueryMultiContextWithMessages(ctx, query)
|
||||
return data, messages, true, err
|
||||
}
|
||||
if q, ok := inst.(agentMultiResultContextRunner); ok {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(effectiveTimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
data, err := q.QueryMultiContext(ctx, query)
|
||||
return data, nil, true, err
|
||||
}
|
||||
}
|
||||
if q, ok := inst.(agentMultiResultMessageRunner); ok {
|
||||
data, messages, err := q.QueryMultiWithMessages(query)
|
||||
return data, messages, true, err
|
||||
}
|
||||
if q, ok := inst.(agentMultiResultRunner); ok {
|
||||
data, err := q.QueryMulti(query)
|
||||
return data, nil, true, err
|
||||
}
|
||||
return nil, nil, false, nil
|
||||
}
|
||||
|
||||
func queryMultiStatementWithMessagesOptionalTimeout(inst db.StatementExecer, query string, timeoutMs int64) ([]connection.ResultSetData, []string, bool, error) {
|
||||
effectiveTimeoutMs := timeoutMs
|
||||
if effectiveTimeoutMs <= 0 && strings.EqualFold(strings.TrimSpace(agentDriverType), "clickhouse") {
|
||||
effectiveTimeoutMs = int64(legacyClickHouseDefaultTimeout / time.Millisecond)
|
||||
}
|
||||
if effectiveTimeoutMs > 0 {
|
||||
if q, ok := inst.(agentMultiResultMessageContextRunner); ok {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(effectiveTimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
data, messages, err := q.QueryMultiContextWithMessages(ctx, query)
|
||||
return data, messages, true, err
|
||||
}
|
||||
if q, ok := inst.(agentMultiResultContextRunner); ok {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(effectiveTimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
data, err := q.QueryMultiContext(ctx, query)
|
||||
return data, nil, true, err
|
||||
}
|
||||
}
|
||||
if q, ok := inst.(agentMultiResultMessageRunner); ok {
|
||||
data, messages, err := q.QueryMultiWithMessages(query)
|
||||
return data, messages, true, err
|
||||
}
|
||||
if q, ok := inst.(agentMultiResultRunner); ok {
|
||||
data, err := q.QueryMulti(query)
|
||||
return data, nil, true, err
|
||||
}
|
||||
return nil, nil, false, nil
|
||||
}
|
||||
|
||||
func streamWithOptionalTimeout(inst db.StreamQueryExecer, query string, timeoutMs int64, consumer db.QueryStreamConsumer) error {
|
||||
effectiveTimeoutMs := timeoutMs
|
||||
if effectiveTimeoutMs <= 0 && strings.EqualFold(strings.TrimSpace(agentDriverType), "clickhouse") {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -4,26 +4,26 @@ package db
|
||||
|
||||
func init() {
|
||||
optionalDriverAgentRevisions = map[string]string{
|
||||
"mariadb": "src-b23e2ce1581a5064",
|
||||
"oceanbase": "src-5067dbdf0ca7b9c4",
|
||||
"diros": "src-db43faca6bf15d9b",
|
||||
"starrocks": "src-01e9f06c0fab09d5",
|
||||
"sphinx": "src-38ee5cae952cc809",
|
||||
"sqlserver": "src-7a87f6deb816f110",
|
||||
"sqlite": "src-d3d439cd788880e2",
|
||||
"duckdb": "src-b11506b8706bfb73",
|
||||
"dameng": "src-1638124bfd7fce09",
|
||||
"kingbase": "src-fb3a404cf4eb1bd9",
|
||||
"highgo": "src-72fe51afa884f6bc",
|
||||
"vastbase": "src-3d48607603bfd8b7",
|
||||
"opengauss": "src-709acf442f016e30",
|
||||
"gaussdb": "src-f6beccc924d71031",
|
||||
"iris": "src-9ebf5b970a73b341",
|
||||
"mongodb": "src-367d11cd04e982c1",
|
||||
"tdengine": "src-3c13c42f18ba01e1",
|
||||
"iotdb": "src-5ba9da13c6a272f9",
|
||||
"clickhouse": "src-99c8babfefdf142c",
|
||||
"elasticsearch": "src-36b2e2b5f49db9d1",
|
||||
"trino": "src-d264ceca132c185c",
|
||||
"mariadb": "src-cc133d2524ceb634",
|
||||
"oceanbase": "src-ac17327184366ff0",
|
||||
"diros": "src-7d4fe439271d0c56",
|
||||
"starrocks": "src-ce9ee22641a32f46",
|
||||
"sphinx": "src-08f5ae54efb3d9df",
|
||||
"sqlserver": "src-33b3b2c6dad5b3e6",
|
||||
"sqlite": "src-96dfa25b3042b2d5",
|
||||
"duckdb": "src-8804eb2cdbc89433",
|
||||
"dameng": "src-016e77082aea6718",
|
||||
"kingbase": "src-17728b2ebda94dc9",
|
||||
"highgo": "src-da2e8a9d2e661d3b",
|
||||
"vastbase": "src-da186ac367206c16",
|
||||
"opengauss": "src-54dc852e4c502947",
|
||||
"gaussdb": "src-3bbbffc6991dc8ae",
|
||||
"iris": "src-e798713e492e9a09",
|
||||
"mongodb": "src-2610395b35c2e708",
|
||||
"tdengine": "src-779b9b537f08856f",
|
||||
"iotdb": "src-7edea4aba8d4869e",
|
||||
"clickhouse": "src-0197342ca5afa8b5",
|
||||
"elasticsearch": "src-08e8e80cb17a409a",
|
||||
"trino": "src-ba947f211ce7b19f",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@ const (
|
||||
optionalAgentMethodOpenSession = "openSession"
|
||||
optionalAgentMethodCloseSession = "closeSession"
|
||||
optionalAgentMethodQuery = "query"
|
||||
optionalAgentMethodQueryMulti = "queryMulti"
|
||||
optionalAgentMethodStreamQuery = "streamQuery"
|
||||
optionalAgentMethodExec = "exec"
|
||||
optionalAgentMethodGetDatabases = "getDatabases"
|
||||
@@ -75,6 +76,7 @@ type optionalAgentResponse struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
Data json.RawMessage `json:"data,omitempty"`
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
Messages []string `json:"messages,omitempty"`
|
||||
ChunkType string `json:"chunkType,omitempty"`
|
||||
RowsAffected int64 `json:"rowsAffected,omitempty"`
|
||||
}
|
||||
@@ -106,7 +108,7 @@ func ProbeOptionalDriverAgentMetadata(driverType string, executablePath string)
|
||||
}()
|
||||
|
||||
var metadata OptionalDriverAgentMetadata
|
||||
if err := client.callWithTimeout(optionalAgentRequest{Method: optionalAgentMethodMetadata}, &metadata, nil, nil, optionalAgentMetadataProbeTimeout); err != nil {
|
||||
if err := client.callWithTimeout(optionalAgentRequest{Method: optionalAgentMethodMetadata}, &metadata, nil, nil, nil, optionalAgentMetadataProbeTimeout); err != nil {
|
||||
return OptionalDriverAgentMetadata{}, err
|
||||
}
|
||||
metadata.DriverType = normalizeRuntimeDriverType(metadata.DriverType)
|
||||
@@ -208,7 +210,7 @@ func (c *optionalDriverAgentClient) stderrText() string {
|
||||
return strings.TrimSpace(c.stderr.String())
|
||||
}
|
||||
|
||||
func (c *optionalDriverAgentClient) call(req optionalAgentRequest, out interface{}, fields *[]string, rowsAffected *int64) error {
|
||||
func (c *optionalDriverAgentClient) call(req optionalAgentRequest, out interface{}, fields *[]string, messages *[]string, rowsAffected *int64) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
@@ -252,6 +254,9 @@ func (c *optionalDriverAgentClient) call(req optionalAgentRequest, out interface
|
||||
if fields != nil {
|
||||
*fields = resp.Fields
|
||||
}
|
||||
if messages != nil {
|
||||
*messages = append((*messages)[:0], resp.Messages...)
|
||||
}
|
||||
if rowsAffected != nil {
|
||||
*rowsAffected = resp.RowsAffected
|
||||
}
|
||||
@@ -263,14 +268,14 @@ func (c *optionalDriverAgentClient) call(req optionalAgentRequest, out interface
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *optionalDriverAgentClient) callWithTimeout(req optionalAgentRequest, out interface{}, fields *[]string, rowsAffected *int64, timeout time.Duration) error {
|
||||
func (c *optionalDriverAgentClient) callWithTimeout(req optionalAgentRequest, out interface{}, fields *[]string, messages *[]string, rowsAffected *int64, timeout time.Duration) error {
|
||||
if timeout <= 0 {
|
||||
return c.call(req, out, fields, rowsAffected)
|
||||
return c.call(req, out, fields, messages, rowsAffected)
|
||||
}
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- c.call(req, out, fields, rowsAffected)
|
||||
errCh <- c.call(req, out, fields, messages, rowsAffected)
|
||||
}()
|
||||
|
||||
timer := time.NewTimer(timeout)
|
||||
@@ -469,7 +474,7 @@ func (d *OptionalDriverAgentDB) Connect(config connection.ConnectionConfig) erro
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodConnect,
|
||||
Config: &config,
|
||||
}, nil, nil, nil); err != nil {
|
||||
}, nil, nil, nil, nil); err != nil {
|
||||
_ = client.close()
|
||||
return err
|
||||
}
|
||||
@@ -482,7 +487,7 @@ func (d *OptionalDriverAgentDB) Close() error {
|
||||
if d.client == nil {
|
||||
return nil
|
||||
}
|
||||
_ = d.client.call(optionalAgentRequest{Method: optionalAgentMethodClose}, nil, nil, nil)
|
||||
_ = d.client.call(optionalAgentRequest{Method: optionalAgentMethodClose}, nil, nil, nil, nil)
|
||||
err := d.client.close()
|
||||
d.client = nil
|
||||
return err
|
||||
@@ -493,10 +498,87 @@ func (d *OptionalDriverAgentDB) Ping() error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return client.call(optionalAgentRequest{Method: optionalAgentMethodPing}, nil, nil, nil)
|
||||
return client.call(optionalAgentRequest{Method: optionalAgentMethodPing}, nil, nil, nil, nil)
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error) {
|
||||
data, fields, _, err := d.QueryContextWithMessages(ctx, query)
|
||||
return data, fields, err
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryContextWithMessages(ctx context.Context, query string) ([]map[string]interface{}, []string, []string, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
var data []map[string]interface{}
|
||||
var fields []string
|
||||
var messages []string
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodQuery,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, &data, &fields, &messages, nil); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
return data, fields, messages, nil
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) Query(query string) ([]map[string]interface{}, []string, error) {
|
||||
data, fields, _, err := d.QueryWithMessages(query)
|
||||
return data, fields, err
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryWithMessages(query string) ([]map[string]interface{}, []string, []string, error) {
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
var data []map[string]interface{}
|
||||
var fields []string
|
||||
var messages []string
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodQuery,
|
||||
Query: query,
|
||||
}, &data, &fields, &messages, nil); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
return data, fields, messages, nil
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryMulti(query string) ([]connection.ResultSetData, error) {
|
||||
results, _, err := d.QueryMultiWithMessages(query)
|
||||
return results, err
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryMultiWithMessages(query string) ([]connection.ResultSetData, []string, error) {
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var results []connection.ResultSetData
|
||||
var messages []string
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodQueryMulti,
|
||||
Query: query,
|
||||
}, &results, nil, &messages, nil); err != nil {
|
||||
if isOptionalAgentMultiResultUnsupportedError(err) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
return nil, nil, err
|
||||
}
|
||||
return results, messages, nil
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryMultiContext(ctx context.Context, query string) ([]connection.ResultSetData, error) {
|
||||
results, _, err := d.QueryMultiContextWithMessages(ctx, query)
|
||||
return results, err
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) QueryMultiContextWithMessages(ctx context.Context, query string) ([]connection.ResultSetData, []string, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
@@ -504,32 +586,19 @@ func (d *OptionalDriverAgentDB) QueryContext(ctx context.Context, query string)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var data []map[string]interface{}
|
||||
var fields []string
|
||||
var results []connection.ResultSetData
|
||||
var messages []string
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodQuery,
|
||||
Method: optionalAgentMethodQueryMulti,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, &data, &fields, nil); err != nil {
|
||||
}, &results, nil, &messages, nil); err != nil {
|
||||
if isOptionalAgentMultiResultUnsupportedError(err) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
return nil, nil, err
|
||||
}
|
||||
return data, fields, nil
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) Query(query string) ([]map[string]interface{}, []string, error) {
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var data []map[string]interface{}
|
||||
var fields []string
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodQuery,
|
||||
Query: query,
|
||||
}, &data, &fields, nil); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return data, fields, nil
|
||||
return results, messages, nil
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
@@ -581,7 +650,7 @@ func (d *OptionalDriverAgentDB) ExecContext(ctx context.Context, query string) (
|
||||
Method: optionalAgentMethodExec,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, nil, nil, &affected); err != nil {
|
||||
}, nil, nil, nil, &affected); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return affected, nil
|
||||
@@ -596,7 +665,7 @@ func (d *OptionalDriverAgentDB) Exec(query string) (int64, error) {
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodExec,
|
||||
Query: query,
|
||||
}, nil, nil, &affected); err != nil {
|
||||
}, nil, nil, nil, &affected); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return affected, nil
|
||||
@@ -611,7 +680,7 @@ func (d *OptionalDriverAgentDB) OpenSessionExecer(ctx context.Context) (Statemen
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodOpenSession,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, &sessionID, nil, nil); err != nil {
|
||||
}, &sessionID, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sessionID = strings.TrimSpace(sessionID)
|
||||
@@ -629,6 +698,10 @@ func (s *optionalDriverAgentSession) Query(query string) ([]map[string]interface
|
||||
return s.QueryContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) QueryWithMessages(query string) ([]map[string]interface{}, []string, []string, error) {
|
||||
return s.QueryContextWithMessages(context.Background(), query)
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return s.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
@@ -663,20 +736,26 @@ func (s *optionalDriverAgentSession) StreamQueryContext(ctx context.Context, que
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error) {
|
||||
data, fields, _, err := s.QueryContextWithMessages(ctx, query)
|
||||
return data, fields, err
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) QueryContextWithMessages(ctx context.Context, query string) ([]map[string]interface{}, []string, []string, error) {
|
||||
if err := s.ensureOpen(); err != nil {
|
||||
return nil, nil, err
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
var data []map[string]interface{}
|
||||
var fields []string
|
||||
var messages []string
|
||||
if err := s.client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodQuery,
|
||||
SessionID: s.sessionID,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, &data, &fields, nil); err != nil {
|
||||
return nil, nil, err
|
||||
}, &data, &fields, &messages, nil); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
return data, fields, nil
|
||||
return data, fields, messages, nil
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) Exec(query string) (int64, error) {
|
||||
@@ -693,7 +772,7 @@ func (s *optionalDriverAgentSession) ExecContext(ctx context.Context, query stri
|
||||
SessionID: s.sessionID,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, nil, nil, &affected); err != nil {
|
||||
}, nil, nil, nil, &affected); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return affected, nil
|
||||
@@ -714,7 +793,7 @@ func (s *optionalDriverAgentSession) Close() error {
|
||||
return s.client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodCloseSession,
|
||||
SessionID: sessionID,
|
||||
}, nil, nil, nil)
|
||||
}, nil, nil, nil, nil)
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) ensureOpen() error {
|
||||
@@ -740,6 +819,19 @@ func isOptionalAgentStreamUnsupportedError(err error) bool {
|
||||
return strings.Contains(text, "不支持的方法") || strings.Contains(text, "不支持流式查询")
|
||||
}
|
||||
|
||||
func isOptionalAgentMultiResultUnsupportedError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
text := strings.TrimSpace(err.Error())
|
||||
if text == "" {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(text, "不支持的方法") ||
|
||||
strings.Contains(text, "不支持原生多结果集查询") ||
|
||||
strings.Contains(text, "不支持多结果集查询")
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) GetDatabases() ([]string, error) {
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
@@ -748,7 +840,7 @@ func (d *OptionalDriverAgentDB) GetDatabases() ([]string, error) {
|
||||
var dbs []string
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodGetDatabases,
|
||||
}, &dbs, nil, nil); err != nil {
|
||||
}, &dbs, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return dbs, nil
|
||||
@@ -763,7 +855,7 @@ func (d *OptionalDriverAgentDB) GetTables(dbName string) ([]string, error) {
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodGetTables,
|
||||
DBName: dbName,
|
||||
}, &tables, nil, nil); err != nil {
|
||||
}, &tables, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return tables, nil
|
||||
@@ -779,7 +871,7 @@ func (d *OptionalDriverAgentDB) GetCreateStatement(dbName, tableName string) (st
|
||||
Method: optionalAgentMethodGetCreateStmt,
|
||||
DBName: dbName,
|
||||
TableName: tableName,
|
||||
}, &sqlText, nil, nil); err != nil {
|
||||
}, &sqlText, nil, nil, nil); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return sqlText, nil
|
||||
@@ -795,7 +887,7 @@ func (d *OptionalDriverAgentDB) GetColumns(dbName, tableName string) ([]connecti
|
||||
Method: optionalAgentMethodGetColumns,
|
||||
DBName: dbName,
|
||||
TableName: tableName,
|
||||
}, &columns, nil, nil); err != nil {
|
||||
}, &columns, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return columns, nil
|
||||
@@ -810,7 +902,7 @@ func (d *OptionalDriverAgentDB) GetAllColumns(dbName string) ([]connection.Colum
|
||||
if err := client.call(optionalAgentRequest{
|
||||
Method: optionalAgentMethodGetAllColumns,
|
||||
DBName: dbName,
|
||||
}, &columns, nil, nil); err != nil {
|
||||
}, &columns, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return columns, nil
|
||||
@@ -826,7 +918,7 @@ func (d *OptionalDriverAgentDB) GetIndexes(dbName, tableName string) ([]connecti
|
||||
Method: optionalAgentMethodGetIndexes,
|
||||
DBName: dbName,
|
||||
TableName: tableName,
|
||||
}, &indexes, nil, nil); err != nil {
|
||||
}, &indexes, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return indexes, nil
|
||||
@@ -842,7 +934,7 @@ func (d *OptionalDriverAgentDB) GetForeignKeys(dbName, tableName string) ([]conn
|
||||
Method: optionalAgentMethodGetForeignKeys,
|
||||
DBName: dbName,
|
||||
TableName: tableName,
|
||||
}, &keys, nil, nil); err != nil {
|
||||
}, &keys, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return keys, nil
|
||||
@@ -858,7 +950,7 @@ func (d *OptionalDriverAgentDB) GetTriggers(dbName, tableName string) ([]connect
|
||||
Method: optionalAgentMethodGetTriggers,
|
||||
DBName: dbName,
|
||||
TableName: tableName,
|
||||
}, &triggers, nil, nil); err != nil {
|
||||
}, &triggers, nil, nil, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return triggers, nil
|
||||
@@ -883,7 +975,7 @@ func (d *OptionalDriverAgentDB) ApplyChanges(tableName string, changes connectio
|
||||
Method: optionalAgentMethodApplyChanges,
|
||||
TableName: tableName,
|
||||
Changes: &changes,
|
||||
}, nil, nil, nil)
|
||||
}, nil, nil, nil, nil)
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) requireClient() (*optionalDriverAgentClient, error) {
|
||||
|
||||
@@ -136,3 +136,68 @@ func TestOptionalDriverAgentClientCallStreamQueryConsumesChunks(t *testing.T) {
|
||||
t.Fatalf("请求未使用 streamQuery 方法: %s", stdin.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOptionalDriverAgentDBQueryWithMessagesParsesAgentMessages(t *testing.T) {
|
||||
var stdin optionalAgentTestWriteCloser
|
||||
stdout := `{"id":1,"success":true,"data":[{"sql_text":"select 1"}],"fields":["sql_text"],"messages":["PRINT sql line 1","PRINT sql line 2"]}` + "\n"
|
||||
|
||||
dbInst := &OptionalDriverAgentDB{
|
||||
driverType: "sqlserver",
|
||||
client: &optionalDriverAgentClient{
|
||||
stdin: &stdin,
|
||||
reader: bufio.NewReader(strings.NewReader(stdout)),
|
||||
driver: "sqlserver",
|
||||
},
|
||||
}
|
||||
|
||||
rows, fields, messages, err := dbInst.QueryWithMessages("exec dbo.p_get_select")
|
||||
if err != nil {
|
||||
t.Fatalf("QueryWithMessages 返回错误: %v", err)
|
||||
}
|
||||
if len(rows) != 1 || rows[0]["sql_text"] != "select 1" {
|
||||
t.Fatalf("查询结果异常: %#v", rows)
|
||||
}
|
||||
if len(fields) != 1 || fields[0] != "sql_text" {
|
||||
t.Fatalf("字段异常: %#v", fields)
|
||||
}
|
||||
if len(messages) != 2 || messages[0] != "PRINT sql line 1" {
|
||||
t.Fatalf("消息异常: %#v", messages)
|
||||
}
|
||||
if !strings.Contains(stdin.String(), `"method":"query"`) {
|
||||
t.Fatalf("请求未使用 query 方法: %s", stdin.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOptionalDriverAgentDBQueryMultiWithMessagesParsesResultSets(t *testing.T) {
|
||||
var stdin optionalAgentTestWriteCloser
|
||||
stdout := `{"id":1,"success":true,"data":[{"statementIndex":1,"rows":[{"name":"master"}],"columns":["name"]},{"statementIndex":1,"rows":[],"columns":[],"messages":["PRINT generated sql"]}],"messages":["batch top-level message"]}` + "\n"
|
||||
|
||||
dbInst := &OptionalDriverAgentDB{
|
||||
driverType: "sqlserver",
|
||||
client: &optionalDriverAgentClient{
|
||||
stdin: &stdin,
|
||||
reader: bufio.NewReader(strings.NewReader(stdout)),
|
||||
driver: "sqlserver",
|
||||
},
|
||||
}
|
||||
|
||||
resultSets, messages, err := dbInst.QueryMultiWithMessages("exec dbo.p_get_select")
|
||||
if err != nil {
|
||||
t.Fatalf("QueryMultiWithMessages 返回错误: %v", err)
|
||||
}
|
||||
if len(resultSets) != 2 {
|
||||
t.Fatalf("结果集数量异常: %#v", resultSets)
|
||||
}
|
||||
if got := resultSets[0].Rows[0]["name"]; got != "master" {
|
||||
t.Fatalf("首个结果集异常,got=%v", got)
|
||||
}
|
||||
if len(resultSets[1].Messages) != 1 || resultSets[1].Messages[0] != "PRINT generated sql" {
|
||||
t.Fatalf("消息结果集异常: %#v", resultSets[1])
|
||||
}
|
||||
if len(messages) != 1 || messages[0] != "batch top-level message" {
|
||||
t.Fatalf("顶层消息异常: %#v", messages)
|
||||
}
|
||||
if !strings.Contains(stdin.String(), `"method":"queryMulti"`) {
|
||||
t.Fatalf("请求未使用 queryMulti 方法: %s", stdin.String())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user