Files
dbx/agents/drivers/kingbase-go/main.go

2052 lines
59 KiB
Go

package main
import (
"bufio"
"context"
"database/sql"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"net/url"
"os"
"strconv"
"strings"
"sync"
"time"
"gitea.com/kingbase/gokb"
)
const (
protocolVersion = 2
defaultMaxRows = 10000
legacyAgentSessionID = "__legacy__"
maxAgentSessions = 256
defaultConnectTimeout = 15 * time.Second
)
type request struct {
ID json.RawMessage `json:"id"`
Method string `json:"method"`
Params map[string]json.RawMessage `json:"params"`
}
type response struct {
JSONRPC string `json:"jsonrpc,omitempty"`
ID json.RawMessage `json:"id,omitempty"`
Result any `json:"result,omitempty"`
Error *rpcError `json:"error,omitempty"`
}
type rpcError struct {
Code int `json:"code"`
Message string `json:"message"`
}
type connectParams struct {
Host string `json:"host"`
Port int `json:"port"`
Database string `json:"database"`
Username string `json:"username"`
Password string `json:"password"`
URLParams string `json:"url_params"`
ConnectionString string `json:"connection_string"`
MySQLCompatMode bool `json:"mysql_compat_mode"`
SSL bool `json:"ssl"`
CACertPath string `json:"ca_cert_path"`
ClientCertPath string `json:"client_cert_path"`
ClientKeyPath string `json:"client_key_path"`
}
type queryOptions struct {
SQL string `json:"sql"`
Database string `json:"database"`
Schema string `json:"schema"`
MaxRows int `json:"maxRows"`
FetchSize int `json:"fetchSize"`
TimeoutSecs int `json:"timeoutSecs"`
}
type completionAssistantRequest struct {
ConnectionID string `json:"connection_id"`
Database string `json:"database"`
Schema string `json:"schema"`
ObjectKinds []string `json:"object_kinds"`
Mask string `json:"mask"`
CaseSensitive bool `json:"case_sensitive"`
GlobalSearch bool `json:"global_search"`
MaxResults int `json:"max_results"`
ParentSchema string `json:"parent_schema"`
ParentName string `json:"parent_name"`
MatchMode string `json:"match_mode"`
}
type completionAssistantCandidate struct {
Name string `json:"name"`
Kind string `json:"kind"`
Database *string `json:"database"`
Schema *string `json:"schema"`
ParentSchema *string `json:"parent_schema"`
ParentName *string `json:"parent_name"`
Comment *string `json:"comment"`
DataType *string `json:"data_type"`
}
type completionAssistantResponse struct {
Candidates []completionAssistantCandidate `json:"candidates"`
Incomplete bool `json:"incomplete"`
FallbackUsed bool `json:"fallback_used"`
}
type queryResult struct {
Columns []string `json:"columns"`
ColumnTypes []string `json:"column_types"`
Rows [][]any `json:"rows"`
AffectedRows int64 `json:"affected_rows"`
ExecutionTimeMS int64 `json:"execution_time_ms"`
Truncated bool `json:"truncated"`
}
type queryPageResult struct {
Columns []string `json:"columns"`
ColumnTypes []string `json:"column_types"`
Rows [][]any `json:"rows"`
AffectedRows int64 `json:"affected_rows"`
ExecutionTimeMS int64 `json:"execution_time_ms"`
Truncated bool `json:"truncated"`
SessionID *string `json:"session_id"`
HasMore bool `json:"has_more"`
}
type querySession struct {
rows *sql.Rows
conn *sql.Conn
columns []string
columnTypes []string
pending []any
remaining int
cancel context.CancelFunc
}
type server struct {
db *sql.DB
openDatabase kingbaseDBOpener
params connectParams
mode kingbaseMode
usePgDefaultExpression bool
usePgViewDefinition bool
usePgFunctionDefinition bool
useLegacyRoutineDefinition bool
catalogIdentityUnsupported bool
catalogOIDUnsupported bool
infoColumnTypeUnsupported bool
infoUdtNameUnsupported bool
indexOrdinalityUnsupported bool
triggerPrettyUnsupported bool
triggerInternalUnsupported bool
constraintDefinitionUnsupported bool
constraintValidatedUnsupported bool
constraintStatusUnsupported bool
currentSchema string
schemaSet bool
sessions map[string]*querySession
nextSessionID uint64
activeCancelMu sync.Mutex
activeCancel context.CancelFunc
}
type agentSession struct {
server *server
mu sync.Mutex
}
type runtimeServer struct {
mu sync.RWMutex
sessions map[string]*agentSession
}
func main() {
runtime := &runtimeServer{sessions: map[string]*agentSession{}}
encoder := json.NewEncoder(os.Stdout)
var encoderMu sync.Mutex
var requests sync.WaitGroup
fmt.Fprintln(os.Stdout, `{"ready":true}`)
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 512*1024*1024)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" {
continue
}
var envelope request
if json.Unmarshal([]byte(line), &envelope) == nil && envelope.Method == "shutdown" {
requests.Wait()
resp, _ := runtime.handleLine(line)
encoderMu.Lock()
_ = encoder.Encode(resp)
encoderMu.Unlock()
return
}
requests.Add(1)
go func(line string) {
defer requests.Done()
resp, _ := runtime.handleLine(line)
encoderMu.Lock()
defer encoderMu.Unlock()
if err := encoder.Encode(resp); err != nil {
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
}
}(line)
}
requests.Wait()
}
func (r *runtimeServer) handleLine(line string) (response, bool) {
var req request
if err := json.Unmarshal([]byte(line), &req); err != nil {
return errorResponse(nil, err), false
}
if len(req.ID) == 0 {
req.ID = json.RawMessage("1")
}
result, shutdown, err := r.dispatch(req.Method, req.Params)
if err != nil {
return errorResponse(req.ID, err), false
}
return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown
}
func (r *runtimeServer) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) {
switch method {
case "handshake":
return map[string]any{
"protocolVersion": protocolVersion,
"agentProtocolVersion": protocolVersion,
"capabilities": []string{
"connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session",
},
}, false, nil
case "open_session":
id := stringParam(params, "agentSessionId")
if id == "" {
return nil, false, errors.New("agentSessionId is required")
}
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
return map[string]bool{"ok": true}, false, r.openSession(id, cp)
case "close_session":
return map[string]bool{"ok": true}, false, r.closeSession(stringParam(params, "agentSessionId"))
case "validate_session":
session, err := r.session(stringParam(params, "agentSessionId"))
if err != nil {
return nil, false, err
}
session.mu.Lock()
defer session.mu.Unlock()
return map[string]bool{"ok": true}, false, session.server.validateConnection()
case "cancel_session":
session, err := r.session(stringParam(params, "agentSessionId"))
if err != nil {
return nil, false, err
}
session.server.cancelActiveQuery()
return map[string]bool{"ok": true}, false, nil
case "test_connection":
return newServer().dispatch(method, params)
case "connect":
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
_ = r.closeSession(legacyAgentSessionID)
return map[string]bool{"ok": true}, false, r.openSession(legacyAgentSessionID, cp)
case "disconnect":
return map[string]bool{"ok": true}, false, r.closeSession(legacyAgentSessionID)
case "shutdown":
return map[string]bool{"ok": true}, true, r.closeAllSessions()
default:
id := stringParam(params, "agentSessionId")
if id == "" {
id = legacyAgentSessionID
}
session, err := r.session(id)
if err != nil {
return nil, false, err
}
session.mu.Lock()
defer session.mu.Unlock()
return session.server.dispatch(method, params)
}
}
func (r *runtimeServer) openSession(id string, cp connectParams) error {
r.mu.Lock()
if _, exists := r.sessions[id]; exists {
r.mu.Unlock()
return fmt.Errorf("agent session already exists: %s", id)
}
if len(r.sessions) >= maxAgentSessions {
r.mu.Unlock()
return fmt.Errorf("agent session limit reached: %d", maxAgentSessions)
}
r.mu.Unlock()
s := newServer()
if err := s.connect(cp); err != nil {
return err
}
r.mu.Lock()
defer r.mu.Unlock()
if _, exists := r.sessions[id]; exists {
_ = s.disconnect()
return fmt.Errorf("agent session already exists: %s", id)
}
r.sessions[id] = &agentSession{server: s}
return nil
}
func (r *runtimeServer) session(id string) (*agentSession, error) {
r.mu.RLock()
session := r.sessions[id]
r.mu.RUnlock()
if session == nil {
return nil, fmt.Errorf("agent session not found: %s", id)
}
return session, nil
}
func (r *runtimeServer) closeSession(id string) error {
r.mu.Lock()
session := r.sessions[id]
delete(r.sessions, id)
r.mu.Unlock()
if session == nil {
return nil
}
session.server.cancelActiveQuery()
session.mu.Lock()
defer session.mu.Unlock()
return session.server.disconnect()
}
func (r *runtimeServer) closeAllSessions() error {
r.mu.RLock()
ids := make([]string, 0, len(r.sessions))
for id := range r.sessions {
ids = append(ids, id)
}
r.mu.RUnlock()
var firstErr error
for _, id := range ids {
if err := r.closeSession(id); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
func newServer() *server {
return &server{openDatabase: openDBWithSSLMode, sessions: map[string]*querySession{}}
}
func (s *server) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) {
switch method {
case "handshake":
return map[string]any{
"protocolVersion": protocolVersion,
"agentProtocolVersion": protocolVersion,
"capabilities": []string{"connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"},
}, false, nil
case "connect":
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
return map[string]bool{"ok": true}, false, s.connect(cp)
case "test_connection":
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
return map[string]bool{"ok": true}, false, s.testConnection(cp)
case "validate_connection":
return map[string]bool{"ok": true}, false, s.validateConnection()
case "connection_info":
info, err := s.connectionInfo()
return info, false, err
case "list_databases":
result, err := s.listDatabases()
return result, false, err
case "list_schemas":
result, err := s.listSchemas(stringSliceParam(params, "visible_schemas"), boolParam(params, "show_system_schemas"))
return result, false, err
case "list_tables":
result, err := s.listTables(stringParam(params, "schema"), metadataListConstraintsFromParams(params))
return result, false, err
case "get_table_comment":
result, err := s.getTableComment(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_objects":
result, err := s.listObjects(stringParam(params, "schema"), metadataListConstraintsFromParams(params))
return result, false, err
case "list_data_types":
return kingbaseDataTypes, false, nil
case "completion_assistant_search_v1":
var request completionAssistantRequest
if err := decodeParams(params, &request); err != nil {
return nil, false, err
}
result, err := s.completionAssistantSearch(request)
return result, false, err
case "get_columns":
result, err := s.getColumns(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_indexes":
result, err := s.listIndexes(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_foreign_keys":
result, err := s.listForeignKeys(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_constraints":
result, err := s.listConstraints(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_triggers":
result, err := s.listTriggers(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "get_object_source":
result, err := s.getObjectSourceForRelation(stringParam(params, "schema"), stringParam(params, "name"), stringParam(params, "object_type"), stringParam(params, "relation_name"))
return result, false, err
case "get_type_details":
result, err := s.getTypeDetails(stringParam(params, "schema"), stringParam(params, "name"))
return result, false, err
case "get_table_ddl":
result, err := s.getTableDDL(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "get_explain_info":
result, err := s.getExplainInfo(stringParam(params, "sql"))
return map[string]any{"plan": result, "has_actual_stats": false}, false, err
case "execute_query":
var opts queryOptions
if err := decodeParams(params, &opts); err != nil {
return nil, false, err
}
result, err := s.executeQuery(opts)
return result, false, err
case "execute_query_page", "start_table_read":
var opts queryOptions
if err := decodeParams(params, &opts); err != nil {
return nil, false, err
}
result, err := s.executeQueryPage(opts, intParam(params, "pageSize"))
return result, false, err
case "fetch_query_page", "fetch_table_read_page":
result, err := s.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize"))
return result, false, err
case "close_query_session", "close_table_read_session":
return s.closeQuerySession(stringParam(params, "sessionId")), false, nil
case "execute_transaction":
result, err := s.executeTransaction(params)
return result, false, err
case "execute_batch":
result, err := s.executeBatch(params)
return result, false, err
case "disconnect":
return map[string]bool{"ok": true}, false, s.disconnect()
case "shutdown":
return map[string]bool{"ok": true}, true, s.disconnect()
default:
return nil, false, fmt.Errorf("unknown method: %s", method)
}
}
func (s *server) connect(cp connectParams) error {
_ = s.disconnect()
db, err := openAndPingDB(cp, defaultConnectTimeout, s.openDatabase)
if err != nil {
return err
}
s.db = db
s.params = cp
s.mode = detectKingbaseMode(db, cp.MySQLCompatMode)
s.mode.legacyV7 = detectKingbaseV7(db)
s.usePgDefaultExpression = false
s.usePgViewDefinition = false
s.usePgFunctionDefinition = false
s.useLegacyRoutineDefinition = false
s.catalogIdentityUnsupported = false
s.catalogOIDUnsupported = false
s.infoColumnTypeUnsupported = false
s.infoUdtNameUnsupported = false
s.indexOrdinalityUnsupported = false
s.triggerPrettyUnsupported = false
s.triggerInternalUnsupported = false
s.constraintDefinitionUnsupported = false
s.constraintValidatedUnsupported = false
s.constraintStatusUnsupported = false
return nil
}
func (s *server) testConnection(cp connectParams) error {
db, err := openAndPingDB(cp, defaultConnectTimeout, s.openDatabase)
if err != nil {
return err
}
defer db.Close()
return nil
}
type kingbaseDBOpener func(connectParams, string) (*sql.DB, error)
func openAndPingDB(cp connectParams, timeout time.Duration, opener kingbaseDBOpener) (*sql.DB, error) {
if endpoints := clusterConnectEndpoints(cp); len(endpoints) > 1 {
return openAndPingClusterEndpoints(cp, timeout, endpoints, opener)
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
return pingDBWithSSLModes(ctx, cp, opener)
}
// pingDBWithSSLModes runs the historical sslmode attempt sequence
// ("prefer" degrades to require then disable) for one set of fields.
func pingDBWithSSLModes(ctx context.Context, cp connectParams, opener kingbaseDBOpener) (*sql.DB, error) {
sslMode := effectiveSSLMode(cp)
attempts := []string{sslMode}
if sslMode == "prefer" {
attempts = []string{"require", "disable"}
}
for index, attempt := range attempts {
db, err := opener(cp, attempt)
if err == nil {
err = db.PingContext(ctx)
}
if err == nil {
return db, nil
}
if db != nil {
_ = db.Close()
}
if index == 0 && len(attempts) == 2 && shouldRetryKingbaseWithoutSSL(err) {
continue
}
return nil, err
}
return nil, errors.New("kingbase connection failed")
}
type kingbaseEndpoint struct {
host string
port int
}
// clusterConnectEndpoints splits the host field into ordered endpoints for
// cluster (multi-IP) configurations. DBX stores cluster hosts as
// comma-separated entries, each optionally embedding its own `:port`
// (mirroring the vastbase/openGauss driver semantics); semicolons are also
// accepted because Kingbase deployments paste both separators. A native
// connection_string builds its own DSN, so the host field is only split for
// the field-based path, and a host without any separator yields a single
// endpoint which keeps the legacy single-host flow untouched.
func clusterConnectEndpoints(cp connectParams) []kingbaseEndpoint {
if value := strings.TrimSpace(cp.ConnectionString); value != "" && !isKingbaseJDBCURL(value) {
return nil
}
defaultPort := cp.Port
if defaultPort <= 0 {
defaultPort = 54321
}
return splitHostEndpoints(cp.Host, defaultPort)
}
// splitHostEndpoints parses `host1[:port1][,;]host2[:port2]...`. Entries
// without an embedded port use fallbackPort; bracketed IPv6 literals may
// carry a port after the closing bracket.
func splitHostEndpoints(rawHost string, fallbackPort int) []kingbaseEndpoint {
parts := strings.FieldsFunc(rawHost, func(r rune) bool { return r == ',' || r == ';' })
endpoints := make([]kingbaseEndpoint, 0, len(parts))
for _, part := range parts {
host, port, ok := parseHostEndpoint(strings.TrimSpace(part), fallbackPort)
if !ok {
continue
}
endpoints = append(endpoints, kingbaseEndpoint{host: host, port: port})
}
return endpoints
}
func parseHostEndpoint(entry string, fallbackPort int) (string, int, bool) {
if entry == "" {
return "", 0, false
}
if strings.HasPrefix(entry, "[") {
if close := strings.Index(entry, "]"); close > 0 {
host := entry[1:close]
suffix := entry[close+1:]
portText, hasPort := strings.CutPrefix(suffix, ":")
if !hasPort && suffix == "" {
return host, fallbackPort, true
}
if hasPort {
if port, err := strconv.Atoi(portText); err == nil && validEndpointPort(port) {
return host, port, true
}
}
// Malformed bracketed entries such as `[::1]x` keep their raw
// text so the eventual dial error stays honest.
return entry, fallbackPort, true
}
}
if strings.Count(entry, ":") == 1 {
if host, rawPort, ok := strings.Cut(entry, ":"); ok && host != "" {
if port, err := strconv.Atoi(rawPort); err == nil && validEndpointPort(port) {
return host, port, true
}
}
}
// Bare IPv6 literals and hostnames (including entries with a non-numeric
// port suffix) are forwarded verbatim, exactly like the single-host path.
return entry, fallbackPort, true
}
func validEndpointPort(port int) bool {
return port > 0 && port <= 65535
}
// openAndPingClusterEndpoints tries each configured endpoint in order and
// keeps the first connection that completes the sslmode sequence, matching
// the failover behavior the openGauss-family drivers implement natively.
// Following libpq semantics, every endpoint gets its own full timeout
// budget so one unreachable node cannot starve the remaining candidates.
func openAndPingClusterEndpoints(
cp connectParams,
timeout time.Duration,
endpoints []kingbaseEndpoint,
opener kingbaseDBOpener,
) (*sql.DB, error) {
failures := make([]error, 0, len(endpoints))
for _, endpoint := range endpoints {
endpointParams := cp
endpointParams.Host = endpoint.host
endpointParams.Port = endpoint.port
ctx, cancel := context.WithTimeout(context.Background(), timeout)
db, err := pingDBWithSSLModes(ctx, endpointParams, opener)
cancel()
if err == nil {
return db, nil
}
failures = append(failures, fmt.Errorf("%s:%d: %w", endpoint.host, endpoint.port, err))
}
return nil, fmt.Errorf("kingbase connection failed after trying %d endpoints: %w", len(endpoints), errors.Join(failures...))
}
func shouldRetryKingbaseWithoutSSL(err error) bool {
if errors.Is(err, gokb.ErrSSLNotSupported) {
return true
}
if err == nil {
return false
}
// KingBase V7 can accept the SSLRequest and then reject the TLS handshake,
// so gokb returns the TLS alert instead of ErrSSLNotSupported.
return strings.Contains(strings.ToLower(err.Error()), "remote error: tls: handshake failure")
}
func openDBWithSSLMode(cp connectParams, sslMode string) (*sql.DB, error) {
dsn := buildDSNWithSSLMode(cp, sslMode)
db, err := sql.Open("kingbase", dsn)
if err != nil {
return nil, err
}
// Each protocol session is serialized and owns one database connection.
// Keeping a single physical connection preserves session state such as
// search_path and avoids extra pool coordination on the hot query path.
db.SetMaxOpenConns(1)
db.SetMaxIdleConns(1)
db.SetConnMaxLifetime(5 * time.Minute)
return db, nil
}
func (s *server) disconnect() error {
s.cancelActiveQuery()
s.closeAllQuerySessions()
s.usePgDefaultExpression = false
s.usePgViewDefinition = false
s.usePgFunctionDefinition = false
s.useLegacyRoutineDefinition = false
s.catalogIdentityUnsupported = false
s.catalogOIDUnsupported = false
s.infoColumnTypeUnsupported = false
s.infoUdtNameUnsupported = false
s.indexOrdinalityUnsupported = false
s.triggerPrettyUnsupported = false
s.triggerInternalUnsupported = false
s.constraintDefinitionUnsupported = false
s.constraintValidatedUnsupported = false
s.constraintStatusUnsupported = false
s.currentSchema = ""
s.schemaSet = false
if s.db == nil {
return nil
}
err := s.db.Close()
s.db = nil
return err
}
func (s *server) validateConnection() error {
db, err := s.requireDB()
if err != nil {
return err
}
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
return db.PingContext(ctx)
}
func (s *server) requireDB() (*sql.DB, error) {
if s.db == nil {
return nil, errors.New("not connected")
}
return s.db, nil
}
func (s *server) beginOperation(timeoutSecs int) (context.Context, context.CancelFunc) {
ctx := context.Background()
var cancel context.CancelFunc
if timeoutSecs > 0 {
ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSecs)*time.Second)
} else {
ctx, cancel = context.WithCancel(ctx)
}
s.activeCancelMu.Lock()
s.activeCancel = cancel
s.activeCancelMu.Unlock()
return ctx, cancel
}
func (s *server) endOperation(cancel context.CancelFunc) {
cancel()
s.activeCancelMu.Lock()
s.activeCancel = nil
s.activeCancelMu.Unlock()
}
func (s *server) cancelActiveQuery() {
s.activeCancelMu.Lock()
cancel := s.activeCancel
s.activeCancelMu.Unlock()
if cancel != nil {
cancel()
}
}
func (s *server) executeQuery(opts queryOptions) (queryResult, error) {
start := time.Now()
sqlText := trimStatementSQL(opts.SQL)
if isQuerySQL(sqlText) {
rows, conn, cancel, err := s.queryRows(sqlText, opts.Schema, opts.TimeoutSecs)
if err != nil {
return queryResult{}, err
}
defer func() {
_ = rows.Close()
_ = conn.Close()
s.endOperation(cancel)
}()
maxRows := opts.MaxRows
if maxRows <= 0 {
maxRows = defaultMaxRows
}
result, err := readRows(rows, maxRows)
result.ExecutionTimeMS = time.Since(start).Milliseconds()
return result, err
}
conn, ctx, cancel, err := s.operationConn(opts.Schema, opts.TimeoutSecs)
if err != nil {
return queryResult{}, err
}
defer func() {
_ = conn.Close()
s.endOperation(cancel)
}()
execResult, err := conn.ExecContext(ctx, sqlText)
if err != nil {
return queryResult{}, err
}
affected, _ := execResult.RowsAffected()
return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil
}
func (s *server) queryRows(sqlText string, schema string, timeoutSecs int) (*sql.Rows, *sql.Conn, context.CancelFunc, error) {
conn, ctx, cancel, err := s.operationConn(schema, timeoutSecs)
if err != nil {
return nil, nil, nil, err
}
rows, err := conn.QueryContext(ctx, sqlText)
if err != nil {
_ = conn.Close()
s.endOperation(cancel)
return nil, nil, nil, err
}
return rows, conn, cancel, nil
}
func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageResult, error) {
start := time.Now()
sqlText := trimStatementSQL(opts.SQL)
if !isQuerySQL(sqlText) {
result, err := s.executeQuery(opts)
return queryPageResult{Columns: result.Columns, ColumnTypes: result.ColumnTypes, Rows: result.Rows, AffectedRows: result.AffectedRows, ExecutionTimeMS: result.ExecutionTimeMS, Truncated: result.Truncated}, err
}
rows, conn, cancel, err := s.queryRows(sqlText, opts.Schema, opts.TimeoutSecs)
if err != nil {
return queryPageResult{}, err
}
columns, err := rows.Columns()
if err != nil {
_ = rows.Close()
_ = conn.Close()
s.endOperation(cancel)
return queryPageResult{}, err
}
columns = nonNilStrings(columns)
maxRows := opts.MaxRows
if maxRows <= 0 {
maxRows = defaultMaxRows
}
session := &querySession{rows: rows, conn: conn, columns: columns, columnTypes: columnTypeNames(rows), remaining: maxRows, cancel: cancel}
result, err := readQuerySessionPage(session, pageSize)
result.ExecutionTimeMS = time.Since(start).Milliseconds()
if err != nil {
_ = rows.Close()
_ = conn.Close()
s.endOperation(cancel)
return queryPageResult{}, err
}
if result.HasMore {
s.nextSessionID++
id := fmt.Sprintf("kingbase-%d", s.nextSessionID)
s.sessions[id] = session
result.SessionID = &id
} else {
_ = rows.Close()
_ = conn.Close()
s.endOperation(cancel)
}
return result, nil
}
func (s *server) fetchQueryPage(id string, pageSize int) (queryPageResult, error) {
session := s.sessions[id]
if session == nil {
return queryPageResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}}, nil
}
result, err := readQuerySessionPage(session, pageSize)
if err != nil {
s.closeQuerySession(id)
return queryPageResult{}, err
}
if result.HasMore {
result.SessionID = &id
} else {
s.closeQuerySession(id)
}
return result, nil
}
func (s *server) closeQuerySession(id string) bool {
session := s.sessions[id]
if session == nil {
return false
}
_ = session.rows.Close()
if session.conn != nil {
_ = session.conn.Close()
}
if session.cancel != nil {
s.endOperation(session.cancel)
}
delete(s.sessions, id)
return true
}
func (s *server) closeAllQuerySessions() {
for id := range s.sessions {
s.closeQuerySession(id)
}
}
func readQuerySessionPage(session *querySession, pageSize int) (queryPageResult, error) {
if pageSize <= 0 {
pageSize = 100
}
capacity := min(pageSize, session.remaining)
result := queryPageResult{Columns: session.columns, ColumnTypes: session.columnTypes, Rows: make([][]any, 0, capacity)}
for len(result.Rows) < pageSize && session.remaining > 0 {
if session.pending != nil {
result.Rows = append(result.Rows, session.pending)
session.pending = nil
session.remaining--
continue
}
if !session.rows.Next() {
return result, session.rows.Err()
}
row, err := scanRow(session.rows, len(session.columns), session.columnTypes)
if err != nil {
return queryPageResult{}, err
}
result.Rows = append(result.Rows, row)
session.remaining--
}
if session.remaining <= 0 {
result.Truncated = true
return result, nil
}
if session.rows.Next() {
row, err := scanRow(session.rows, len(session.columns), session.columnTypes)
if err != nil {
return queryPageResult{}, err
}
session.pending = row
result.HasMore = true
}
return result, session.rows.Err()
}
func readRows(rows *sql.Rows, maxRows int) (queryResult, error) {
columns, err := rows.Columns()
if err != nil {
return queryResult{}, err
}
columns = nonNilStrings(columns)
columnTypes := columnTypeNames(rows)
result := queryResult{Columns: columns, ColumnTypes: columnTypes, Rows: make([][]any, 0, min(maxRows, 1024))}
for rows.Next() {
if len(result.Rows) >= maxRows {
result.Truncated = true
break
}
row, err := scanRow(rows, len(columns), columnTypes)
if err != nil {
return queryResult{}, err
}
result.Rows = append(result.Rows, row)
}
return result, rows.Err()
}
func scanRow(rows *sql.Rows, count int, columnTypes []string) ([]any, error) {
storage := make([]any, count*2)
values := storage[:count]
dest := storage[count:]
for i := range values {
dest[i] = &values[i]
}
if err := rows.Scan(dest...); err != nil {
return nil, err
}
for i, value := range values {
values[i] = normalizeValue(value, columnTypeAt(columnTypes, i))
}
return values, nil
}
func columnTypeAt(columnTypes []string, index int) string {
if index < 0 || index >= len(columnTypes) {
return ""
}
return columnTypes[index]
}
func columnTypeNames(rows *sql.Rows) []string {
types, err := rows.ColumnTypes()
if err != nil {
return []string{}
}
result := make([]string, len(types))
for i, columnType := range types {
result[i] = columnType.DatabaseTypeName()
}
return result
}
func (s *server) executeTransaction(params map[string]json.RawMessage) (queryResult, error) {
statements := stringSliceParam(params, "statements")
conn, ctx, cancel, err := s.operationConn(stringParam(params, "schema"), intParam(params, "timeoutSecs"))
if err != nil {
return queryResult{}, err
}
defer func() {
_ = conn.Close()
s.endOperation(cancel)
}()
start := time.Now()
tx, err := conn.BeginTx(ctx, nil)
if err != nil {
return queryResult{}, err
}
var affected int64
for _, statement := range statements {
result, execErr := tx.ExecContext(ctx, trimStatementSQL(statement))
if execErr != nil {
_ = tx.Rollback()
return queryResult{}, execErr
}
rows, _ := result.RowsAffected()
affected += rows
}
if err := tx.Commit(); err != nil {
return queryResult{}, err
}
return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil
}
func (s *server) executeBatch(params map[string]json.RawMessage) (queryResult, error) {
start := time.Now()
var affected int64
for _, statement := range stringSliceParam(params, "statements") {
result, err := s.executeQuery(queryOptions{SQL: statement, Schema: stringParam(params, "schema")})
if err != nil {
return queryResult{}, err
}
affected += result.AffectedRows
}
return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil
}
func (s *server) operationConn(schema string, timeoutSecs int) (*sql.Conn, context.Context, context.CancelFunc, error) {
ctx, cancel := s.beginOperation(timeoutSecs)
conn, err := s.schemaConn(ctx, schema)
if err != nil {
s.endOperation(cancel)
return nil, nil, nil, err
}
return conn, ctx, cancel, nil
}
func (s *server) schemaConn(ctx context.Context, schema string) (*sql.Conn, error) {
db, err := s.requireDB()
if err != nil {
return nil, err
}
conn, err := db.Conn(ctx)
if err != nil {
return nil, err
}
if err := s.setSchema(ctx, conn, schema); err != nil {
_ = conn.Close()
return nil, err
}
return conn, nil
}
func (s *server) setSchema(ctx context.Context, conn *sql.Conn, schema string) error {
schema = strings.TrimSpace(schema)
// An omitted schema leaves the session search_path under user control.
// Reset it only after DBX applied an explicit schema.
if schema == "" && !s.schemaSet {
return nil
}
statement := "RESET search_path"
if schema != "" {
// Kingbase implicitly prioritizes its system catalog when it is not
// listed explicitly, matching the JDBC agent and DBeaver behavior.
statement = "SET search_path TO " + s.quoteIdentifier(schema)
}
if _, err := conn.ExecContext(ctx, statement); err != nil {
return err
}
s.currentSchema = schema
s.schemaSet = schema != ""
return nil
}
func buildDSN(cp connectParams) string {
sslMode := effectiveSSLMode(cp)
if sslMode == "prefer" {
sslMode = "require"
}
return buildDSNWithSSLMode(cp, sslMode)
}
func buildDSNWithSSLMode(cp connectParams, sslMode string) string {
if value := strings.TrimSpace(cp.ConnectionString); value != "" && !isKingbaseJDBCURL(value) {
return rewriteNativeConnectionStringSSLMode(value, sslMode)
}
port := cp.Port
if port <= 0 {
port = 54321
}
parts := []string{
"host=" + quoteDSNValue(cp.Host),
fmt.Sprintf("port=%d", port),
"user=" + quoteDSNValue(cp.Username),
"password=" + quoteDSNValue(cp.Password),
"dbname=" + quoteDSNValue(cp.Database),
"sslmode=" + sslMode,
}
if cp.CACertPath != "" {
parts = append(parts, "sslrootcert="+quoteDSNValue(cp.CACertPath))
}
if cp.ClientCertPath != "" {
parts = append(parts, "sslcert="+quoteDSNValue(cp.ClientCertPath))
}
if cp.ClientKeyPath != "" {
parts = append(parts, "sslkey="+quoteDSNValue(cp.ClientKeyPath))
}
// Classify and de-duplicate the app-supplied url_params. The connect_timeout
// default is only applied when the user did not provide one (natively or via
// the connectTimeout alias), so the parameter is never emitted twice.
urlParams := normalizeURLParams(cp.URLParams)
if !hasDSNParam(urlParams, "connect_timeout") {
parts = append(parts, "connect_timeout=15")
}
for _, p := range urlParams {
parts = append(parts, p.key+"="+quoteDSNValue(p.value))
}
return strings.Join(parts, " ")
}
func effectiveSSLMode(cp connectParams) string {
if value := strings.TrimSpace(cp.ConnectionString); value != "" && !isKingbaseJDBCURL(value) {
if sslMode, ok := nativeConnectionStringSSLMode(value); ok && sslMode != "" {
return sslMode
}
return "prefer"
}
sslMode := ""
for _, pair := range strings.FieldsFunc(cp.URLParams, func(r rune) bool { return r == '&' || r == ';' }) {
key, value, ok := strings.Cut(pair, "=")
if ok && strings.EqualFold(strings.TrimSpace(key), "sslmode") {
sslMode = strings.ToLower(strings.TrimSpace(value))
}
}
if sslMode != "" {
return sslMode
}
if cp.SSL {
return "verify-full"
}
return "prefer"
}
func nativeConnectionStringSSLMode(value string) (string, bool) {
if strings.HasPrefix(strings.ToLower(value), "kingbase://") {
query := value
if _, after, ok := strings.Cut(query, "?"); ok {
query = after
} else {
return "", false
}
query, _, _ = strings.Cut(query, "#")
sslMode := ""
found := false
for _, pair := range strings.Split(query, "&") {
key, rawValue, ok := strings.Cut(pair, "=")
if !ok {
continue
}
decodedKey, err := url.QueryUnescape(key)
if err != nil || !strings.EqualFold(decodedKey, "sslmode") {
continue
}
decodedValue, err := url.QueryUnescape(rawValue)
if err != nil {
decodedValue = rawValue
}
sslMode = strings.ToLower(strings.TrimSpace(decodedValue))
found = true
}
return sslMode, found
}
sslMode := ""
found := false
for _, field := range splitNativeDSNFields(value) {
key, rawValue, ok := strings.Cut(field, "=")
if !ok || !strings.EqualFold(strings.TrimSpace(key), "sslmode") {
continue
}
sslMode = strings.ToLower(unquoteNativeDSNValue(rawValue))
found = true
}
return sslMode, found
}
func rewriteNativeConnectionStringSSLMode(value, sslMode string) string {
if strings.HasPrefix(strings.ToLower(value), "kingbase://") {
baseAndQuery, fragment, hasFragment := strings.Cut(value, "#")
base, query, hasQuery := strings.Cut(baseAndQuery, "?")
params := make([]dsnParam, 0)
if hasQuery {
for _, pair := range strings.Split(query, "&") {
if pair == "" {
continue
}
rawKey, rawValue, _ := strings.Cut(pair, "=")
decodedKey, err := url.QueryUnescape(rawKey)
if err != nil {
decodedKey = rawKey
}
if strings.EqualFold(strings.TrimSpace(decodedKey), "sslmode") {
continue
}
decodedValue, err := url.QueryUnescape(rawValue)
if err != nil {
decodedValue = rawValue
}
nativeKey, keep := classifyDSNParam(decodedKey, decodedValue)
if !keep {
continue
}
params = append(params, dsnParam{
key: nativeKey,
value: rawValue, // preserve the original percent-encoding
fromAlias: !strings.EqualFold(strings.TrimSpace(decodedKey), nativeKey),
})
}
}
pairs := make([]string, 0, len(params)+1)
for _, p := range mergeDSNParams(params) {
pairs = append(pairs, url.QueryEscape(p.key)+"="+p.value)
}
pairs = append(pairs, "sslmode="+url.QueryEscape(sslMode))
result := base + "?" + strings.Join(pairs, "&")
if hasFragment {
result += "#" + fragment
}
return result
}
fields := splitNativeDSNFields(value)
params := make([]dsnParam, 0, len(fields))
passthrough := make([]string, 0)
for _, field := range fields {
key, rawValue, ok := strings.Cut(field, "=")
if !ok {
passthrough = append(passthrough, field)
continue
}
if strings.EqualFold(strings.TrimSpace(key), "sslmode") {
continue
}
nativeKey, keep := classifyDSNParam(key, unquoteNativeDSNValue(rawValue))
if !keep {
continue
}
params = append(params, dsnParam{
key: nativeKey,
value: rawValue, // preserve the original quoting
fromAlias: !strings.EqualFold(strings.TrimSpace(key), nativeKey),
})
}
result := make([]string, 0, len(params)+len(passthrough)+1)
for _, p := range mergeDSNParams(params) {
result = append(result, p.key+"="+p.value)
}
result = append(result, passthrough...)
result = append(result, "sslmode="+sslMode)
return strings.Join(result, " ")
}
func splitNativeDSNFields(value string) []string {
fields := make([]string, 0)
for index := 0; index < len(value); {
for index < len(value) && isNativeDSNSpace(value[index]) {
index++
}
if index >= len(value) {
break
}
start := index
for index < len(value) && value[index] != '=' {
index++
}
if index >= len(value) {
fields = append(fields, strings.TrimSpace(value[start:]))
break
}
index++
for index < len(value) && isNativeDSNSpace(value[index]) {
index++
}
quoted := index < len(value) && value[index] == '\''
if quoted {
index++
}
for index < len(value) {
if value[index] == '\\' && index+1 < len(value) {
index += 2
continue
}
if quoted {
if value[index] == '\'' {
index++
break
}
} else if isNativeDSNSpace(value[index]) {
break
}
index++
}
for index < len(value) && !isNativeDSNSpace(value[index]) {
index++
}
fields = append(fields, strings.TrimSpace(value[start:index]))
}
return fields
}
func unquoteNativeDSNValue(value string) string {
value = strings.TrimSpace(value)
if len(value) >= 2 && value[0] == '\'' && value[len(value)-1] == '\'' {
value = value[1 : len(value)-1]
}
return strings.TrimSpace(value)
}
func isNativeDSNSpace(value byte) bool {
return value == ' ' || value == '\t' || value == '\n' || value == '\r' || value == '\f'
}
func isKingbaseJDBCURL(value string) bool {
return strings.HasPrefix(strings.ToLower(strings.TrimSpace(value)), "jdbc:kingbase8://")
}
func quoteDSNValue(value string) string {
return "'" + strings.ReplaceAll(strings.ReplaceAll(value, `\`, `\\`), "'", `\'`) + "'"
}
// supportedDSNParams is the curated set of parameters known to be understood by
// the gokb driver or the Kingbase server. It is no longer a strict allow-list:
// classifyDSNParam also forwards unknown lower_snake_case names to the server as
// run-time parameters, because gokb passes every non-driver-setting to the
// startup packet (conn.go startup()). This set is what classifyDSNParam treats
// as definitely native, which short-circuits the camelCase JDBC heuristic so
// CamelCase GUCs such as DateStyle/TimeZone are still forwarded rather than
// dropped.
//
// The list mirrors the driver's own surface:
// - gokb conn.go isDriverSetting(): host, port, password, sslmode, sslcert,
// sslkey, sslrootcert, fallback_application_name, connect_timeout,
// disable_prepared_binary_result, binary_parameters, krbsrvname, krbspn;
// - the standard startup keywords user and dbname;
// - connector.go special handling: client_encoding (must be UTF8),
// datestyle, extra_float_digits;
// - common Kingbase/PostgreSQL run-time parameters that can be set in the
// startup packet: application_name, options, search_path,
// statement_timeout, work_mem, timezone and friends.
var supportedDSNParams = map[string]struct{}{
// gokb driver settings (conn.go isDriverSetting) and startup keywords
"host": {},
"port": {},
"user": {},
"password": {},
"dbname": {},
"sslmode": {},
"sslcert": {},
"sslkey": {},
"sslrootcert": {},
"fallback_application_name": {},
"connect_timeout": {},
"disable_prepared_binary_result": {},
"binary_parameters": {},
"krbsrvname": {},
"krbspn": {},
// connector.go special handling
"client_encoding": {},
"datestyle": {},
"extra_float_digits": {},
// Common run-time parameters the Kingbase server accepts in the startup
// packet (PostgreSQL-compatible GUCs).
"application_name": {},
"options": {},
"search_path": {},
"statement_timeout": {},
"lock_timeout": {},
"idle_in_transaction_session_timeout": {},
"idle_session_timeout": {},
"work_mem": {},
"maintenance_work_mem": {},
"temp_buffers": {},
"effective_cache_size": {},
"timezone": {},
"intervalstyle": {},
"lc_messages": {},
"lc_monetary": {},
"lc_numeric": {},
"lc_time": {},
"default_transaction_isolation": {},
"default_transaction_read_only": {},
"default_transaction_deferrable": {},
"synchronous_commit": {},
"client_min_messages": {},
"standard_conforming_strings": {},
"xmloption": {},
"role": {},
"session_replication_role": {},
"default_tablespace": {},
"temp_tablespaces": {},
"default_table_access_method": {},
"max_parallel_workers_per_gather": {},
}
func isSupportedDSNParam(key string) bool {
_, ok := supportedDSNParams[strings.ToLower(strings.TrimSpace(key))]
return ok
}
func isSafeParamKey(value string) bool {
value = strings.TrimSpace(value)
if value == "" {
return false
}
for _, char := range value {
if !(char == '_' || char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || char >= '0' && char <= '9') {
return false
}
}
return true
}
// dsnParam is a single normalized connection parameter ready to be emitted into
// a DSN. value carries the surface-specific text (single-quoted for keyword
// DSNs, percent-encoded for kingbase:// URLs, raw for url_params) so callers can
// preserve the original quoting/encoding when only the key was rewritten.
type dsnParam struct {
key string
value string
fromAlias bool
}
// jdbcAliasParams maps a lowercased JDBC property to the native gokb/server
// parameter with equivalent semantics. clientEncoding is handled separately in
// classifyDSNParam because it also has to validate the value.
var jdbcAliasParams = map[string]string{
"connecttimeout": "connect_timeout", // both measured in seconds
"currentschema": "search_path", // both accept a comma-separated list
"applicationname": "application_name",
}
// jdbcOnlyParams lists client-side JDBC/driver properties that have no meaning to
// the Kingbase server. gokb forwards every non-driver-setting to the startup
// packet, so a value the server does not recognize fails the whole connection
// with "unrecognized configuration parameter". camelCase names are also caught by
// the heuristic in classifyDSNParam; this set additionally covers the lowercase
// JDBC properties the heuristic cannot detect and documents intent for the common
// MySQL/JDBC-style names.
var jdbcOnlyParams = map[string]struct{}{
"usessl": {},
"autoreconnect": {},
"characterencoding": {},
"servertimezone": {},
"rewritebatchedstatements": {},
"useserverprepstmts": {},
"sockettimeout": {},
"usecompression": {},
"zerodatetimebehavior": {},
"useaffectedrows": {},
"usecursorfetch": {},
"defaultfetchsize": {},
"allowmultiqueries": {},
"useunicode": {},
// Lowercase PgJDBC/Kingbase-JDBC client properties the camelCase heuristic
// would otherwise forward and break the connection.
"ssl": {},
"sslfactory": {},
"stringtype": {},
"gsslib": {},
"sspiservicename": {},
"protocolversion": {},
"loglevel": {},
}
// classifyDSNParam decides how one connection parameter should be treated and
// returns the native parameter name to emit plus whether to keep it. sslmode is
// handled separately by the callers and must not be passed here. decodedValue is
// the already-unquoted/decoded value, used only for the client_encoding check.
func classifyDSNParam(key, decodedValue string) (nativeKey string, keep bool) {
trimmed := strings.TrimSpace(key)
if !isSafeParamKey(trimmed) {
return "", false
}
lower := strings.ToLower(trimmed)
// client_encoding (native, or via the clientEncoding alias): gokb only
// accepts UTF-8, so map compatible values and drop everything else — a
// non-UTF8 value would otherwise fail the whole connection.
if lower == "client_encoding" || lower == "clientencoding" {
if isUTF8Encoding(decodedValue) {
return "client_encoding", true
}
return "", false
}
// JDBC properties with a direct native equivalent.
if native, ok := jdbcAliasParams[lower]; ok {
return native, true
}
// Curated native/server parameters are always forwarded. Matching here also
// keeps CamelCase GUCs such as DateStyle/TimeZone from being mistaken for JDBC
// camelCase properties by the heuristic below.
if isSupportedDSNParam(lower) {
return lower, true
}
// Known JDBC-only client properties never reach the server.
if _, ok := jdbcOnlyParams[lower]; ok {
return "", false
}
// Unknown parameter. Server GUCs are conventionally lower_snake_case while
// JDBC properties are camelCase, so forward snake_case names as run-time
// parameters (gokb passes them to the startup packet) and drop names carrying
// an uppercase letter as presumed client-side JDBC settings.
if hasUpperASCII(trimmed) {
return "", false
}
return lower, true
}
// mergeDSNParams applies duplicate-parameter precedence: an explicit native
// parameter beats a JDBC alias for the same key, and within the same class the
// first occurrence wins to preserve gokb's existing DSN behavior. Output order
// follows each key's first appearance.
func mergeDSNParams(params []dsnParam) []dsnParam {
result := make([]dsnParam, 0, len(params))
pos := make(map[string]int, len(params))
for _, p := range params {
if i, ok := pos[p.key]; ok {
// A later explicit native parameter may replace an earlier alias, but
// same-class duplicates keep the first value just as gokb does.
if result[i].fromAlias && !p.fromAlias {
result[i] = p
}
continue
}
pos[p.key] = len(result)
result = append(result, p)
}
return result
}
// normalizeURLParams classifies and de-duplicates the app-supplied url_params
// blob (a &/;-separated key=value list), excluding sslmode which is handled
// separately. Values are kept raw for later single-quoting.
func normalizeURLParams(raw string) []dsnParam {
params := make([]dsnParam, 0)
for _, pair := range strings.FieldsFunc(raw, func(r rune) bool { return r == '&' || r == ';' }) {
key, value, ok := strings.Cut(pair, "=")
if !ok {
continue
}
if strings.EqualFold(strings.TrimSpace(key), "sslmode") {
continue
}
val := strings.TrimSpace(value)
nativeKey, keep := classifyDSNParam(key, val)
if !keep {
continue
}
params = append(params, dsnParam{
key: nativeKey,
value: val,
fromAlias: !strings.EqualFold(strings.TrimSpace(key), nativeKey),
})
}
return mergeDSNParams(params)
}
func hasDSNParam(params []dsnParam, key string) bool {
for _, p := range params {
if p.key == key {
return true
}
}
return false
}
func hasUpperASCII(value string) bool {
for i := 0; i < len(value); i++ {
if value[i] >= 'A' && value[i] <= 'Z' {
return true
}
}
return false
}
// isUTF8Encoding mirrors gokb's isUTF8: it recognizes fuzzy variants of "UTF-8"
// (dropping non-alphanumerics, case-insensitively) as well as "unicode".
func isUTF8Encoding(name string) bool {
var b strings.Builder
for _, ch := range name {
switch {
case ch >= 'A' && ch <= 'Z':
b.WriteRune(ch + ('a' - 'A'))
case ch >= 'a' && ch <= 'z', ch >= '0' && ch <= '9':
b.WriteRune(ch)
}
}
s := b.String()
return s == "utf8" || s == "unicode"
}
func normalizeValue(value any, columnTypeName string) any {
switch typed := value.(type) {
case nil:
return nil
case []byte:
if isTextBytes(typed) {
return string(typed)
}
return map[string]string{"$binary": base64.StdEncoding.EncodeToString(typed)}
case time.Time:
// KingBase "timestamp"/"date" columns are wall-clock values with no timezone
// meaning; the gokb driver decodes them into a time.Time labeled with the
// process-local zone, which is not a real UTC instant. Formatting those with
// an offset (RFC3339Nano) makes clients double-apply the timezone shift.
if isKingbaseTimezoneLessDateTime(columnTypeName) {
return typed.Format("2006-01-02T15:04:05.999999999")
}
return typed.Format(time.RFC3339Nano)
case int8:
return int64(typed)
case int16:
return int64(typed)
case int32:
return int64(typed)
case float32:
return float64(typed)
default:
return typed
}
}
func isKingbaseTimezoneLessDateTime(columnTypeName string) bool {
normalized := strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(columnTypeName), " ", ""))
switch normalized {
case "DATE", "TIMESTAMP", "TIME":
return true
default:
return false
}
}
func isTextBytes(value []byte) bool {
for _, char := range value {
if char == 0 || char < 0x09 || char > 0x0d && char < 0x20 {
return false
}
}
return true
}
func decodeParams(params map[string]json.RawMessage, target any) error {
data, err := json.Marshal(params)
if err != nil {
return err
}
return json.Unmarshal(data, target)
}
func stringParam(params map[string]json.RawMessage, key string) string {
var value string
_ = json.Unmarshal(params[key], &value)
return value
}
func intParam(params map[string]json.RawMessage, key string) int {
var value int
_ = json.Unmarshal(params[key], &value)
return value
}
func boolParam(params map[string]json.RawMessage, key string) bool {
var value bool
if raw, ok := params[key]; ok {
_ = json.Unmarshal(raw, &value)
}
return value
}
func stringSliceParam(params map[string]json.RawMessage, key string) []string {
var values []string
if json.Unmarshal(params[key], &values) == nil {
return values
}
return nil
}
func metadataListConstraintsFromParams(params map[string]json.RawMessage) metadataListConstraints {
return metadataListConstraints{
Filter: stringParam(params, "filter"),
Limit: intParam(params, "limit"),
Offset: intParam(params, "offset"),
ObjectTypes: stringSliceParam(params, "object_types"),
}
}
func errorResponse(id json.RawMessage, err error) response {
return response{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: -1, Message: err.Error()}}
}
func trimStatementSQL(sqlText string) string {
return strings.TrimRight(strings.TrimSpace(sqlText), "; \t\r\n")
}
func isQuerySQL(sqlText string) bool {
keyword, next := sqlKeywordAt(sqlText, 0)
if keyword == "with" {
terminal, terminalEnd := withTerminalKeyword(sqlText, next)
if terminal == "" || !isStatementKeyword(terminal) {
return true
}
keyword = terminal
next = terminalEnd
}
if keyword == "select" || keyword == "show" || keyword == "explain" || keyword == "values" || keyword == "table" {
return true
}
return isDMLKeyword(keyword) && hasTopLevelSQLKeyword(sqlText, next, "returning")
}
func withTerminalKeyword(sqlText string, index int) (string, int) {
if keyword, next := sqlKeywordAt(sqlText, index); keyword == "recursive" {
index = next
}
for {
index = skipSQLTrivia(sqlText, index)
index = skipSQLIdentifier(sqlText, index)
if index < 0 {
return "", index
}
index = skipSQLTrivia(sqlText, index)
if index < len(sqlText) && sqlText[index] == '(' {
index = skipSQLParenthesized(sqlText, index)
if index < 0 {
return "", index
}
}
keyword, next := sqlKeywordAt(sqlText, index)
if keyword != "as" {
return "", index
}
index = next
if keyword, next = sqlKeywordAt(sqlText, index); keyword == "not" {
index = next
keyword, next = sqlKeywordAt(sqlText, index)
if keyword != "materialized" {
return "", index
}
index = next
} else if keyword == "materialized" {
index = next
}
index = skipSQLTrivia(sqlText, index)
if index >= len(sqlText) || sqlText[index] != '(' {
return "", index
}
index = skipSQLParenthesized(sqlText, index)
if index < 0 {
return "", index
}
index = skipSQLTrivia(sqlText, index)
if index < len(sqlText) && sqlText[index] == ',' {
index++
continue
}
keyword, next = sqlKeywordAt(sqlText, index)
return keyword, next
}
}
func hasTopLevelSQLKeyword(sqlText string, index int, target string) bool {
depth := 0
for index < len(sqlText) {
switch sqlText[index] {
case '\'':
index = skipSQLQuoted(sqlText, index, '\'')
if index < 0 {
return false
}
continue
case '"':
index = skipSQLQuoted(sqlText, index, '"')
if index < 0 {
return false
}
continue
case '$':
if next := skipSQLDollarQuoted(sqlText, index); next != index {
if next < 0 {
return false
}
index = next
continue
}
case '-':
if index+1 < len(sqlText) && sqlText[index+1] == '-' {
index = skipSQLLineComment(sqlText, index+2)
continue
}
case '/':
if index+1 < len(sqlText) && sqlText[index+1] == '*' {
index = skipSQLBlockComment(sqlText, index)
if index < 0 {
return false
}
continue
}
case '(':
depth++
case ')':
depth = max(depth-1, 0)
default:
if depth == 0 && isSQLWordStartByte(sqlText[index]) {
keyword, next := sqlKeywordAt(sqlText, index)
if keyword == target {
return true
}
index = next
continue
}
}
index++
}
return false
}
func isDMLKeyword(keyword string) bool {
return keyword == "insert" || keyword == "update" || keyword == "delete" || keyword == "merge"
}
func isStatementKeyword(keyword string) bool {
return isDMLKeyword(keyword) || keyword == "select" || keyword == "values" || keyword == "table"
}
func sqlKeywordAt(sqlText string, index int) (string, int) {
index = skipSQLTrivia(sqlText, index)
if index >= len(sqlText) || !isSQLWordStartByte(sqlText[index]) {
return "", index
}
start := index
for index < len(sqlText) && isSQLIdentifierByte(sqlText[index]) {
index++
}
return strings.ToLower(sqlText[start:index]), index
}
func skipSQLIdentifier(sqlText string, index int) int {
index = skipSQLTrivia(sqlText, index)
if index >= len(sqlText) {
return -1
}
if sqlText[index] == '"' {
return skipSQLQuoted(sqlText, index, '"')
}
start := index
for index < len(sqlText) && isSQLIdentifierByte(sqlText[index]) {
index++
}
if start == index {
return -1
}
return index
}
func skipSQLParenthesized(sqlText string, index int) int {
if index >= len(sqlText) || sqlText[index] != '(' {
return -1
}
depth := 0
for index < len(sqlText) {
switch sqlText[index] {
case '\'':
index = skipSQLQuoted(sqlText, index, '\'')
if index < 0 {
return -1
}
continue
case '"':
index = skipSQLQuoted(sqlText, index, '"')
if index < 0 {
return -1
}
continue
case '$':
if next := skipSQLDollarQuoted(sqlText, index); next != index {
if next < 0 {
return -1
}
index = next
continue
}
case '-':
if index+1 < len(sqlText) && sqlText[index+1] == '-' {
index = skipSQLLineComment(sqlText, index+2)
continue
}
case '/':
if index+1 < len(sqlText) && sqlText[index+1] == '*' {
index = skipSQLBlockComment(sqlText, index)
if index < 0 {
return -1
}
continue
}
case '(':
depth++
case ')':
depth--
if depth == 0 {
return index + 1
}
}
index++
}
return -1
}
func skipSQLTrivia(sqlText string, index int) int {
for index < len(sqlText) {
if isSQLSpace(sqlText[index]) {
index++
continue
}
if index+1 < len(sqlText) && sqlText[index] == '-' && sqlText[index+1] == '-' {
index = skipSQLLineComment(sqlText, index+2)
continue
}
if index+1 < len(sqlText) && sqlText[index] == '/' && sqlText[index+1] == '*' {
index = skipSQLBlockComment(sqlText, index)
if index < 0 {
return len(sqlText)
}
continue
}
break
}
return index
}
func skipSQLLineComment(sqlText string, index int) int {
for index < len(sqlText) && sqlText[index] != '\n' {
index++
}
return index
}
func skipSQLBlockComment(sqlText string, index int) int {
depth := 0
for index+1 < len(sqlText) {
if sqlText[index] == '/' && sqlText[index+1] == '*' {
depth++
index += 2
continue
}
if sqlText[index] == '*' && sqlText[index+1] == '/' {
depth--
index += 2
if depth == 0 {
return index
}
continue
}
index++
}
return -1
}
func skipSQLQuoted(sqlText string, index int, quote byte) int {
index++
for index < len(sqlText) {
if sqlText[index] == '\\' && quote == '\'' && index+1 < len(sqlText) {
index += 2
continue
}
if sqlText[index] == quote {
if index+1 < len(sqlText) && sqlText[index+1] == quote {
index += 2
continue
}
return index + 1
}
index++
}
return -1
}
func skipSQLDollarQuoted(sqlText string, index int) int {
endTag := index + 1
for endTag < len(sqlText) && isSQLDollarTagByte(sqlText[endTag]) {
endTag++
}
if endTag >= len(sqlText) || sqlText[endTag] != '$' {
return index
}
tag := sqlText[index : endTag+1]
if end := strings.Index(sqlText[endTag+1:], tag); end >= 0 {
return endTag + 1 + end + len(tag)
}
return -1
}
func nonNilStrings(values []string) []string {
if values == nil {
return []string{}
}
return values
}
func isSQLSpace(value byte) bool {
return value == ' ' || value == '\t' || value == '\r' || value == '\n' || value == '\f'
}
func isSQLWordStartByte(value byte) bool {
return value >= 'a' && value <= 'z' || value >= 'A' && value <= 'Z' || value == '_'
}
func isSQLIdentifierByte(value byte) bool {
return isSQLWordStartByte(value) || value >= '0' && value <= '9' || value == '$' || value >= 0x80
}
func isSQLDollarTagByte(value byte) bool {
return isSQLWordStartByte(value) || value >= '0' && value <= '9'
}
func quoteIdentifier(value string) string {
return `"` + strings.ReplaceAll(value, `"`, `""`) + `"`
}
func (s *server) quoteIdentifier(value string) string {
if s.mode.mysqlCompat {
return "`" + strings.ReplaceAll(value, "`", "``") + "`"
}
return quoteIdentifier(value)
}
func quoteLiteral(value string) string {
return "'" + strings.ReplaceAll(value, "'", "''") + "'"
}
func stringPtr(value string) *string {
if value == "" {
return nil
}
return &value
}