mirror of
https://github.com/Syngnat/GoNavi.git
synced 2026-07-09 14:31:35 +08:00
Merge remote-tracking branch 'origin/dev' into feature/20260602_connection_driver_i18n
# Conflicts: # frontend/src/App.tsx # frontend/src/components/AISettingsModal.tsx # frontend/src/components/ConnectionModal.edit-password.test.tsx # frontend/src/components/ConnectionModal.tsx # frontend/src/components/DataSyncModal.i18n.test.ts # frontend/src/components/DataSyncModal.tsx # frontend/src/components/QueryEditor.external-sql-save.test.tsx # frontend/src/components/QueryEditor.tsx # frontend/src/components/Sidebar.locate-toolbar.test.tsx # frontend/src/components/Sidebar.tsx # frontend/src/components/SnippetSettingsModal.tsx # frontend/src/components/TableOverview.tsx # frontend/src/components/ai/AIChatHeader.test.tsx # frontend/src/components/ai/AISettingsProvidersSection.tsx # frontend/src/components/ai/aiChatPayloadDispatch.ts # frontend/src/components/ai/aiChatReadiness.ts # frontend/src/components/ai/aiSettingsModalConfig.tsx # frontend/src/components/ai/messageBubble/AIMessageCodeBlock.tsx # frontend/src/components/sidebarV2Utils.ts # frontend/src/i18n/catalog.test.ts # frontend/src/utils/connectionTypeCatalog.test.ts # frontend/src/utils/connectionTypeCatalog.ts # frontend/src/utils/tabDisplay.ts # internal/ai/provider/custom.go # internal/ai/service/service.go # internal/app/methods_driver.go # internal/app/methods_file.go # internal/db/custom_impl.go # internal/db/iris_impl.go # internal/db/mariadb_impl.go # internal/db/sqlserver_impl.go # shared/i18n/de-DE.json # shared/i18n/en-US.json # shared/i18n/ja-JP.json # shared/i18n/ru-RU.json # shared/i18n/zh-CN.json # shared/i18n/zh-TW.json
This commit is contained in:
@@ -3,15 +3,18 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
@@ -208,6 +211,10 @@ func (c *ClickHouseDB) buildClickHouseOptionsWithHTTPCompatibility(config connec
|
||||
|
||||
type clickHouseHTTPClientProtocolVersionStripper struct {
|
||||
next http.RoundTripper
|
||||
// serverHelloRewritten 保证只对每个连接的首个握手探测请求改写一次,
|
||||
// 避免连接建立之后误改写恰好相同的用户查询(clickhouse-go 的 queryHello
|
||||
// 始终是连接上的第一个 HTTP 请求)。
|
||||
serverHelloRewritten *atomic.Bool
|
||||
}
|
||||
|
||||
func (rt clickHouseHTTPClientProtocolVersionStripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
@@ -218,19 +225,81 @@ func (rt clickHouseHTTPClientProtocolVersionStripper) RoundTrip(req *http.Reques
|
||||
if req == nil || req.URL == nil {
|
||||
return next.RoundTrip(req)
|
||||
}
|
||||
|
||||
query := req.URL.Query()
|
||||
if _, ok := query["client_protocol_version"]; !ok {
|
||||
stripParam := false
|
||||
if _, ok := query["client_protocol_version"]; ok {
|
||||
stripParam = true
|
||||
}
|
||||
|
||||
var (
|
||||
rewrittenBody []byte
|
||||
hadServerInfoQuery bool
|
||||
err error
|
||||
)
|
||||
// 仅在握手阶段(首个匹配请求)改写探测查询;后续用户查询一律放行。
|
||||
if rt.serverHelloRewritten == nil || !rt.serverHelloRewritten.Load() {
|
||||
rewrittenBody, hadServerInfoQuery, err = rewriteClickHouseServerHelloRequestBody(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if hadServerInfoQuery && rt.serverHelloRewritten != nil {
|
||||
rt.serverHelloRewritten.Store(true)
|
||||
}
|
||||
}
|
||||
|
||||
if !stripParam && !hadServerInfoQuery {
|
||||
return next.RoundTrip(req)
|
||||
}
|
||||
|
||||
cloned := req.Clone(req.Context())
|
||||
clonedURL := *req.URL
|
||||
query.Del("client_protocol_version")
|
||||
clonedURL.RawQuery = query.Encode()
|
||||
cloned.URL = &clonedURL
|
||||
if stripParam {
|
||||
clonedURL := *req.URL
|
||||
query.Del("client_protocol_version")
|
||||
clonedURL.RawQuery = query.Encode()
|
||||
cloned.URL = &clonedURL
|
||||
}
|
||||
if hadServerInfoQuery {
|
||||
cloned.Body = io.NopCloser(bytes.NewReader(rewrittenBody))
|
||||
cloned.ContentLength = int64(len(rewrittenBody))
|
||||
cloned.GetBody = func() (io.ReadCloser, error) {
|
||||
return io.NopCloser(bytes.NewReader(rewrittenBody)), nil
|
||||
}
|
||||
}
|
||||
return next.RoundTrip(cloned)
|
||||
}
|
||||
|
||||
// clickHouseServerHelloQuery 是 clickhouse-go HTTP 驱动在握手阶段发送的服务端信息探测语句。
|
||||
// 旧版本服务端(如 ClickHouse 22.8)没有 displayName() 函数,会直接返回 UNKNOWN_FUNCTION。
|
||||
const clickHouseServerHelloQuery = "SELECT displayName(), version(), revision(), timezone()"
|
||||
|
||||
// clickHouseServerHelloCompatQuery 使用 hostName() 替换不存在的 displayName()。
|
||||
// hostName() 在所有受支持的 ClickHouse 版本上都可用,并返回服务端主机名,
|
||||
// 足以填充驱动握手所需的显示名称字段,其余 version()/revision()/timezone() 保持不变。
|
||||
const clickHouseServerHelloCompatQuery = "SELECT hostName(), version(), revision(), timezone()"
|
||||
|
||||
// rewriteClickHouseServerHelloRequestBody 检测并改写握手探测请求体,将 displayName() 替换为
|
||||
// hostName()。仅当请求体恰好是驱动的握手探测语句时才改写,其它请求体一律原样放行。
|
||||
func rewriteClickHouseServerHelloRequestBody(req *http.Request) ([]byte, bool, error) {
|
||||
if req == nil || req.Body == nil || req.Body == http.NoBody {
|
||||
return nil, false, nil
|
||||
}
|
||||
body, err := io.ReadAll(req.Body)
|
||||
closeErr := req.Body.Close()
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if closeErr != nil {
|
||||
return nil, false, closeErr
|
||||
}
|
||||
// 恢复原始请求体,保证非握手请求不受影响。
|
||||
req.Body = io.NopCloser(bytes.NewReader(body))
|
||||
if strings.TrimSpace(string(body)) != clickHouseServerHelloQuery {
|
||||
return nil, false, nil
|
||||
}
|
||||
return []byte(clickHouseServerHelloCompatQuery), true, nil
|
||||
}
|
||||
|
||||
func installClickHouseHTTPClientProtocolVersionStripper(opts *clickhouse.Options) {
|
||||
if opts == nil {
|
||||
return
|
||||
@@ -247,7 +316,10 @@ func installClickHouseHTTPClientProtocolVersionStripper(opts *clickhouse.Options
|
||||
next = wrapped
|
||||
}
|
||||
}
|
||||
return clickHouseHTTPClientProtocolVersionStripper{next: next}, nil
|
||||
return clickHouseHTTPClientProtocolVersionStripper{
|
||||
next: next,
|
||||
serverHelloRewritten: &atomic.Bool{},
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
@@ -462,9 +534,32 @@ func isClickHouseHTTPClientProtocolVersionUnsupported(err error) bool {
|
||||
strings.Contains(text, "code: 115")
|
||||
}
|
||||
|
||||
// isClickHouseHTTPServerInfoFunctionUnsupported 识别 clickhouse-go 在 HTTP 握手阶段
|
||||
// 执行 "SELECT displayName(), version(), revision(), timezone()" 时,旧版本服务端
|
||||
// (如 ClickHouse 22.8)因不存在 displayName() 函数而返回的 Code 46 / UNKNOWN_FUNCTION 错误。
|
||||
func isClickHouseHTTPServerInfoFunctionUnsupported(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
text := strings.ToLower(strings.TrimSpace(err.Error()))
|
||||
if text == "" || !strings.Contains(text, "displayname") {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(text, "unknown function") ||
|
||||
strings.Contains(text, "unknown_function") ||
|
||||
strings.Contains(text, "code: 46")
|
||||
}
|
||||
|
||||
// shouldRetryClickHouseHTTPCompatibility 判断 HTTP 协议下的失败是否可以通过
|
||||
// HTTP 兼容模式(移除 client_protocol_version 并改写握手探测查询)重试解决。
|
||||
func shouldRetryClickHouseHTTPCompatibility(err error) bool {
|
||||
return isClickHouseHTTPClientProtocolVersionUnsupported(err) ||
|
||||
isClickHouseHTTPServerInfoFunctionUnsupported(err)
|
||||
}
|
||||
|
||||
func shouldTryNextClickHouseProtocol(protocol clickhouse.Protocol, err error) bool {
|
||||
return isClickHouseProtocolMismatch(err) ||
|
||||
(protocol == clickhouse.HTTP && isClickHouseHTTPClientProtocolVersionUnsupported(err))
|
||||
(protocol == clickhouse.HTTP && shouldRetryClickHouseHTTPCompatibility(err))
|
||||
}
|
||||
|
||||
func clickHouseProtocolName(protocol clickhouse.Protocol) string {
|
||||
@@ -510,6 +605,9 @@ func clickHouseAttemptFailureMessage(protocol clickhouse.Protocol, err error) st
|
||||
if protocol == clickhouse.HTTP && isClickHouseHTTPClientProtocolVersionUnsupported(err) {
|
||||
return localizedDriverRuntimeText("db.backend.error.clickhouse_http_client_protocol_version_unsupported", nil)
|
||||
}
|
||||
if protocol == clickhouse.HTTP && isClickHouseHTTPServerInfoFunctionUnsupported(err) {
|
||||
return "当前 ClickHouse HTTP 端口不支持 displayName() 握手探测函数(常见于 ClickHouse 22.8),将使用 HTTP 兼容模式重试;如仍失败请确认连接协议和端口"
|
||||
}
|
||||
if isClickHouseProtocolMismatch(err) {
|
||||
if protocol == clickhouse.Native {
|
||||
return localizedDriverRuntimeText("db.backend.error.clickhouse_native_protocol_mismatch", nil)
|
||||
@@ -665,6 +763,7 @@ func (c *ClickHouseDB) Connect(config connection.ConnectionConfig) error {
|
||||
break
|
||||
}
|
||||
c.conn = clickhouse.OpenDB(opts)
|
||||
configureSQLConnectionPool(c.conn, "clickhouse")
|
||||
if err := c.Ping(); err != nil {
|
||||
lastProtocolErr = err
|
||||
failureMessage := clickHouseAttemptFailureMessage(protocol, err)
|
||||
@@ -677,9 +776,13 @@ func (c *ClickHouseDB) Connect(config connection.ConnectionConfig) error {
|
||||
}
|
||||
if protocol == clickhouse.HTTP &&
|
||||
!stripHTTPClientProtocolVersion &&
|
||||
isClickHouseHTTPClientProtocolVersionUnsupported(err) &&
|
||||
shouldRetryClickHouseHTTPCompatibility(err) &&
|
||||
compatIdx+1 < len(compatibilityModes) {
|
||||
logger.Warnf("ClickHouse HTTP 端口不支持 client_protocol_version,改用 HTTP 兼容模式重试")
|
||||
if isClickHouseHTTPServerInfoFunctionUnsupported(err) {
|
||||
logger.Warnf("ClickHouse HTTP 端口不支持 displayName() 握手探测函数,改用 HTTP 兼容模式重试")
|
||||
} else {
|
||||
logger.Warnf("ClickHouse HTTP 端口不支持 client_protocol_version,改用 HTTP 兼容模式重试")
|
||||
}
|
||||
continue
|
||||
}
|
||||
break
|
||||
@@ -799,6 +902,22 @@ func (c *ClickHouseDB) Query(query string) ([]map[string]interface{}, []string,
|
||||
return scanRows(rows)
|
||||
}
|
||||
|
||||
func (c *ClickHouseDB) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if c.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
rows, err := c.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
return streamRows(rows, consumer)
|
||||
}
|
||||
|
||||
func (c *ClickHouseDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return c.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (c *ClickHouseDB) ExecContext(ctx context.Context, query string) (int64, error) {
|
||||
if c.conn == nil {
|
||||
return 0, fmt.Errorf("连接未打开")
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -724,6 +725,170 @@ func TestClickHouseHTTPClientProtocolVersionStripperRemovesDriverQueryParam(t *t
|
||||
}
|
||||
}
|
||||
|
||||
func TestClickHouseHTTPServerInfoFunctionUnsupportedEnablesCompatibilityRetry(t *testing.T) {
|
||||
err := errors.New(`failed to query server hello: failed to query server hello info: sendQuery: [HTTP 404] response body: "Code: 46. DB::Exception: Unknown function displayName: While processing displayName(), version(), revision(), timezone(). (UNKNOWN_FUNCTION)"`)
|
||||
if !isClickHouseHTTPServerInfoFunctionUnsupported(err) {
|
||||
t.Fatalf("expected displayName unknown function to be treated as HTTP server-info compatibility issue")
|
||||
}
|
||||
if !shouldRetryClickHouseHTTPCompatibility(err) {
|
||||
t.Fatalf("expected displayName unknown function to permit HTTP compatibility retry")
|
||||
}
|
||||
if !shouldTryNextClickHouseProtocol(clickhouse.HTTP, err) {
|
||||
t.Fatalf("expected HTTP displayName issue to permit protocol fallback")
|
||||
}
|
||||
if shouldTryNextClickHouseProtocol(clickhouse.Native, err) {
|
||||
t.Fatalf("native protocol should not treat HTTP displayName issue as retryable")
|
||||
}
|
||||
|
||||
message := clickHouseAttemptFailureMessage(clickhouse.HTTP, err)
|
||||
if !strings.Contains(message, "displayName") || !strings.Contains(message, "兼容模式") {
|
||||
t.Fatalf("expected displayName compatibility retry hint, got %q", message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsClickHouseHTTPServerInfoFunctionUnsupportedIgnoresUnrelatedErrors(t *testing.T) {
|
||||
if isClickHouseHTTPServerInfoFunctionUnsupported(nil) {
|
||||
t.Fatal("nil error should not be treated as server-info function issue")
|
||||
}
|
||||
if isClickHouseHTTPServerInfoFunctionUnsupported(errors.New("[HTTP 404] page not found")) {
|
||||
t.Fatal("plain 404 without displayName signal should not be treated as server-info function issue")
|
||||
}
|
||||
if isClickHouseHTTPServerInfoFunctionUnsupported(errors.New("Code: 60. DB::Exception: Unknown function someOtherFn")) {
|
||||
t.Fatal("unknown function error without displayName should not be treated as server-info function issue")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClickHouseHTTPCompatibilityStripperRewritesServerHelloQuery(t *testing.T) {
|
||||
var seenBody string
|
||||
stripper := clickHouseHTTPClientProtocolVersionStripper{
|
||||
next: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.Body != nil {
|
||||
data, err := io.ReadAll(req.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
seenBody = string(data)
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
req, err := http.NewRequest(
|
||||
http.MethodPost,
|
||||
"http://clickhouse.local:8123/?database=default",
|
||||
strings.NewReader(clickHouseServerHelloQuery),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("new request: %v", err)
|
||||
}
|
||||
|
||||
res, err := stripper.RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatalf("round trip: %v", err)
|
||||
}
|
||||
if res != nil && res.Body != nil {
|
||||
res.Body.Close()
|
||||
}
|
||||
if strings.Contains(seenBody, "displayName()") {
|
||||
t.Fatalf("expected displayName() rewritten out of server hello query, got %q", seenBody)
|
||||
}
|
||||
if seenBody != clickHouseServerHelloCompatQuery {
|
||||
t.Fatalf("expected compatibility server hello query, got %q", seenBody)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClickHouseHTTPCompatibilityStripperRewritesServerHelloOnlyOnce(t *testing.T) {
|
||||
var seenBodies []string
|
||||
stripper := clickHouseHTTPClientProtocolVersionStripper{
|
||||
serverHelloRewritten: &atomic.Bool{},
|
||||
next: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.Body != nil {
|
||||
data, err := io.ReadAll(req.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
seenBodies = append(seenBodies, string(data))
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
req, err := http.NewRequest(
|
||||
http.MethodPost,
|
||||
"http://clickhouse.local:8123/?database=default",
|
||||
strings.NewReader(clickHouseServerHelloQuery),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("new request %d: %v", i, err)
|
||||
}
|
||||
res, err := stripper.RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatalf("round trip %d: %v", i, err)
|
||||
}
|
||||
if res != nil && res.Body != nil {
|
||||
res.Body.Close()
|
||||
}
|
||||
}
|
||||
|
||||
if len(seenBodies) != 2 {
|
||||
t.Fatalf("expected two forwarded requests, got %d", len(seenBodies))
|
||||
}
|
||||
if seenBodies[0] != clickHouseServerHelloCompatQuery {
|
||||
t.Fatalf("expected first (handshake) request rewritten, got %q", seenBodies[0])
|
||||
}
|
||||
if seenBodies[1] != clickHouseServerHelloQuery {
|
||||
t.Fatalf("expected second identical query left unchanged after handshake, got %q", seenBodies[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestClickHouseHTTPCompatibilityStripperLeavesOtherBodiesUnchanged(t *testing.T) {
|
||||
const userQuery = "SELECT count() FROM system.tables"
|
||||
var seenBody string
|
||||
stripper := clickHouseHTTPClientProtocolVersionStripper{
|
||||
next: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.Body != nil {
|
||||
data, err := io.ReadAll(req.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
seenBody = string(data)
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
req, err := http.NewRequest(
|
||||
http.MethodPost,
|
||||
"http://clickhouse.local:8123/?database=default",
|
||||
strings.NewReader(userQuery),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("new request: %v", err)
|
||||
}
|
||||
|
||||
res, err := stripper.RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatalf("round trip: %v", err)
|
||||
}
|
||||
if res != nil && res.Body != nil {
|
||||
res.Body.Close()
|
||||
}
|
||||
if seenBody != userQuery {
|
||||
t.Fatalf("expected user query body untouched, got %q", seenBody)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithClickHouseProtocolForcesProtocolSelection(t *testing.T) {
|
||||
httpConfig := withClickHouseProtocol(connection.ConnectionConfig{
|
||||
Type: "clickhouse",
|
||||
|
||||
@@ -34,10 +34,13 @@ func (c *CustomDB) Connect(config connection.ConnectionConfig) error {
|
||||
if err != nil {
|
||||
return formatCustomDriverOpenError(driver, err)
|
||||
}
|
||||
configureSQLConnectionPool(db, driver)
|
||||
c.conn = db
|
||||
c.driver = driver
|
||||
c.pingTimeout = getConnectTimeout(config)
|
||||
if err := c.Ping(); err != nil {
|
||||
_ = db.Close()
|
||||
c.conn = nil
|
||||
return wrapDatabaseConnectionVerifyError(err)
|
||||
}
|
||||
return nil
|
||||
@@ -115,6 +118,24 @@ func (c *CustomDB) Query(query string) ([]map[string]interface{}, []string, erro
|
||||
return scanRowsForDialect(rows, c.scanDialect())
|
||||
}
|
||||
|
||||
func (c *CustomDB) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if c.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
|
||||
rows, err := c.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return streamRowsForDialect(rows, c.scanDialect(), consumer)
|
||||
}
|
||||
|
||||
func (c *CustomDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return c.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (c *CustomDB) scanDialect() string {
|
||||
if strings.EqualFold(strings.TrimSpace(c.driver), "mysql") {
|
||||
return "mysql"
|
||||
|
||||
@@ -110,6 +110,7 @@ func (d *DamengDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("第%d次连接打开失败: %v", idx+1, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "dameng")
|
||||
d.conn = db
|
||||
d.pingTimeout = getConnectTimeout(attempt)
|
||||
if err := d.Ping(); err != nil {
|
||||
@@ -182,6 +183,24 @@ func (d *DamengDB) Query(query string) ([]map[string]interface{}, []string, erro
|
||||
return scanRows(rows)
|
||||
}
|
||||
|
||||
func (d *DamengDB) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if d.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
|
||||
rows, err := d.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return streamRows(rows, consumer)
|
||||
}
|
||||
|
||||
func (d *DamengDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return d.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (d *DamengDB) ExecContext(ctx context.Context, query string) (int64, error) {
|
||||
if d.conn == nil {
|
||||
return 0, fmt.Errorf("连接未打开")
|
||||
|
||||
@@ -76,6 +76,48 @@ type StatementQueryExecer interface {
|
||||
QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error)
|
||||
}
|
||||
|
||||
// QueryStreamConsumer receives query metadata and rows incrementally.
|
||||
// Implementations can stream rows directly to files to avoid buffering entire result sets in memory.
|
||||
type QueryStreamConsumer interface {
|
||||
SetColumns(columns []string) error
|
||||
ConsumeRow(row map[string]interface{}) error
|
||||
}
|
||||
|
||||
// QueryStreamValueConsumer is an optional fast path for stream consumers that
|
||||
// can consume normalized row values in column order without requiring a
|
||||
// map[string]interface{} allocation per row.
|
||||
type QueryStreamValueConsumer interface {
|
||||
SetColumns(columns []string) error
|
||||
ConsumeRowValues(values []interface{}) error
|
||||
}
|
||||
|
||||
// StreamQueryExecer is an optional interface for drivers or pinned sessions that can
|
||||
// stream query rows incrementally instead of materializing []map rows in memory.
|
||||
type StreamQueryExecer interface {
|
||||
StreamQuery(query string, consumer QueryStreamConsumer) error
|
||||
StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error
|
||||
}
|
||||
|
||||
// ExplainExecer is an optional interface for drivers that can run EXPLAIN and
|
||||
// return the dialect-native output (JSON text, table rows as JSON, or XML).
|
||||
//
|
||||
// Drivers that implement this interface own the full EXPLAIN lifecycle:
|
||||
// - MySQL: prefer EXPLAIN FORMAT=JSON, fallback to vanilla EXPLAIN on 5.7
|
||||
// - PostgreSQL: EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON)
|
||||
// - Oracle: EXPLAIN PLAN SET STATEMENT_ID ... + DBMS_XPLAN.DISPLAY + cleanup
|
||||
// - SQLServer: SET SHOWPLAN_XML ON + sql + SET OFF (defer cleanup mandatory)
|
||||
// - SQLite: EXPLAIN QUERY PLAN
|
||||
// - ClickHouse: EXPLAIN JSON
|
||||
//
|
||||
// The driver decides which format to use and returns the raw payload plus the
|
||||
// detected format tag; the app layer parses via the corresponding parser.
|
||||
//
|
||||
// Drivers that do NOT implement this interface fall back to the generic path
|
||||
// in app.DiagnoseQuery: wrap the SQL as "EXPLAIN <sql>" and run via QueryMulti.
|
||||
type ExplainExecer interface {
|
||||
Explain(ctx context.Context, query string) (raw string, format connection.ExplainFormat, err error)
|
||||
}
|
||||
|
||||
// StatementQueryMessageExecer can run queries on a pinned session and return
|
||||
// extra server messages/notices alongside rows.
|
||||
type StatementQueryMessageExecer interface {
|
||||
@@ -196,6 +238,22 @@ func (e *sqlConnStatementExecer) Query(query string) ([]map[string]interface{},
|
||||
return e.QueryContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (e *sqlConnStatementExecer) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if e == nil || e.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
rows, err := e.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
return streamRowsForDialect(rows, e.scanDialect, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlConnStatementExecer) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return e.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlConnStatementExecer) QueryMultiContext(ctx context.Context, query string) ([]connection.ResultSetData, error) {
|
||||
if e == nil || e.conn == nil {
|
||||
return nil, localizedDatabaseRuntimeError("db.backend.error.connection_not_open", nil)
|
||||
@@ -293,6 +351,23 @@ func (e *sqlConnTransactionExecer) Query(query string) ([]map[string]interface{}
|
||||
return e.QueryContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (e *sqlConnTransactionExecer) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
conn, err := e.activeConn()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rows, err := conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
return streamRowsForDialect(rows, e.scanDialect, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlConnTransactionExecer) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return e.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlConnTransactionExecer) QueryMultiContext(ctx context.Context, query string) ([]connection.ResultSetData, error) {
|
||||
conn, err := e.activeConn()
|
||||
if err != nil {
|
||||
@@ -419,6 +494,23 @@ func (e *sqlTxStatementExecer) Query(query string) ([]map[string]interface{}, []
|
||||
return e.QueryContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (e *sqlTxStatementExecer) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
tx, err := e.activeTx()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rows, err := tx.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
return streamRows(rows, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlTxStatementExecer) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return e.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlTxStatementExecer) QueryMultiContext(ctx context.Context, query string) ([]connection.ResultSetData, error) {
|
||||
tx, err := e.activeTx()
|
||||
if err != nil {
|
||||
|
||||
@@ -23,4 +23,5 @@ func registerOptionalDatabaseFactories() {
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("iotdb"), "iotdb", "apache-iotdb", "apache_iotdb")
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("clickhouse"), "clickhouse")
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("elasticsearch"), "elasticsearch", "elastic")
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("trino"), "trino")
|
||||
}
|
||||
|
||||
@@ -23,4 +23,5 @@ func registerOptionalDatabaseFactories() {
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("iotdb"), "iotdb", "apache-iotdb", "apache_iotdb")
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("clickhouse"), "clickhouse")
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("elasticsearch"), "elasticsearch", "elastic")
|
||||
registerDatabaseFactory(newOptionalDriverAgentDatabase("trino"), "trino")
|
||||
}
|
||||
|
||||
@@ -187,6 +187,7 @@ func (d *DirosDB) Connect(config connection.ConnectionConfig) error {
|
||||
errorDetails = append(errorDetails, fmt.Sprintf("%s 打开失败: %v", address, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "diros")
|
||||
|
||||
timeout := getConnectTimeout(candidateConfig)
|
||||
ctx, cancel := utils.ContextWithTimeout(timeout)
|
||||
|
||||
@@ -4,25 +4,26 @@ package db
|
||||
|
||||
func init() {
|
||||
optionalDriverAgentRevisions = map[string]string{
|
||||
"mariadb": "src-0a4176f4b5743323",
|
||||
"oceanbase": "src-7cb0f2c4dc0510a5",
|
||||
"diros": "src-cc11b882e28fa5d4",
|
||||
"starrocks": "src-83a6d81c91c7f5c8",
|
||||
"sphinx": "src-a70c2cd4d223dac2",
|
||||
"sqlserver": "src-6d5cf334034bce41",
|
||||
"sqlite": "src-762863d48f653b89",
|
||||
"duckdb": "src-df5d60ebb175bbbc",
|
||||
"dameng": "src-596bebeaa016fc74",
|
||||
"kingbase": "src-2e5a1337b0405c57",
|
||||
"highgo": "src-5a29a1d3685eb6b4",
|
||||
"vastbase": "src-e3cfef65512feb23",
|
||||
"opengauss": "src-58227ba3bc1ec894",
|
||||
"gaussdb": "src-1458564993a9d455",
|
||||
"iris": "src-1b072c57af08bec4",
|
||||
"mongodb": "src-57fdd8bfebdcd46e",
|
||||
"tdengine": "src-939715f94df1ec9c",
|
||||
"iotdb": "src-473c39891f926db2",
|
||||
"clickhouse": "src-482d62ed565b3e69",
|
||||
"elasticsearch": "src-2fb00b94d7067c56",
|
||||
"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",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -49,6 +49,7 @@ var optionalGoDrivers = map[string]struct{}{
|
||||
"iotdb": {},
|
||||
"clickhouse": {},
|
||||
"elasticsearch": {},
|
||||
"trino": {},
|
||||
}
|
||||
|
||||
// optionalDriverAgentRevisions 记录 GoNavi 对各可选 driver-agent 包装逻辑的兼容版本。
|
||||
@@ -151,6 +152,8 @@ func driverDisplayName(driverType string) string {
|
||||
return "ClickHouse"
|
||||
case "elasticsearch":
|
||||
return "Elasticsearch"
|
||||
case "trino":
|
||||
return "Trino"
|
||||
case "chroma":
|
||||
return "Chroma"
|
||||
case "qdrant":
|
||||
|
||||
@@ -179,6 +179,7 @@ func (g *GaussDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("%s 数据库=%s 打开连接失败: %v", sslLabel, dbName, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(dbConn, "gaussdb")
|
||||
g.conn = dbConn
|
||||
|
||||
if err := g.Ping(); err != nil {
|
||||
@@ -233,6 +234,7 @@ func (g *GaussDB) ensureSearchPath(baseDSN string) {
|
||||
|
||||
newDB, err := sql.Open("gaussdb", newDSN)
|
||||
if err == nil {
|
||||
configureSQLConnectionPool(newDB, "gaussdb")
|
||||
newDB.SetConnMaxLifetime(5 * time.Minute)
|
||||
oldConn := g.conn
|
||||
g.conn = newDB
|
||||
|
||||
@@ -103,6 +103,7 @@ func (h *HighGoDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("第%d次连接打开失败: %v", idx+1, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "highgo")
|
||||
h.conn = db
|
||||
h.pingTimeout = getConnectTimeout(attempt)
|
||||
if err := h.Ping(); err != nil {
|
||||
|
||||
@@ -141,9 +141,12 @@ func (i *IrisDB) Connect(config connection.ConnectionConfig) error {
|
||||
if err != nil {
|
||||
return wrapDatabaseConnectionOpenError(err)
|
||||
}
|
||||
configureSQLConnectionPool(db, "iris")
|
||||
i.conn = db
|
||||
i.pingTimeout = getConnectTimeout(runConfig)
|
||||
if err := i.Ping(); err != nil {
|
||||
_ = db.Close()
|
||||
i.conn = nil
|
||||
return wrapDatabaseConnectionVerifyError(err)
|
||||
}
|
||||
cleanupOnFailure = false
|
||||
|
||||
@@ -157,6 +157,7 @@ func (k *KingbaseDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("第%d次连接打开失败: %v", idx+1, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "kingbase")
|
||||
k.conn = db
|
||||
k.pingTimeout = getConnectTimeout(attempt)
|
||||
if err := k.Ping(); err != nil {
|
||||
@@ -175,8 +176,9 @@ func (k *KingbaseDB) Connect(config connection.ConnectionConfig) error {
|
||||
// 将 search_path 参数拼入 DSN
|
||||
finalDSN := dsn + " search_path=" + quoteConnValue(searchPathStr)
|
||||
if finalDB, err := sql.Open("kingbase", finalDSN); err == nil {
|
||||
k.pingTimeout = getConnectTimeout(attempt)
|
||||
configureSQLConnectionPool(finalDB, "kingbase")
|
||||
finalDB.SetConnMaxLifetime(5 * time.Minute)
|
||||
k.pingTimeout = getConnectTimeout(attempt)
|
||||
|
||||
// 临时将 k.conn 指向 finalDB 来做 ping 测试
|
||||
oldConn := k.conn
|
||||
|
||||
@@ -49,10 +49,13 @@ func (m *MariaDB) Connect(config connection.ConnectionConfig) error {
|
||||
if err != nil {
|
||||
return wrapDatabaseConnectionOpenError(err)
|
||||
}
|
||||
configureSQLConnectionPool(db, "mariadb")
|
||||
m.conn = db
|
||||
m.pingTimeout = getConnectTimeout(config)
|
||||
|
||||
if err := m.Ping(); err != nil {
|
||||
_ = db.Close()
|
||||
m.conn = nil
|
||||
return wrapDatabaseConnectionVerifyError(err)
|
||||
}
|
||||
return nil
|
||||
|
||||
117
internal/db/memory_limit_autoscale.go
Normal file
117
internal/db/memory_limit_autoscale.go
Normal file
@@ -0,0 +1,117 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"runtime/debug"
|
||||
"sync/atomic"
|
||||
|
||||
"GoNavi-Wails/internal/logger"
|
||||
)
|
||||
|
||||
// 本文件实现 driver-agent 进程的 GOMEMLIMIT 自适应策略。
|
||||
//
|
||||
// 背景:driver-agent 是独立子进程,主进程无法控制其内存。
|
||||
// 静态 limit(如固定 2GB)在大结果集场景下会触发 GC 硬模式,导出速度降 15-25%;
|
||||
// 而 limit 设太大又失去约束意义。
|
||||
//
|
||||
// 策略:起步保守(2GB),运行时监控 HeapAlloc,逼近当前 limit 时按 1GB 步长抬升,
|
||||
// 上限 8GB 防止无限制增长。配合 SetGCPercent(50) + 周期 GC,正常场景下稳态堆
|
||||
// 仅几百 MB,limit 不会被触发;只有 GC 真的跟不上时才抬升。
|
||||
|
||||
const (
|
||||
// MemorySoftLimitInitialBytes 是进程启动时的默认 soft limit。
|
||||
// 2GB 覆盖绝大多数导出场景的稳态堆需求。
|
||||
MemorySoftLimitInitialBytes int64 = 2 * 1024 * 1024 * 1024
|
||||
|
||||
// MemorySoftLimitMaxBytes 是自适应抬升的绝对上限。
|
||||
// 8GB 防止失控;用户机器内存 < 16GB 时也留有余量给主进程和系统。
|
||||
MemorySoftLimitMaxBytes int64 = 8 * 1024 * 1024 * 1024
|
||||
|
||||
// MemorySoftLimitStepBytes 是每次抬升的步长。
|
||||
// 1GB 粒度足够平滑(不会一次跳太多),又不会频繁触发(HeapAlloc 1GB 量级才需要再抬)。
|
||||
MemorySoftLimitStepBytes int64 = 1 * 1024 * 1024 * 1024
|
||||
|
||||
// MemoryAutoscaleTriggerPercent 控制 HeapAlloc 达到当前 limit 的多少百分比时触发抬升。
|
||||
// 80% 留出 20% 缓冲,避免 GC 噪声导致频繁抖动抬升。
|
||||
MemoryAutoscaleTriggerPercent = 80
|
||||
)
|
||||
|
||||
// currentMemorySoftLimit 记录当前已应用的 soft limit。
|
||||
// atomic 以便 MaybeGrowMemoryLimit 在并发流式查询中安全调用。
|
||||
var currentMemorySoftLimit atomic.Int64
|
||||
|
||||
// InitMemorySoftLimit 在进程启动时调用,应用初始 soft limit。
|
||||
// 重复调用安全:以最后一次为准。
|
||||
func InitMemorySoftLimit(initial int64) {
|
||||
if initial <= 0 {
|
||||
initial = MemorySoftLimitInitialBytes
|
||||
}
|
||||
if initial > MemorySoftLimitMaxBytes {
|
||||
initial = MemorySoftLimitMaxBytes
|
||||
}
|
||||
debug.SetMemoryLimit(initial)
|
||||
currentMemorySoftLimit.Store(initial)
|
||||
}
|
||||
|
||||
// CurrentMemorySoftLimit 返回当前已应用的 soft limit,主要供测试和监控使用。
|
||||
func CurrentMemorySoftLimit() int64 {
|
||||
return currentMemorySoftLimit.Load()
|
||||
}
|
||||
|
||||
// MaybeGrowMemoryLimit 在大结果集流式处理时周期性调用(建议与周期 GC 同节奏),
|
||||
// 当堆用量达到当前 limit 的 MemoryAutoscaleTriggerPercent 时按步长抬升。
|
||||
//
|
||||
// 设计要点:
|
||||
// - 仅对调用过 InitMemorySoftLimit 的进程生效(driver-agent);主进程未初始化时 currentMemorySoftLimit=0,
|
||||
// 本函数直接返回,不影响主进程的 GC 行为
|
||||
// - 读 HeapAlloc 用 runtime.ReadMemStats(短暂 STW,每 5W 行一次可忽略)
|
||||
// - 抬升通过 debug.SetMemoryLimit 应用,原子记录新值
|
||||
// - 达到 MemorySoftLimitMaxBytes 后不再抬升,让 GC 硬模式接管
|
||||
// - 不做"降级":进程 long-running,下次任务可能同样需要;soft limit 大不浪费内存
|
||||
//
|
||||
// 返回 true 表示触发了抬升(用于日志观测)。
|
||||
func MaybeGrowMemoryLimit() bool {
|
||||
current := currentMemorySoftLimit.Load()
|
||||
if current <= 0 {
|
||||
// 进程未启用 soft limit(如主进程),跳过自适应
|
||||
return false
|
||||
}
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, readHeapAlloc())
|
||||
if !grown {
|
||||
return false
|
||||
}
|
||||
|
||||
currentHeap := readHeapAlloc()
|
||||
debug.SetMemoryLimit(next)
|
||||
currentMemorySoftLimit.Store(next)
|
||||
logger.Infof("内存 soft limit 自适应抬升:%dMB → %dMB(HeapAlloc=%dMB)",
|
||||
current/1024/1024, next/1024/1024, currentHeap/1024/1024)
|
||||
return true
|
||||
}
|
||||
|
||||
// shouldGrowMemoryLimit 是 MaybeGrowMemoryLimit 的纯逻辑核心,便于单元测试。
|
||||
// 输入:当前 limit、当前 HeapAlloc;输出:是否抬升、抬升后的新 limit。
|
||||
func shouldGrowMemoryLimit(currentLimit, heapAlloc int64) (bool, int64) {
|
||||
if currentLimit >= MemorySoftLimitMaxBytes {
|
||||
return false, currentLimit
|
||||
}
|
||||
if heapAlloc < currentLimit*MemoryAutoscaleTriggerPercent/100 {
|
||||
return false, currentLimit
|
||||
}
|
||||
next := currentLimit + MemorySoftLimitStepBytes
|
||||
if next > MemorySoftLimitMaxBytes {
|
||||
next = MemorySoftLimitMaxBytes
|
||||
}
|
||||
if next == currentLimit {
|
||||
return false, currentLimit
|
||||
}
|
||||
return true, next
|
||||
}
|
||||
|
||||
// readHeapAlloc 封装 runtime.ReadMemStats,便于测试 mock。
|
||||
func readHeapAlloc() int64 {
|
||||
var m runtime.MemStats
|
||||
runtime.ReadMemStats(&m)
|
||||
return int64(m.HeapAlloc)
|
||||
}
|
||||
129
internal/db/memory_limit_autoscale_test.go
Normal file
129
internal/db/memory_limit_autoscale_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestShouldGrowMemoryLimit_NoActionWhenBelowThreshold(t *testing.T) {
|
||||
current := int64(2 * 1024 * 1024 * 1024) // 2GB
|
||||
// HeapAlloc 仅占 50%,远低于 80% 阈值
|
||||
heapAlloc := current * 50 / 100
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, heapAlloc)
|
||||
if grown {
|
||||
t.Fatalf("HeapAlloc=%dB 低于 80%% 阈值,不应抬升", heapAlloc)
|
||||
}
|
||||
if next != current {
|
||||
t.Fatalf("未抬升时 next 应等于 current,got=%d want=%d", next, current)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldGrowMemoryLimit_NoActionAtExactThreshold(t *testing.T) {
|
||||
current := int64(2 * 1024 * 1024 * 1024)
|
||||
// HeapAlloc 正好等于 80% 阈值:heapAlloc < current*80/100 为假时才抬升
|
||||
// current*80/100 = 1.6GB;heapAlloc = 1.6GB 时 heapAlloc < 1.6GB 为假 → 抬升
|
||||
heapAlloc := current * MemoryAutoscaleTriggerPercent / 100
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, heapAlloc)
|
||||
if !grown {
|
||||
t.Fatalf("HeapAlloc=%dB 已达 80%% 阈值,应抬升", heapAlloc)
|
||||
}
|
||||
wantNext := current + MemorySoftLimitStepBytes
|
||||
if next != wantNext {
|
||||
t.Fatalf("抬升步长错误:got=%d want=%d", next, wantNext)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldGrowMemoryLimit_StepByGB(t *testing.T) {
|
||||
current := int64(2 * 1024 * 1024 * 1024) // 2GB
|
||||
heapAlloc := int64(3 * 1024 * 1024 * 1024) // 3GB > 2GB * 80% = 1.6GB
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, heapAlloc)
|
||||
if !grown {
|
||||
t.Fatalf("HeapAlloc=%dB 超过 80%% 阈值,应抬升", heapAlloc)
|
||||
}
|
||||
wantNext := int64(3 * 1024 * 1024 * 1024) // 2GB + 1GB step = 3GB
|
||||
if next != wantNext {
|
||||
t.Fatalf("抬升后 limit 应为 3GB,got=%d want=%d", next, wantNext)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldGrowMemoryLimit_CapAtMax(t *testing.T) {
|
||||
// 当前 limit 已等于上限
|
||||
current := MemorySoftLimitMaxBytes
|
||||
heapAlloc := current * 2 // 即使 HeapAlloc 远超 limit 也不再抬升
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, heapAlloc)
|
||||
if grown {
|
||||
t.Fatalf("已达上限 %dB,不应再抬升", MemorySoftLimitMaxBytes)
|
||||
}
|
||||
if next != current {
|
||||
t.Fatalf("已达上限时 next 应等于 current,got=%d want=%d", next, current)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldGrowMemoryLimit_CapWhenStepExceedsMax(t *testing.T) {
|
||||
// 当前 limit 距上限不足 1GB 步长:7.5GB
|
||||
current := MemorySoftLimitMaxBytes - 512*1024*1024 // 7.5GB
|
||||
heapAlloc := current + 1 // 超过 80% 阈值
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, heapAlloc)
|
||||
if !grown {
|
||||
t.Fatalf("HeapAlloc 已逼近 limit,应触发抬升(即便步长会触及上限)")
|
||||
}
|
||||
if next != MemorySoftLimitMaxBytes {
|
||||
t.Fatalf("抬升后应 cap 在 max,got=%d want=%d", next, MemorySoftLimitMaxBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldGrowMemoryLimit_NoActionWhenCurrentExceedsMax(t *testing.T) {
|
||||
// 异常情况:current > max(理论不会发生,但应防御性处理)
|
||||
current := MemorySoftLimitMaxBytes + 1
|
||||
heapAlloc := current * 2
|
||||
|
||||
grown, next := shouldGrowMemoryLimit(current, heapAlloc)
|
||||
if grown {
|
||||
t.Fatalf("current 已超过 max,不应再抬升")
|
||||
}
|
||||
if next != current {
|
||||
t.Fatalf("next 应等于 current,got=%d want=%d", next, current)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInitMemorySoftLimit_ClampToMax(t *testing.T) {
|
||||
// 初始化值超过 max 时应被截断
|
||||
overMax := MemorySoftLimitMaxBytes * 2
|
||||
InitMemorySoftLimit(overMax)
|
||||
if got := CurrentMemorySoftLimit(); got != MemorySoftLimitMaxBytes {
|
||||
t.Fatalf("初始化超过 max 应被截断:got=%d want=%d", got, MemorySoftLimitMaxBytes)
|
||||
}
|
||||
// 恢复默认值,避免污染其他测试
|
||||
InitMemorySoftLimit(MemorySoftLimitInitialBytes)
|
||||
}
|
||||
|
||||
func TestInitMemorySoftLimit_DefaultWhenZeroOrNegative(t *testing.T) {
|
||||
InitMemorySoftLimit(0)
|
||||
if got := CurrentMemorySoftLimit(); got != MemorySoftLimitInitialBytes {
|
||||
t.Fatalf("initial=0 应使用默认值:got=%d want=%d", got, MemorySoftLimitInitialBytes)
|
||||
}
|
||||
InitMemorySoftLimit(-1)
|
||||
if got := CurrentMemorySoftLimit(); got != MemorySoftLimitInitialBytes {
|
||||
t.Fatalf("initial<0 应使用默认值:got=%d want=%d", got, MemorySoftLimitInitialBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaybeGrowMemoryLimit_NoOpWhenUninitialized(t *testing.T) {
|
||||
// 模拟主进程未初始化的场景:
|
||||
// 通过将 currentMemorySoftLimit 直接置零(绕过 InitMemorySoftLimit)来测试
|
||||
// 注意:这是一个破坏性测试,需在测试末尾恢复状态
|
||||
saved := currentMemorySoftLimit.Load()
|
||||
defer currentMemorySoftLimit.Store(saved)
|
||||
|
||||
currentMemorySoftLimit.Store(0)
|
||||
if MaybeGrowMemoryLimit() {
|
||||
t.Fatalf("currentMemorySoftLimit=0 时应直接返回 false,不主动初始化")
|
||||
}
|
||||
if got := CurrentMemorySoftLimit(); got != 0 {
|
||||
t.Fatalf("未初始化时不应被 MaybeGrowMemoryLimit 改写,got=%d want=0", got)
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,21 @@ func parseMySQLDriverCharsetsForTest(t *testing.T, dsn string) []string {
|
||||
return charsets
|
||||
}
|
||||
|
||||
func TestConfigureSQLConnectionPoolCapsOpenConnections(t *testing.T) {
|
||||
dbConn, err := sql.Open("mysql", "root@tcp(127.0.0.1:1)/test")
|
||||
if err != nil {
|
||||
t.Fatalf("sql.Open failed: %v", err)
|
||||
}
|
||||
defer dbConn.Close()
|
||||
|
||||
configureSQLConnectionPool(dbConn, "mysql")
|
||||
|
||||
stats := dbConn.Stats()
|
||||
if stats.MaxOpenConnections != defaultSQLMaxOpenConns {
|
||||
t.Fatalf("expected max open connections %d, got %d", defaultSQLMaxOpenConns, stats.MaxOpenConnections)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMySQLDSN_MergesConnectionParamsWithDefaults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -847,6 +847,7 @@ func (m *MySQLDB) Connect(config connection.ConnectionConfig) error {
|
||||
}
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, candidateConfig.Type)
|
||||
|
||||
timeout := getConnectTimeout(candidateConfig)
|
||||
ctx, cancel := utils.ContextWithTimeout(timeout)
|
||||
|
||||
@@ -621,6 +621,7 @@ func (o *OceanBaseDB) connectOracleViaOBClient(config connection.ConnectionConfi
|
||||
errorDetails = append(errorDetails, fmt.Sprintf("%s 打开失败:%v", address, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "oceanbase")
|
||||
|
||||
timeout := getConnectTimeout(candidateConfig)
|
||||
ctx, cancel := utils.ContextWithTimeout(timeout)
|
||||
@@ -741,6 +742,7 @@ func (o *OceanBaseDB) Connect(config connection.ConnectionConfig) error {
|
||||
errorDetails = append(errorDetails, fmt.Sprintf("%s 打开失败:%v", address, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "oceanbase")
|
||||
|
||||
timeout := getConnectTimeout(candidateConfig)
|
||||
ctx, cancel := utils.ContextWithTimeout(timeout)
|
||||
|
||||
@@ -2,6 +2,7 @@ package db
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -28,6 +29,7 @@ const (
|
||||
optionalAgentMethodOpenSession = "openSession"
|
||||
optionalAgentMethodCloseSession = "closeSession"
|
||||
optionalAgentMethodQuery = "query"
|
||||
optionalAgentMethodStreamQuery = "streamQuery"
|
||||
optionalAgentMethodExec = "exec"
|
||||
optionalAgentMethodGetDatabases = "getDatabases"
|
||||
optionalAgentMethodGetTables = "getTables"
|
||||
@@ -40,6 +42,19 @@ const (
|
||||
optionalAgentMethodApplyChanges = "applyChanges"
|
||||
optionalAgentDefaultScannerMaxBytes = 8 << 20
|
||||
optionalAgentMetadataProbeTimeout = 5 * time.Second
|
||||
// callStreamQueryGCInterval 控制 callStreamQuery 每接收多少行 driver-agent 数据触发一次 runtime.GC。
|
||||
//
|
||||
// 该路径不走 sql.Rows(scan_rows.go 的周期 GC 覆盖不到),但每个 chunk 解码
|
||||
// [][]interface{} + normalizeQueryValue 转换会产生大量临时字符串,需要主动回收。
|
||||
// 取 50000 与 scan_rows.go 的 streamRowsPeriodicGCInterval 保持一致,
|
||||
// 让两端在相近节奏下分别 GC,避免内存峰值叠加。
|
||||
callStreamQueryGCInterval = 50000
|
||||
)
|
||||
|
||||
const (
|
||||
optionalAgentChunkColumns = "columns"
|
||||
optionalAgentChunkRows = "rows"
|
||||
optionalAgentChunkDone = "done"
|
||||
)
|
||||
|
||||
type optionalAgentRequest struct {
|
||||
@@ -60,6 +75,7 @@ type optionalAgentResponse struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
Data json.RawMessage `json:"data,omitempty"`
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
ChunkType string `json:"chunkType,omitempty"`
|
||||
RowsAffected int64 `json:"rowsAffected,omitempty"`
|
||||
}
|
||||
|
||||
@@ -269,6 +285,125 @@ func (c *optionalDriverAgentClient) callWithTimeout(req optionalAgentRequest, ou
|
||||
}
|
||||
}
|
||||
|
||||
func (c *optionalDriverAgentClient) callStreamQuery(req optionalAgentRequest, consumer QueryStreamConsumer) error {
|
||||
if consumer == nil {
|
||||
return fmt.Errorf("query stream consumer required")
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.nextID++
|
||||
req.ID = c.nextID
|
||||
|
||||
payload, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload = append(payload, '\n')
|
||||
if _, err := c.stdin.Write(payload); err != nil {
|
||||
stderrText := c.stderrText()
|
||||
if stderrText == "" {
|
||||
return fmt.Errorf("调用 %s 驱动代理失败:%w", driverDisplayName(c.driver), err)
|
||||
}
|
||||
return fmt.Errorf("调用 %s 驱动代理失败:%w(stderr: %s)", driverDisplayName(c.driver), err, stderrText)
|
||||
}
|
||||
|
||||
var columns []string
|
||||
valueConsumer, useValueConsumer := consumer.(QueryStreamValueConsumer)
|
||||
|
||||
// processedRows 用于周期性触发 GC。
|
||||
// 该路径不走 sql.Rows,scan_rows.go 的周期 GC 覆盖不到。
|
||||
// 每个 chunk 解码会分配 [][]interface{} + normalizeQueryValue 转换副本,
|
||||
// 88W 行场景下不主动 GC 会让主进程 RSS 单调爬升。
|
||||
var processedRows int64
|
||||
|
||||
for {
|
||||
line, err := c.reader.ReadBytes('\n')
|
||||
if err != nil {
|
||||
stderrText := c.stderrText()
|
||||
if stderrText == "" {
|
||||
return fmt.Errorf("读取 %s 驱动代理响应失败:%w", driverDisplayName(c.driver), err)
|
||||
}
|
||||
return fmt.Errorf("读取 %s 驱动代理响应失败:%w(stderr: %s)", driverDisplayName(c.driver), err, stderrText)
|
||||
}
|
||||
|
||||
var resp optionalAgentResponse
|
||||
if err := json.Unmarshal(line, &resp); err != nil {
|
||||
return fmt.Errorf("解析 %s 驱动代理响应失败:%w", driverDisplayName(c.driver), err)
|
||||
}
|
||||
if !resp.Success {
|
||||
errText := strings.TrimSpace(resp.Error)
|
||||
if errText == "" {
|
||||
errText = fmt.Sprintf("%s 驱动代理返回失败", driverDisplayName(c.driver))
|
||||
}
|
||||
return errors.New(errText)
|
||||
}
|
||||
|
||||
switch resp.ChunkType {
|
||||
case optionalAgentChunkColumns:
|
||||
columns = append(columns[:0], resp.Fields...)
|
||||
if err := consumer.SetColumns(columns); err != nil {
|
||||
return err
|
||||
}
|
||||
case optionalAgentChunkRows:
|
||||
if len(columns) == 0 {
|
||||
return fmt.Errorf("%s 驱动代理流式响应缺少列信息", driverDisplayName(c.driver))
|
||||
}
|
||||
rows, err := decodeOptionalAgentRowValueBatch(resp.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析 %s 驱动代理流式数据失败:%w", driverDisplayName(c.driver), err)
|
||||
}
|
||||
for _, row := range rows {
|
||||
if useValueConsumer {
|
||||
if err := valueConsumer.ConsumeRowValues(row); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
entry := make(map[string]interface{}, len(columns))
|
||||
for i, column := range columns {
|
||||
if i < len(row) {
|
||||
entry[column] = row[i]
|
||||
} else {
|
||||
entry[column] = nil
|
||||
}
|
||||
}
|
||||
if err := consumer.ConsumeRow(entry); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
processedRows += int64(len(rows))
|
||||
if processedRows >= callStreamQueryGCInterval {
|
||||
runtime.GC()
|
||||
processedRows = 0
|
||||
}
|
||||
case optionalAgentChunkDone:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("%s 驱动代理返回未知流式分片类型:%s", driverDisplayName(c.driver), strings.TrimSpace(resp.ChunkType))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeOptionalAgentRowValueBatch(data []byte) ([][]interface{}, error) {
|
||||
if len(data) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
decoder := json.NewDecoder(bytes.NewReader(data))
|
||||
decoder.UseNumber()
|
||||
var rows [][]interface{}
|
||||
if err := decoder.Decode(&rows); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for rowIdx := range rows {
|
||||
for colIdx := range rows[rowIdx] {
|
||||
rows[rowIdx][colIdx] = normalizeQueryValue(rows[rowIdx][colIdx])
|
||||
}
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
func (c *optionalDriverAgentClient) forceTerminate() {
|
||||
if c.stdin != nil {
|
||||
_ = c.stdin.Close()
|
||||
@@ -397,6 +532,42 @@ func (d *OptionalDriverAgentDB) Query(query string) ([]map[string]interface{}, [
|
||||
return data, fields, nil
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return d.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = client.callStreamQuery(optionalAgentRequest{
|
||||
Method: optionalAgentMethodStreamQuery,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, consumer)
|
||||
if isOptionalAgentStreamUnsupportedError(err) {
|
||||
logger.Warnf("%s 驱动代理暂不支持流式查询,回退到缓冲模式:err=%v", driverDisplayName(d.driverType), err)
|
||||
data, columns, queryErr := d.QueryContext(ctx, query)
|
||||
if queryErr != nil {
|
||||
return queryErr
|
||||
}
|
||||
if err := consumer.SetColumns(columns); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range data {
|
||||
if err := consumer.ConsumeRow(row); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) ExecContext(ctx context.Context, query string) (int64, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return 0, err
|
||||
@@ -458,6 +629,39 @@ func (s *optionalDriverAgentSession) Query(query string) ([]map[string]interface
|
||||
return s.QueryContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return s.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if err := s.ensureOpen(); err != nil {
|
||||
return err
|
||||
}
|
||||
err := s.client.callStreamQuery(optionalAgentRequest{
|
||||
Method: optionalAgentMethodStreamQuery,
|
||||
SessionID: s.sessionID,
|
||||
Query: query,
|
||||
TimeoutMs: timeoutMsFromContext(ctx),
|
||||
}, consumer)
|
||||
if isOptionalAgentStreamUnsupportedError(err) {
|
||||
logger.Warnf("%s 驱动代理事务会话暂不支持流式查询,回退到缓冲模式:err=%v", driverDisplayName(s.driver), err)
|
||||
data, columns, queryErr := s.QueryContext(ctx, query)
|
||||
if queryErr != nil {
|
||||
return queryErr
|
||||
}
|
||||
if err := consumer.SetColumns(columns); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range data {
|
||||
if err := consumer.ConsumeRow(row); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *optionalDriverAgentSession) QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error) {
|
||||
if err := s.ensureOpen(); err != nil {
|
||||
return nil, nil, err
|
||||
@@ -525,6 +729,17 @@ func (s *optionalDriverAgentSession) ensureOpen() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func isOptionalAgentStreamUnsupportedError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
text := strings.TrimSpace(err.Error())
|
||||
if text == "" {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(text, "不支持的方法") || strings.Contains(text, "不支持流式查询")
|
||||
}
|
||||
|
||||
func (d *OptionalDriverAgentDB) GetDatabases() ([]string, error) {
|
||||
client, err := d.requireClient()
|
||||
if err != nil {
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"GoNavi-Wails/internal/connection"
|
||||
@@ -65,3 +68,71 @@ func TestNormalizeKingbaseAgentChangeSetByColumns(t *testing.T) {
|
||||
t.Fatalf("unexpected update value key \"event name\" after normalization")
|
||||
}
|
||||
}
|
||||
|
||||
type optionalAgentTestWriteCloser struct {
|
||||
bytes.Buffer
|
||||
}
|
||||
|
||||
func (w *optionalAgentTestWriteCloser) Close() error { return nil }
|
||||
|
||||
type optionalAgentTestStreamConsumer struct {
|
||||
columns []string
|
||||
rows [][]interface{}
|
||||
}
|
||||
|
||||
func (c *optionalAgentTestStreamConsumer) SetColumns(columns []string) error {
|
||||
c.columns = append([]string(nil), columns...)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *optionalAgentTestStreamConsumer) ConsumeRow(row map[string]interface{}) error {
|
||||
values := make([]interface{}, len(c.columns))
|
||||
for idx, column := range c.columns {
|
||||
values[idx] = row[column]
|
||||
}
|
||||
c.rows = append(c.rows, values)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *optionalAgentTestStreamConsumer) ConsumeRowValues(values []interface{}) error {
|
||||
c.rows = append(c.rows, append([]interface{}(nil), values...))
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestOptionalDriverAgentClientCallStreamQueryConsumesChunks(t *testing.T) {
|
||||
var stdin optionalAgentTestWriteCloser
|
||||
stdout := strings.Join([]string{
|
||||
`{"id":1,"success":true,"chunkType":"columns","fields":["id","name"]}`,
|
||||
`{"id":1,"success":true,"chunkType":"rows","data":[[1,"alice"],[2,"bob"]]}`,
|
||||
`{"id":1,"success":true,"chunkType":"done"}`,
|
||||
}, "\n") + "\n"
|
||||
|
||||
client := &optionalDriverAgentClient{
|
||||
stdin: &stdin,
|
||||
reader: bufio.NewReader(strings.NewReader(stdout)),
|
||||
driver: "oceanbase",
|
||||
}
|
||||
consumer := &optionalAgentTestStreamConsumer{}
|
||||
if err := client.callStreamQuery(optionalAgentRequest{
|
||||
Method: optionalAgentMethodStreamQuery,
|
||||
Query: "SELECT 1",
|
||||
}, consumer); err != nil {
|
||||
t.Fatalf("callStreamQuery 返回错误: %v", err)
|
||||
}
|
||||
|
||||
if len(consumer.columns) != 2 || consumer.columns[0] != "id" || consumer.columns[1] != "name" {
|
||||
t.Fatalf("流式列定义异常: %#v", consumer.columns)
|
||||
}
|
||||
if len(consumer.rows) != 2 {
|
||||
t.Fatalf("流式行数异常: %#v", consumer.rows)
|
||||
}
|
||||
if got := consumer.rows[0][1]; got != "alice" {
|
||||
t.Fatalf("第 1 行数据异常,want=%q got=%v", "alice", got)
|
||||
}
|
||||
if got := consumer.rows[1][0]; got != int64(2) {
|
||||
t.Fatalf("第 2 行 ID 异常,want=%d got=%v (%T)", 2, got, got)
|
||||
}
|
||||
if !strings.Contains(stdin.String(), `"method":"streamQuery"`) {
|
||||
t.Fatalf("请求未使用 streamQuery 方法: %s", stdin.String())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -165,6 +165,7 @@ func (o *OracleDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("第%d次连接打开失败: %v", idx+1, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "oracle")
|
||||
o.conn = db
|
||||
o.pingTimeout = getConnectTimeout(attempt)
|
||||
if err := o.Ping(); err != nil {
|
||||
|
||||
@@ -159,6 +159,7 @@ func (p *PostgresDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("%s 数据库=%s 打开连接失败: %v", sslLabel, dbName, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(dbConn, "postgres")
|
||||
p.conn = dbConn
|
||||
|
||||
// Force verification
|
||||
@@ -604,6 +605,7 @@ func (p *PostgresDB) ensureSearchPath(baseDSN string) {
|
||||
|
||||
newDB, err := sql.Open("postgres", newDSN)
|
||||
if err == nil {
|
||||
configureSQLConnectionPool(newDB, "postgres")
|
||||
newDB.SetConnMaxLifetime(5 * time.Minute)
|
||||
oldConn := p.conn
|
||||
p.conn = newDB
|
||||
|
||||
@@ -3,14 +3,38 @@ package db
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"runtime"
|
||||
|
||||
"GoNavi-Wails/internal/connection"
|
||||
)
|
||||
|
||||
// streamRowsPeriodicGCInterval 控制 streamRowsForDialect 每处理多少行主动触发一次 runtime.GC。
|
||||
//
|
||||
// 背景:大结果集(88W+ 行)流式扫描时,每行 scanner 会分配 []interface{} 和 map[string]interface{},
|
||||
// Go 默认 GOGC=100 下堆翻倍才触发 GC,瞬时峰值可达数据总量 5-8 倍。
|
||||
// 这里周期性主动 GC,让内存在扫描过程中及时回收,避免 RSS 单调爬升。
|
||||
//
|
||||
// 取值 50000:每 5W 行触发一次 GC,对 88W 行导出场景约触发 18 次,CPU 开销可忽略;
|
||||
// 同时保证单次 GC 之间累积的临时对象不超过几百 MB,避免 GC 间隙堆膨胀。
|
||||
const streamRowsPeriodicGCInterval = 50000
|
||||
|
||||
func scanRows(rows *sql.Rows) ([]map[string]interface{}, []string, error) {
|
||||
return scanRowsForDialect(rows, "")
|
||||
}
|
||||
|
||||
func streamRows(rows *sql.Rows, consumer QueryStreamConsumer) error {
|
||||
return streamRowsForDialect(rows, "", consumer)
|
||||
}
|
||||
|
||||
type queryRowScanner struct {
|
||||
columns []string
|
||||
dbTypeNames []string
|
||||
dialect string
|
||||
values []interface{}
|
||||
normalized []interface{}
|
||||
valuePtrs []interface{}
|
||||
}
|
||||
|
||||
func scanRowsForDialect(rows *sql.Rows, dialect string) ([]map[string]interface{}, []string, error) {
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
@@ -23,27 +47,14 @@ func scanRowsForDialect(rows *sql.Rows, dialect string) ([]map[string]interface{
|
||||
colTypes = nil
|
||||
}
|
||||
|
||||
scanner := newQueryRowScanner(columns, colTypes, dialect)
|
||||
resultData := make([]map[string]interface{}, 0)
|
||||
|
||||
for rows.Next() {
|
||||
values := make([]interface{}, len(columns))
|
||||
valuePtrs := make([]interface{}, len(columns))
|
||||
for i := range columns {
|
||||
valuePtrs[i] = &values[i]
|
||||
}
|
||||
|
||||
if err := rows.Scan(valuePtrs...); err != nil {
|
||||
entry, err := scanner.scanCurrentRow(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
entry := make(map[string]interface{}, len(columns))
|
||||
for i, col := range columns {
|
||||
dbTypeName := ""
|
||||
if colTypes != nil && i < len(colTypes) && colTypes[i] != nil {
|
||||
dbTypeName = colTypes[i].DatabaseTypeName()
|
||||
}
|
||||
entry[col] = normalizeQueryValueWithDBTypeAndDialect(values[i], dbTypeName, dialect)
|
||||
}
|
||||
resultData = append(resultData, entry)
|
||||
}
|
||||
|
||||
@@ -53,6 +64,108 @@ func scanRowsForDialect(rows *sql.Rows, dialect string) ([]map[string]interface{
|
||||
return resultData, columns, nil
|
||||
}
|
||||
|
||||
func streamRowsForDialect(rows *sql.Rows, dialect string, consumer QueryStreamConsumer) error {
|
||||
if consumer == nil {
|
||||
return fmt.Errorf("query stream consumer required")
|
||||
}
|
||||
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
columns = ensureUniqueQueryColumnNames(columns)
|
||||
|
||||
colTypes, err := rows.ColumnTypes()
|
||||
if err != nil || len(colTypes) != len(columns) {
|
||||
colTypes = nil
|
||||
}
|
||||
|
||||
scanner := newQueryRowScanner(columns, colTypes, dialect)
|
||||
if err := consumer.SetColumns(columns); err != nil {
|
||||
return err
|
||||
}
|
||||
valueConsumer, useValueConsumer := consumer.(QueryStreamValueConsumer)
|
||||
|
||||
// processedRows 用于周期性触发 GC,见 streamRowsPeriodicGCInterval 注释。
|
||||
// 注意:此路径同时被 driver-agent 进程(OceanBase 等 optional driver)和
|
||||
// 主进程的 in-process 流式查询调用,所以一处加 GC 即可覆盖两端。
|
||||
var processedRows int64
|
||||
|
||||
for rows.Next() {
|
||||
if useValueConsumer {
|
||||
values, err := scanner.scanCurrentRowValues(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if err := valueConsumer.ConsumeRowValues(values); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
entry, err := scanner.scanCurrentRow(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if err := consumer.ConsumeRow(entry); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
processedRows++
|
||||
if processedRows%streamRowsPeriodicGCInterval == 0 {
|
||||
runtime.GC()
|
||||
// 自适应抬升 driver-agent 进程的内存 soft limit。
|
||||
// 主进程未启用 soft limit(未调 InitMemorySoftLimit),此调用是 no-op。
|
||||
MaybeGrowMemoryLimit()
|
||||
}
|
||||
}
|
||||
|
||||
return rows.Err()
|
||||
}
|
||||
|
||||
func newQueryRowScanner(columns []string, colTypes []*sql.ColumnType, dialect string) *queryRowScanner {
|
||||
values := make([]interface{}, len(columns))
|
||||
valuePtrs := make([]interface{}, len(columns))
|
||||
for i := range columns {
|
||||
valuePtrs[i] = &values[i]
|
||||
}
|
||||
dbTypeNames := make([]string, len(columns))
|
||||
for i := range columns {
|
||||
if colTypes != nil && i < len(colTypes) && colTypes[i] != nil {
|
||||
dbTypeNames[i] = colTypes[i].DatabaseTypeName()
|
||||
}
|
||||
}
|
||||
return &queryRowScanner{
|
||||
columns: columns,
|
||||
dbTypeNames: dbTypeNames,
|
||||
dialect: dialect,
|
||||
values: values,
|
||||
normalized: make([]interface{}, len(columns)),
|
||||
valuePtrs: valuePtrs,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *queryRowScanner) scanCurrentRowValues(rows *sql.Rows) ([]interface{}, error) {
|
||||
if err := rows.Scan(s.valuePtrs...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range s.columns {
|
||||
s.normalized[i] = normalizeQueryValueWithDBTypeAndDialect(s.values[i], s.dbTypeNames[i], s.dialect)
|
||||
}
|
||||
return s.normalized, nil
|
||||
}
|
||||
|
||||
func (s *queryRowScanner) scanCurrentRow(rows *sql.Rows) (map[string]interface{}, error) {
|
||||
normalized, err := s.scanCurrentRowValues(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entry := make(map[string]interface{}, len(s.columns))
|
||||
for i, col := range s.columns {
|
||||
entry[col] = normalized[i]
|
||||
}
|
||||
return entry, nil
|
||||
}
|
||||
|
||||
func ensureUniqueQueryColumnNames(columns []string) []string {
|
||||
if len(columns) == 0 {
|
||||
return columns
|
||||
|
||||
27
internal/db/sql_pool.go
Normal file
27
internal/db/sql_pool.go
Normal file
@@ -0,0 +1,27 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultSQLMaxOpenConns = 4
|
||||
defaultSQLConnMaxLifetime = 30 * time.Minute
|
||||
defaultSQLConnMaxIdleTime = 30 * time.Second
|
||||
)
|
||||
|
||||
func configureSQLConnectionPool(db *sql.DB, dbType string) {
|
||||
if db == nil {
|
||||
return
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(dbType)) {
|
||||
case "sqlite", "duckdb":
|
||||
return
|
||||
}
|
||||
db.SetMaxOpenConns(defaultSQLMaxOpenConns)
|
||||
db.SetMaxIdleConns(0)
|
||||
db.SetConnMaxIdleTime(defaultSQLConnMaxIdleTime)
|
||||
db.SetConnMaxLifetime(defaultSQLConnMaxLifetime)
|
||||
}
|
||||
@@ -176,10 +176,13 @@ func (s *SqlServerDB) Connect(config connection.ConnectionConfig) error {
|
||||
if err != nil {
|
||||
return wrapDatabaseConnectionOpenError(err)
|
||||
}
|
||||
configureSQLConnectionPool(db, "sqlserver")
|
||||
s.conn = db
|
||||
s.pingTimeout = getConnectTimeout(config)
|
||||
|
||||
if err := s.Ping(); err != nil {
|
||||
_ = db.Close()
|
||||
s.conn = nil
|
||||
return wrapDatabaseConnectionVerifyError(err)
|
||||
}
|
||||
return nil
|
||||
@@ -385,6 +388,23 @@ func (e *sqlServerSessionExecer) QueryContext(ctx context.Context, query string)
|
||||
return rows, columns, err
|
||||
}
|
||||
|
||||
func (e *sqlServerSessionExecer) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if e == nil || e.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
retmsg := &sqlexp.ReturnMessage{}
|
||||
rows, err := e.conn.QueryContext(ctx, query, retmsg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
return streamRows(rows, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlServerSessionExecer) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return e.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (e *sqlServerSessionExecer) QueryWithMessages(query string) ([]map[string]interface{}, []string, []string, error) {
|
||||
return e.QueryContextWithMessages(context.Background(), query)
|
||||
}
|
||||
|
||||
@@ -270,6 +270,7 @@ func (s *StarRocksDB) Connect(config connection.ConnectionConfig) error {
|
||||
errorDetails = append(errorDetails, fmt.Sprintf("%s 打开失败: %v", address, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "starrocks")
|
||||
|
||||
timeout := getConnectTimeout(candidateConfig)
|
||||
ctx, cancel := utils.ContextWithTimeout(timeout)
|
||||
|
||||
@@ -96,6 +96,7 @@ func (t *TDengineDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("第%d次连接打开失败: %v", idx+1, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "tdengine")
|
||||
t.conn = db
|
||||
t.pingTimeout = getConnectTimeout(attempt)
|
||||
|
||||
@@ -168,6 +169,24 @@ func (t *TDengineDB) Query(query string) ([]map[string]interface{}, []string, er
|
||||
return scanRows(rows)
|
||||
}
|
||||
|
||||
func (t *TDengineDB) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if t.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
|
||||
rows, err := t.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return streamRows(rows, consumer)
|
||||
}
|
||||
|
||||
func (t *TDengineDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return t.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (t *TDengineDB) ExecContext(ctx context.Context, query string) (int64, error) {
|
||||
if t.conn == nil {
|
||||
return 0, fmt.Errorf("连接未打开")
|
||||
|
||||
661
internal/db/trino_impl.go
Normal file
661
internal/db/trino_impl.go
Normal file
@@ -0,0 +1,661 @@
|
||||
//go:build gonavi_full_drivers || gonavi_trino_driver
|
||||
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"GoNavi-Wails/internal/connection"
|
||||
"GoNavi-Wails/internal/logger"
|
||||
"GoNavi-Wails/internal/ssh"
|
||||
|
||||
trinodriver "github.com/trinodb/trino-go-client/trino"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultTrinoPort = 8080
|
||||
defaultTrinoSource = "GoNavi"
|
||||
)
|
||||
|
||||
type TrinoDB struct {
|
||||
conn *sql.DB
|
||||
pingTimeout time.Duration
|
||||
forwarder *ssh.LocalForwarder
|
||||
namespace string
|
||||
customClientName string
|
||||
}
|
||||
|
||||
func normalizeTrinoConfig(config connection.ConnectionConfig) connection.ConnectionConfig {
|
||||
normalized := applyTrinoURI(config)
|
||||
normalized = applyTrinoHostURI(normalized)
|
||||
if strings.TrimSpace(normalized.Host) == "" {
|
||||
normalized.Host = "localhost"
|
||||
}
|
||||
if normalized.Port <= 0 {
|
||||
normalized.Port = defaultTrinoPort
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
func applyTrinoURI(config connection.ConnectionConfig) connection.ConnectionConfig {
|
||||
return applyTrinoEndpointURI(config, config.URI, false)
|
||||
}
|
||||
|
||||
func applyTrinoHostURI(config connection.ConnectionConfig) connection.ConnectionConfig {
|
||||
return applyTrinoEndpointURI(config, config.Host, true)
|
||||
}
|
||||
|
||||
func applyTrinoEndpointURI(config connection.ConnectionConfig, raw string, fromHostField bool) connection.ConnectionConfig {
|
||||
uriText := strings.TrimSpace(raw)
|
||||
if uriText == "" {
|
||||
return config
|
||||
}
|
||||
parsed, err := url.Parse(uriText)
|
||||
if err != nil {
|
||||
return config
|
||||
}
|
||||
scheme := strings.ToLower(strings.TrimSpace(parsed.Scheme))
|
||||
if scheme != "trino" && scheme != "http" && scheme != "https" {
|
||||
return config
|
||||
}
|
||||
if strings.TrimSpace(parsed.Host) == "" {
|
||||
return config
|
||||
}
|
||||
|
||||
if parsed.User != nil {
|
||||
if strings.TrimSpace(config.User) == "" {
|
||||
config.User = parsed.User.Username()
|
||||
}
|
||||
if pass, ok := parsed.User.Password(); ok && config.Password == "" {
|
||||
config.Password = pass
|
||||
}
|
||||
}
|
||||
|
||||
params := url.Values{}
|
||||
mergeConnectionParamValues(params, parsed.Query())
|
||||
mergeConnectionParamValues(params, connectionParamsFromText(config.ConnectionParams))
|
||||
|
||||
catalog := strings.TrimSpace(parsed.Query().Get("catalog"))
|
||||
schema := strings.TrimSpace(parsed.Query().Get("schema"))
|
||||
if strings.TrimSpace(config.Database) == "" {
|
||||
config.Database = joinTrinoNamespace(catalog, schema)
|
||||
}
|
||||
|
||||
if scheme == "https" {
|
||||
config.UseSSL = true
|
||||
if normalizeSSLModeValue(config.SSLMode) == sslModeDisable || strings.TrimSpace(config.SSLMode) == "" {
|
||||
config.SSLMode = sslModeRequired
|
||||
}
|
||||
}
|
||||
|
||||
defaultPort := config.Port
|
||||
if defaultPort <= 0 {
|
||||
defaultPort = defaultTrinoPort
|
||||
}
|
||||
if fromHostField || strings.TrimSpace(config.Host) == "" {
|
||||
host, port, ok := parseHostPortWithDefault(parsed.Host, defaultPort)
|
||||
if ok {
|
||||
config.Host = host
|
||||
config.Port = port
|
||||
}
|
||||
}
|
||||
if config.Port <= 0 {
|
||||
config.Port = defaultPort
|
||||
}
|
||||
config.ConnectionParams = params.Encode()
|
||||
return config
|
||||
}
|
||||
|
||||
func splitTrinoNamespace(raw string) (string, string) {
|
||||
text := strings.TrimSpace(raw)
|
||||
if text == "" {
|
||||
return "", ""
|
||||
}
|
||||
parts := strings.SplitN(text, ".", 2)
|
||||
catalog := strings.TrimSpace(parts[0])
|
||||
if len(parts) == 1 {
|
||||
return catalog, ""
|
||||
}
|
||||
return catalog, strings.TrimSpace(parts[1])
|
||||
}
|
||||
|
||||
func joinTrinoNamespace(catalog, schema string) string {
|
||||
c := strings.TrimSpace(catalog)
|
||||
s := strings.TrimSpace(schema)
|
||||
switch {
|
||||
case c == "":
|
||||
return s
|
||||
case s == "":
|
||||
return c
|
||||
default:
|
||||
return c + "." + s
|
||||
}
|
||||
}
|
||||
|
||||
func resolveTrinoNamespace(raw string, fallback string) (string, string) {
|
||||
catalog, schema := splitTrinoNamespace(raw)
|
||||
if catalog != "" || schema != "" {
|
||||
return catalog, schema
|
||||
}
|
||||
return splitTrinoNamespace(fallback)
|
||||
}
|
||||
|
||||
func quoteTrinoIdentifier(ident string) string {
|
||||
return `"` + strings.ReplaceAll(strings.TrimSpace(ident), `"`, `""`) + `"`
|
||||
}
|
||||
|
||||
func quoteTrinoQualifiedTable(catalog, schema, table string) string {
|
||||
quoted := make([]string, 0, 3)
|
||||
if trimmed := strings.TrimSpace(catalog); trimmed != "" {
|
||||
quoted = append(quoted, quoteTrinoIdentifier(trimmed))
|
||||
}
|
||||
if trimmed := strings.TrimSpace(schema); trimmed != "" {
|
||||
quoted = append(quoted, quoteTrinoIdentifier(trimmed))
|
||||
}
|
||||
quoted = append(quoted, quoteTrinoIdentifier(table))
|
||||
return strings.Join(quoted, ".")
|
||||
}
|
||||
|
||||
func escapeTrinoSQLLiteral(value string) string {
|
||||
return "'" + strings.ReplaceAll(strings.TrimSpace(value), "'", "''") + "'"
|
||||
}
|
||||
|
||||
func trinoRowValue(row map[string]interface{}, keys ...string) (interface{}, bool) {
|
||||
if len(row) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
for _, key := range keys {
|
||||
for current, value := range row {
|
||||
if strings.EqualFold(strings.TrimSpace(current), strings.TrimSpace(key)) {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func trinoRowString(row map[string]interface{}, keys ...string) string {
|
||||
value, ok := trinoRowValue(row, keys...)
|
||||
if !ok || value == nil {
|
||||
return ""
|
||||
}
|
||||
text := strings.TrimSpace(fmt.Sprintf("%v", value))
|
||||
if strings.EqualFold(text, "<nil>") {
|
||||
return ""
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func firstTrinoMapValueAsString(row map[string]interface{}) string {
|
||||
for _, value := range row {
|
||||
text := strings.TrimSpace(fmt.Sprintf("%v", value))
|
||||
if !strings.EqualFold(text, "<nil>") {
|
||||
return text
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func firstTrinoRowValueAsString(data []map[string]interface{}) string {
|
||||
if len(data) == 0 {
|
||||
return ""
|
||||
}
|
||||
return firstTrinoMapValueAsString(data[0])
|
||||
}
|
||||
|
||||
func (t *TrinoDB) buildTrinoHTTPClient(config connection.ConnectionConfig) (*http.Client, error) {
|
||||
transport := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: (&net.Dialer{
|
||||
Timeout: getConnectTimeout(config),
|
||||
KeepAlive: 30 * time.Second,
|
||||
}).DialContext,
|
||||
MaxIdleConns: 32,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: time.Second,
|
||||
}
|
||||
tlsConfig, err := resolveGenericTLSConfig(config)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if tlsConfig != nil {
|
||||
transport.TLSClientConfig = tlsConfig
|
||||
}
|
||||
return &http.Client{Transport: transport}, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) registerTrinoCustomClient(config connection.ConnectionConfig) (string, error) {
|
||||
client, err := t.buildTrinoHTTPClient(config)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
name := fmt.Sprintf("gonavi-trino-%d", time.Now().UnixNano())
|
||||
if err := trinodriver.RegisterCustomClient(name, client); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
|
||||
func buildTrinoDSN(config connection.ConnectionConfig, customClientName string) (string, error) {
|
||||
user := strings.TrimSpace(config.User)
|
||||
if user == "" {
|
||||
return "", fmt.Errorf("Trino 用户名不能为空")
|
||||
}
|
||||
|
||||
scheme := "http"
|
||||
if config.UseSSL {
|
||||
scheme = "https"
|
||||
}
|
||||
if config.Password != "" && scheme != "https" {
|
||||
return "", fmt.Errorf("Trino 启用密码认证时必须使用 HTTPS")
|
||||
}
|
||||
|
||||
params := connectionParamsFromText(config.ConnectionParams)
|
||||
catalog, schema := resolveTrinoNamespace(config.Database, "")
|
||||
if catalog != "" {
|
||||
params.Set("catalog", catalog)
|
||||
}
|
||||
if schema != "" {
|
||||
params.Set("schema", schema)
|
||||
}
|
||||
if strings.TrimSpace(params.Get("source")) == "" {
|
||||
params.Set("source", defaultTrinoSource)
|
||||
}
|
||||
if strings.TrimSpace(params.Get("explicitPrepare")) == "" {
|
||||
params.Set("explicitPrepare", "false")
|
||||
}
|
||||
if strings.TrimSpace(params.Get("query_timeout")) == "" {
|
||||
params.Set("query_timeout", fmt.Sprintf("%ds", getConnectTimeoutSeconds(config)))
|
||||
}
|
||||
if strings.TrimSpace(customClientName) != "" {
|
||||
params.Set("custom_client", strings.TrimSpace(customClientName))
|
||||
}
|
||||
|
||||
endpoint := &url.URL{
|
||||
Scheme: scheme,
|
||||
Host: net.JoinHostPort(strings.TrimSpace(config.Host), strconv.Itoa(config.Port)),
|
||||
RawQuery: params.Encode(),
|
||||
}
|
||||
if config.Password != "" {
|
||||
endpoint.User = url.UserPassword(user, config.Password)
|
||||
} else {
|
||||
endpoint.User = url.User(user)
|
||||
}
|
||||
return endpoint.String(), nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) Close() error {
|
||||
if t.conn != nil {
|
||||
if err := t.conn.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
t.conn = nil
|
||||
}
|
||||
if t.forwarder != nil {
|
||||
if err := t.forwarder.Close(); err != nil {
|
||||
logger.Warnf("关闭 Trino SSH 端口转发失败:%v", err)
|
||||
}
|
||||
t.forwarder = nil
|
||||
}
|
||||
if t.customClientName != "" {
|
||||
trinodriver.DeregisterCustomClient(t.customClientName)
|
||||
t.customClientName = ""
|
||||
}
|
||||
t.namespace = ""
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) Connect(config connection.ConnectionConfig) error {
|
||||
_ = t.Close()
|
||||
|
||||
runConfig := normalizeTrinoConfig(config)
|
||||
t.pingTimeout = getConnectTimeout(runConfig)
|
||||
|
||||
if runConfig.UseSSH {
|
||||
forwarder, err := ssh.GetOrCreateLocalForwarder(runConfig.SSH, runConfig.Host, runConfig.Port)
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建 SSH 隧道失败:%w", err)
|
||||
}
|
||||
t.forwarder = forwarder
|
||||
|
||||
host, portText, err := net.SplitHostPort(forwarder.LocalAddr)
|
||||
if err != nil {
|
||||
_ = t.Close()
|
||||
return fmt.Errorf("解析本地转发地址失败:%w", err)
|
||||
}
|
||||
port, err := strconv.Atoi(portText)
|
||||
if err != nil {
|
||||
_ = t.Close()
|
||||
return fmt.Errorf("解析本地端口失败:%w", err)
|
||||
}
|
||||
runConfig.Host = host
|
||||
runConfig.Port = port
|
||||
runConfig.UseSSH = false
|
||||
logger.Infof("Trino 通过本地端口转发连接:%s -> %s:%d", forwarder.LocalAddr, config.Host, config.Port)
|
||||
}
|
||||
|
||||
customClientName, err := t.registerTrinoCustomClient(runConfig)
|
||||
if err != nil {
|
||||
_ = t.Close()
|
||||
return fmt.Errorf("注册 Trino 自定义 HTTP 客户端失败:%w", err)
|
||||
}
|
||||
t.customClientName = customClientName
|
||||
|
||||
dsn, err := buildTrinoDSN(runConfig, customClientName)
|
||||
if err != nil {
|
||||
_ = t.Close()
|
||||
return err
|
||||
}
|
||||
conn, err := sql.Open("trino", dsn)
|
||||
if err != nil {
|
||||
_ = t.Close()
|
||||
return err
|
||||
}
|
||||
t.conn = conn
|
||||
t.namespace = strings.TrimSpace(runConfig.Database)
|
||||
if err := t.Ping(); err != nil {
|
||||
_ = t.Close()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) Ping() error {
|
||||
if t.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
timeout := t.pingTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = 10 * time.Second
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
rows, err := t.conn.QueryContext(ctx, "SELECT 1")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
if !rows.Next() {
|
||||
if err := rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return fmt.Errorf("连接查询验证未返回结果")
|
||||
}
|
||||
var value sql.NullInt64
|
||||
if err := rows.Scan(&value); err != nil {
|
||||
return err
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
|
||||
func (t *TrinoDB) QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error) {
|
||||
if t.conn == nil {
|
||||
return nil, nil, fmt.Errorf("连接未打开")
|
||||
}
|
||||
rows, err := t.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanRows(rows)
|
||||
}
|
||||
|
||||
func (t *TrinoDB) Query(query string) ([]map[string]interface{}, []string, error) {
|
||||
return t.QueryContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (t *TrinoDB) StreamQueryContext(ctx context.Context, query string, consumer QueryStreamConsumer) error {
|
||||
if t.conn == nil {
|
||||
return fmt.Errorf("连接未打开")
|
||||
}
|
||||
rows, err := t.conn.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
return streamRows(rows, consumer)
|
||||
}
|
||||
|
||||
func (t *TrinoDB) StreamQuery(query string, consumer QueryStreamConsumer) error {
|
||||
return t.StreamQueryContext(context.Background(), query, consumer)
|
||||
}
|
||||
|
||||
func (t *TrinoDB) ExecContext(ctx context.Context, query string) (int64, error) {
|
||||
if t.conn == nil {
|
||||
return 0, fmt.Errorf("连接未打开")
|
||||
}
|
||||
res, err := t.conn.ExecContext(ctx, query)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
affected, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return 0, nil
|
||||
}
|
||||
return affected, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) Exec(query string) (int64, error) {
|
||||
return t.ExecContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (t *TrinoDB) queryTrinoSingleColumnStrings(query string) ([]string, error) {
|
||||
data, _, err := t.Query(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]string, 0, len(data))
|
||||
for _, row := range data {
|
||||
text := firstTrinoMapValueAsString(row)
|
||||
if text != "" {
|
||||
result = append(result, text)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetDatabases() ([]string, error) {
|
||||
catalogs, err := t.queryTrinoSingleColumnStrings("SHOW CATALOGS")
|
||||
if err != nil {
|
||||
if strings.TrimSpace(t.namespace) != "" {
|
||||
return []string{t.namespace}, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
namespaces := make([]string, 0, len(catalogs)*2)
|
||||
seen := make(map[string]struct{}, len(catalogs)*4)
|
||||
var lastErr error
|
||||
for _, catalog := range catalogs {
|
||||
query := fmt.Sprintf("SHOW SCHEMAS FROM %s", quoteTrinoIdentifier(catalog))
|
||||
schemas, schemaErr := t.queryTrinoSingleColumnStrings(query)
|
||||
if schemaErr != nil {
|
||||
lastErr = schemaErr
|
||||
continue
|
||||
}
|
||||
for _, schema := range schemas {
|
||||
namespace := joinTrinoNamespace(catalog, schema)
|
||||
if namespace == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(namespace)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
namespaces = append(namespaces, namespace)
|
||||
}
|
||||
}
|
||||
|
||||
if len(namespaces) == 0 {
|
||||
if strings.TrimSpace(t.namespace) != "" {
|
||||
return []string{t.namespace}, nil
|
||||
}
|
||||
if lastErr != nil {
|
||||
return nil, lastErr
|
||||
}
|
||||
}
|
||||
sort.Strings(namespaces)
|
||||
return namespaces, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetTables(dbName string) ([]string, error) {
|
||||
catalog, schema := resolveTrinoNamespace(dbName, t.namespace)
|
||||
if catalog == "" || schema == "" {
|
||||
return nil, fmt.Errorf("Trino 默认命名空间必须使用 catalog.schema")
|
||||
}
|
||||
query := fmt.Sprintf("SHOW TABLES FROM %s.%s", quoteTrinoIdentifier(catalog), quoteTrinoIdentifier(schema))
|
||||
tables, err := t.queryTrinoSingleColumnStrings(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sort.Strings(tables)
|
||||
return tables, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetCreateStatement(dbName, tableName string) (string, error) {
|
||||
catalog, schema := resolveTrinoNamespace(dbName, t.namespace)
|
||||
if catalog == "" || schema == "" {
|
||||
return "", fmt.Errorf("Trino 默认命名空间必须使用 catalog.schema")
|
||||
}
|
||||
query := fmt.Sprintf("SHOW CREATE TABLE %s", quoteTrinoQualifiedTable(catalog, schema, tableName))
|
||||
data, _, err := t.Query(query)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
ddl := firstTrinoRowValueAsString(data)
|
||||
if ddl == "" {
|
||||
return "", fmt.Errorf("未返回建表语句")
|
||||
}
|
||||
return ddl, nil
|
||||
}
|
||||
|
||||
func buildTrinoColumnsQuery(catalog, schema, tableName string) string {
|
||||
return fmt.Sprintf(`SELECT
|
||||
column_name,
|
||||
data_type,
|
||||
is_nullable,
|
||||
column_default
|
||||
FROM %s.information_schema.columns
|
||||
WHERE table_schema = %s AND table_name = %s
|
||||
ORDER BY ordinal_position`,
|
||||
quoteTrinoIdentifier(catalog),
|
||||
escapeTrinoSQLLiteral(schema),
|
||||
escapeTrinoSQLLiteral(tableName),
|
||||
)
|
||||
}
|
||||
|
||||
func buildTrinoColumnDefinitions(data []map[string]interface{}) []connection.ColumnDefinition {
|
||||
result := make([]connection.ColumnDefinition, 0, len(data))
|
||||
for _, row := range data {
|
||||
column := connection.ColumnDefinition{
|
||||
Name: trinoRowString(row, "column_name", "Column", "Field"),
|
||||
Type: trinoRowString(row, "data_type", "Type"),
|
||||
Nullable: strings.ToUpper(trinoRowString(row, "is_nullable", "Null")),
|
||||
}
|
||||
if rawDefault, ok := trinoRowValue(row, "column_default", "Default"); ok && rawDefault != nil {
|
||||
def := strings.TrimSpace(fmt.Sprintf("%v", rawDefault))
|
||||
if !strings.EqualFold(def, "<nil>") && def != "" {
|
||||
column.Default = &def
|
||||
}
|
||||
}
|
||||
if column.Nullable == "" {
|
||||
column.Nullable = "YES"
|
||||
}
|
||||
result = append(result, column)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetColumns(dbName, tableName string) ([]connection.ColumnDefinition, error) {
|
||||
catalog, schema := resolveTrinoNamespace(dbName, t.namespace)
|
||||
if catalog == "" || schema == "" {
|
||||
return nil, fmt.Errorf("Trino 默认命名空间必须使用 catalog.schema")
|
||||
}
|
||||
data, _, err := t.Query(buildTrinoColumnsQuery(catalog, schema, tableName))
|
||||
if err == nil {
|
||||
return buildTrinoColumnDefinitions(data), nil
|
||||
}
|
||||
|
||||
describeQuery := fmt.Sprintf("DESCRIBE %s", quoteTrinoQualifiedTable(catalog, schema, tableName))
|
||||
describeRows, _, describeErr := t.Query(describeQuery)
|
||||
if describeErr != nil {
|
||||
return nil, err
|
||||
}
|
||||
columns := make([]connection.ColumnDefinition, 0, len(describeRows))
|
||||
for _, row := range describeRows {
|
||||
name := trinoRowString(row, "Column", "column_name", "Field")
|
||||
if name == "" || strings.HasPrefix(name, "#") {
|
||||
continue
|
||||
}
|
||||
columns = append(columns, connection.ColumnDefinition{
|
||||
Name: name,
|
||||
Type: trinoRowString(row, "Type", "data_type"),
|
||||
Nullable: "YES",
|
||||
})
|
||||
}
|
||||
return columns, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetAllColumns(dbName string) ([]connection.ColumnDefinitionWithTable, error) {
|
||||
catalog, schema := resolveTrinoNamespace(dbName, t.namespace)
|
||||
if catalog == "" || schema == "" {
|
||||
return nil, fmt.Errorf("Trino 默认命名空间必须使用 catalog.schema")
|
||||
}
|
||||
query := fmt.Sprintf(`SELECT
|
||||
table_name,
|
||||
column_name,
|
||||
data_type
|
||||
FROM %s.information_schema.columns
|
||||
WHERE table_schema = %s
|
||||
ORDER BY table_name, ordinal_position`,
|
||||
quoteTrinoIdentifier(catalog),
|
||||
escapeTrinoSQLLiteral(schema),
|
||||
)
|
||||
data, _, err := t.Query(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]connection.ColumnDefinitionWithTable, 0, len(data))
|
||||
for _, row := range data {
|
||||
result = append(result, connection.ColumnDefinitionWithTable{
|
||||
TableName: trinoRowString(row, "table_name", "TABLE_NAME"),
|
||||
Name: trinoRowString(row, "column_name", "COLUMN_NAME"),
|
||||
Type: trinoRowString(row, "data_type", "DATA_TYPE"),
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetIndexes(dbName, tableName string) ([]connection.IndexDefinition, error) {
|
||||
return []connection.IndexDefinition{}, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetForeignKeys(dbName, tableName string) ([]connection.ForeignKeyDefinition, error) {
|
||||
return []connection.ForeignKeyDefinition{}, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) GetTriggers(dbName, tableName string) ([]connection.TriggerDefinition, error) {
|
||||
return []connection.TriggerDefinition{}, nil
|
||||
}
|
||||
|
||||
func (t *TrinoDB) OpenSessionExecer(ctx context.Context) (StatementExecer, error) {
|
||||
if t.conn == nil {
|
||||
return nil, fmt.Errorf("连接未打开")
|
||||
}
|
||||
conn, err := t.conn.Conn(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return NewSQLConnStatementExecer(conn), nil
|
||||
}
|
||||
@@ -94,6 +94,7 @@ func (v *VastbaseDB) Connect(config connection.ConnectionConfig) error {
|
||||
failures = append(failures, fmt.Sprintf("第%d次连接打开失败: %v", idx+1, err))
|
||||
continue
|
||||
}
|
||||
configureSQLConnectionPool(db, "vastbase")
|
||||
v.conn = db
|
||||
v.pingTimeout = getConnectTimeout(attempt)
|
||||
if err := v.Ping(); err != nil {
|
||||
|
||||
Reference in New Issue
Block a user