feat: add offline DuckLake service wiring (#488)

* feat: add offline DuckLake service wiring

Task #71 reviewed DuckLake service wiring candidate, rebased onto current main.\n\nAssembled through a temporary ref so the final commit is platform-signed; no task #75 changes.

* fix: parse executable SQL comments at the protocol boundary

Task #71 follow-up: allow hosted MySQL executable comments that carry ordinary CREATE DATABASE clauses, while rejecting persistent secret creation. This replacement changes only catalog/sensitive_sql.go and is based on the existing signed PR head.

Author: weicao

backend/handler.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 
 	"github.com/apecloud/myduckserver/catalog"
+	"github.com/apecloud/myduckserver/mycontext"
 	sqle "github.com/dolthub/go-mysql-server"
 	"github.com/dolthub/go-mysql-server/server"
 	"github.com/dolthub/go-mysql-server/sql"
@@ -45,7 +46,8 @@ func (h *MyHandler) ConnectionClosed(c *mysql.Conn) {
 }
 
 func (h *MyHandler) ComInitDB(c *mysql.Conn, schemaName string) error {
-	_, err := h.provider.Pool().GetConnForSchema(context.Background(), c.ConnectionID, schemaName)
+	ctx := mycontext.WithFrontendQuery(context.Background())
+	_, err := h.provider.Pool().GetConnForSchema(ctx, c.ConnectionID, schemaName)
 	if err != nil {
 		return err
 	}
@@ -79,7 +81,11 @@ func (h *MyHandler) ComMultiQuery(
 	query string,
 	callback mysql.ResultSpoolFn,
 ) (rest string, returnErr error) {
-	audit := NewQueryAudit(c, "mysql", query)
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return query, err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
+	audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
 	defer func() {
 		audit.Complete(returnErr)
 	}()
@@ -121,7 +127,11 @@ func (h *MyHandler) ComQuery(
 	query string,
 	callback mysql.ResultSpoolFn,
 ) (returnErr error) {
-	audit := NewQueryAudit(c, "mysql", query)
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
+	audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
 	defer func() {
 		audit.Complete(returnErr)
 	}()
@@ -160,7 +170,14 @@ func (h *MyHandler) ComStmtExecute(ctx context.Context, c *mysql.Conn, prepare *
 	if prepare != nil {
 		query = prepare.PrepareStmt
 	}
-	audit := NewQueryAudit(c, "mysql", query)
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
+	if returnErr = h.rejectReadOnly(ctx, c, query); returnErr != nil {
+		return returnErr
+	}
+	audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
 	defer func() {
 		audit.Complete(returnErr)
 	}()
@@ -185,28 +202,44 @@ func (h *MyHandler) ComStmtExecute(ctx context.Context, c *mysql.Conn, prepare *
 }
 
 func (h *MyHandler) ComPrepare(ctx context.Context, c *mysql.Conn, query string, prepare *mysql.PrepareData) ([]*querypb.Field, error) {
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return nil, err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
 	if err := h.rejectReadOnly(ctx, c, query); err != nil {
 		return nil, err
 	}
 	return h.Handler.ComPrepare(ctx, c, query, prepare)
 }
 
 func (h *MyHandler) ComPrepareParsed(ctx context.Context, c *mysql.Conn, query string, parsed sqlparser.Statement, prepare *mysql.PrepareData) (mysql.ParsedQuery, []*querypb.Field, error) {
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return nil, nil, err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
 	if err := h.rejectReadOnly(ctx, c, query); err != nil {
 		return nil, nil, err
 	}
 	return h.Handler.ComPrepareParsed(ctx, c, query, parsed, prepare)
 }
 
 func (h *MyHandler) ComBind(ctx context.Context, c *mysql.Conn, query string, parsedQuery mysql.ParsedQuery, prepare *mysql.PrepareData) (mysql.BoundQuery, []*querypb.Field, error) {
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return nil, nil, err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
 	if err := h.rejectReadOnly(ctx, c, query); err != nil {
 		return nil, nil, err
 	}
 	return h.Handler.ComBind(ctx, c, query, parsedQuery, prepare)
 }
 
 func (h *MyHandler) ComExecuteBound(ctx context.Context, c *mysql.Conn, query string, boundQuery mysql.BoundQuery, callback mysql.ResultSpoolFn) (returnErr error) {
-	audit := NewQueryAudit(c, "mysql", query)
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
+	audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
 	defer func() {
 		audit.Complete(returnErr)
 	}()
@@ -227,7 +260,11 @@ func (h *MyHandler) ComExecuteBound(ctx context.Context, c *mysql.Conn, query st
 }
 
 func (h *MyHandler) ComParsedQuery(ctx context.Context, c *mysql.Conn, query string, parsed sqlparser.Statement, callback mysql.ResultSpoolFn) (returnErr error) {
-	audit := NewQueryAudit(c, "mysql", query)
+	if err := catalog.RejectSensitiveSQL(query); err != nil {
+		return err
+	}
+	ctx = mycontext.WithFrontendQuery(ctx)
+	audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
 	defer func() {
 		audit.Complete(returnErr)
 	}()

backend/query_audit.go

@@ -1,18 +1,21 @@
 package backend
 
 import (
+	"github.com/apecloud/myduckserver/catalog"
 	"github.com/dolthub/vitess/go/mysql"
 	"github.com/sirupsen/logrus"
 )
 
 // QueryAudit records the outcome of one ordinary protocol query.
 type QueryAudit struct {
-	entry *logrus.Entry
-	rows  uint64
+	entry        *logrus.Entry
+	rows         uint64
+	redactErrors bool
 }
 
 // NewQueryAudit starts an audit record for a user-facing protocol query.
 func NewQueryAudit(conn *mysql.Conn, protocol, query string) *QueryAudit {
+	query = catalog.RedactSensitiveSQL(query)
 	fields := logrus.Fields{
 		"audit":    "query",
 		"protocol": protocol,
@@ -22,7 +25,10 @@ func NewQueryAudit(conn *mysql.Conn, protocol, query string) *QueryAudit {
 	if conn != nil {
 		fields["user"] = conn.User
 	}
-	return &QueryAudit{entry: logrus.WithFields(fields)}
+	return &QueryAudit{
+		entry:        logrus.WithFields(fields),
+		redactErrors: query == catalog.RedactedSensitiveSQL,
+	}
 }
 
 // AddRows records rows successfully handed to the protocol callback.
@@ -34,7 +40,15 @@ func (audit *QueryAudit) AddRows(rows int) {
 func (audit *QueryAudit) Complete(err error) {
 	entry := audit.entry.WithField("rows", audit.rows)
 	if err != nil {
-		entry = entry.WithError(err)
+		if audit.redactErrors {
+			// Legacy BACKUP/RESTORE keeps its historical storage path, but its
+			// downstream errors can echo endpoint or credential material. Keep the
+			// audit useful without retaining the raw error alongside the redacted
+			// query.
+			entry = entry.WithField("error", "service-managed operation failed")
+		} else {
+			entry = entry.WithError(err)
+		}
 	}
 	entry.Info("query audit")
 }

backend/query_audit_test.go

@@ -0,0 +1,50 @@
+package backend
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+
+	"github.com/apecloud/myduckserver/catalog"
+	"github.com/dolthub/vitess/go/mysql"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLegacyObjectStorageAuditRedactsQueryAndCompletionError(t *testing.T) {
+	queries := []string{
+		"BACKUP DATABASE app TO 's3://bucket/app/' ENDPOINT='https://s3.example.test:9443' ACCESS_KEY_ID='access-key-71' SECRET_ACCESS_KEY='secret-71'",
+		"RESTORE DATABASE app FROM 's3://bucket/app/' ENDPOINT='https://s3.example.test:9443' ACCESS_KEY_ID='access-key-71' SECRET_ACCESS_KEY='secret-71'",
+		"BACKUP DATABASE app TO 's3://bucket/app/' ENDPOINT='https://s3.example.test:9443' ACCESS_KEY_ID='access-key-71' SECRET_ACCESS_KEY='secret-71'; SELECT 1",
+		"SELECT 1; RESTORE DATABASE app FROM 's3://bucket/app/' ENDPOINT='https://s3.example.test:9443' ACCESS_KEY_ID='access-key-71' SECRET_ACCESS_KEY='secret-71'",
+	}
+	forbidden := []string{"s3.example.test", "access-key-71", "secret-71", "s3://bucket/app/"}
+
+	logger := logrus.StandardLogger()
+	oldOut, oldLevel, oldFormatter := logger.Out, logger.Level, logger.Formatter
+	defer func() {
+		logger.Out = oldOut
+		logger.Level = oldLevel
+		logger.Formatter = oldFormatter
+	}()
+	logger.SetLevel(logrus.InfoLevel)
+	logger.SetFormatter(&logrus.TextFormatter{DisableTimestamp: true})
+
+	for _, query := range queries {
+		var output bytes.Buffer
+		logger.SetOutput(&output)
+		require.Equal(t, catalog.RedactedSensitiveSQL, catalog.RedactSensitiveSQL(query))
+
+		// Pass the raw protocol text to exercise the audit boundary itself; callers
+		// should not need to remember a separate redaction step.
+		audit := NewQueryAudit(&mysql.Conn{}, "mysql", query)
+		audit.Complete(errors.New("upload failed endpoint=https://s3.example.test:9443 access=access-key-71 secret=secret-71"))
+
+		logged := output.String()
+		require.Contains(t, logged, catalog.RedactedSensitiveSQL)
+		require.Contains(t, logged, "service-managed operation failed")
+		for _, value := range forbidden {
+			require.NotContains(t, logged, value)
+		}
+	}
+}

catalog/connpool.go

@@ -34,6 +34,8 @@ type ConnectionPool struct {
 	conns                 sync.Map // concurrent-safe map[uint32]*stdsql.Conn
 	txns                  sync.Map // concurrent-safe map[uint32]*stdsql.Tx
 	closedConns           sync.Map // connection IDs that completed their lifecycle
+	initializerMu         sync.RWMutex
+	connectionInitializer func(context.Context, *stdsql.Conn) error
 	registerMySQLUDFsOnce sync.Once
 	registerMySQLUDFsErr  error
 }
@@ -50,6 +52,32 @@ func (p *ConnectionPool) Connector() *duckdb.Connector {
 	return p.connector
 }
 
+// SetConnectionInitializer installs a hook that runs whenever a logical
+// session acquires a connection outside an active session transaction. The
+// hook receives the acquisition context, including its query-origin
+// classification. Running it for both new and reused logical connections
+// prevents a connection that was previously used by one origin from carrying
+// session settings into another origin; GetTxn runs it before BeginTx and then
+// keeps transaction-scoped state stable until that transaction closes.
+func (p *ConnectionPool) SetConnectionInitializer(initializer func(context.Context, *stdsql.Conn) error) {
+	p.initializerMu.Lock()
+	p.connectionInitializer = initializer
+	p.initializerMu.Unlock()
+}
+
+func (p *ConnectionPool) initializeConnection(ctx context.Context, conn *stdsql.Conn) error {
+	p.initializerMu.RLock()
+	initializer := p.connectionInitializer
+	p.initializerMu.RUnlock()
+	if initializer == nil {
+		return nil
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	return initializer(ctx, conn)
+}
+
 // CurrentSchema retrieves the current schema of the connection.
 // Returns an empty string if the connection is not established
 // or the schema cannot be retrieved.
@@ -88,13 +116,22 @@ func (p *ConnectionPool) CurrentCatalog(id uint32) string {
 }
 
 func (p *ConnectionPool) GetConn(ctx context.Context, id uint32) (*stdsql.Conn, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	var conn *stdsql.Conn
 	entry, ok := p.conns.Load(id)
 	if !ok {
 		c, err := p.DB.Conn(ctx)
 		if err != nil {
 			return nil, err
 		}
+		if _, transactionActive := p.txns.Load(id); !transactionActive {
+			if err := p.initializeConnection(ctx, c); err != nil {
+				_ = c.Close()
+				return nil, err
+			}
+		}
 		if err := p.registerMySQLUDFs(c); err != nil {
 			_ = c.Close()
 			return nil, err
@@ -104,6 +141,22 @@ func (p *ConnectionPool) GetConn(ctx context.Context, id uint32) (*stdsql.Conn,
 		conn = c
 	} else {
 		conn = entry.(*stdsql.Conn)
+		// A session transaction owns this connection's transaction-scoped
+		// state. Re-running the initializer here would execute LOAD/CREATE
+		// SECRET inside that transaction and could alter or roll back with user
+		// work. GetTxn initializes before BeginTx; keep the state stable until
+		// the transaction is closed.
+		if _, transactionActive := p.txns.Load(id); !transactionActive {
+			if err := p.initializeConnection(ctx, conn); err != nil {
+				// Do not leave a failed or partially initialized connection available
+				// to a later request. CompareAndDelete avoids removing a replacement
+				// installed by a concurrent recovery path.
+				p.conns.CompareAndDelete(id, conn)
+				p.closedConns.Store(id, struct{}{})
+				_ = conn.Close()
+				return nil, err
+			}
+		}
 	}
 	return conn, nil
 }
@@ -127,18 +180,25 @@ func (p *ConnectionPool) registerMySQLUDFs(conn *stdsql.Conn) error {
 }
 
 func (p *ConnectionPool) GetConnForSchema(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	conn, err := p.GetConn(ctx, id)
 	if err != nil {
 		return nil, err
 	}
 
 	if schemaName != "" {
+		// Schema selection is session state, but it should retain the origin
+		// value while avoiding cancellation of a request that is already being
+		// serviced.
+		schemaCtx := context.WithoutCancel(ctx)
 		var currentSchema string
-		if err := conn.QueryRowContext(context.Background(), "SELECT CURRENT_SCHEMA").Scan(&currentSchema); err != nil {
+		if err := conn.QueryRowContext(schemaCtx, "SELECT CURRENT_SCHEMA").Scan(&currentSchema); err != nil {
 			logrus.WithError(err).Error("Failed to get current schema")
 			return nil, err
 		} else if currentSchema != schemaName {
-			if _, err := conn.ExecContext(context.Background(), "USE "+FullSchemaName(p.CurrentCatalog(id), schemaName)); err != nil {
+			if _, err := conn.ExecContext(schemaCtx, "USE "+FullSchemaName(p.CurrentCatalog(id), schemaName)); err != nil {
 				if IsDuckDBSetSchemaNotFoundError(err) {
 					return nil, sql.ErrDatabaseNotFound.New(schemaName)
 				}
@@ -184,6 +244,9 @@ func (p *ConnectionPool) CloseConn(id uint32) error {
 }
 
 func (p *ConnectionPool) GetTxn(ctx context.Context, id uint32, schemaName string, options *stdsql.TxOptions) (*stdsql.Tx, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	var tx *stdsql.Tx
 	entry, ok := p.txns.Load(id)
 	if !ok {
@@ -242,6 +305,9 @@ func (p *ConnectionPool) Close() error {
 			lastErr = err
 		}
 	}
+	p.conns.Clear()
+	p.txns.Clear()
+	p.closedConns.Clear()
 	return errors.Join(lastErr, p.DB.Close())
 }