Skip to content

Commit 9b9a5ef

Browse files
weicaoweicao
authored
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.
1 parent 2eb4143 commit 9b9a5ef

29 files changed

Lines changed: 2793 additions & 123 deletions

‎backend/handler.go‎

Lines changed: 43 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"fmt"
2121

2222
"github.com/apecloud/myduckserver/catalog"
23+
"github.com/apecloud/myduckserver/mycontext"
2324
sqle "github.com/dolthub/go-mysql-server"
2425
"github.com/dolthub/go-mysql-server/server"
2526
"github.com/dolthub/go-mysql-server/sql"
@@ -45,7 +46,8 @@ func (h *MyHandler) ConnectionClosed(c *mysql.Conn) {
4546
}
4647

4748
func (h *MyHandler) ComInitDB(c *mysql.Conn, schemaName string) error {
48-
_, err := h.provider.Pool().GetConnForSchema(context.Background(), c.ConnectionID, schemaName)
49+
ctx := mycontext.WithFrontendQuery(context.Background())
50+
_, err := h.provider.Pool().GetConnForSchema(ctx, c.ConnectionID, schemaName)
4951
if err != nil {
5052
return err
5153
}
@@ -79,7 +81,11 @@ func (h *MyHandler) ComMultiQuery(
7981
query string,
8082
callback mysql.ResultSpoolFn,
8183
) (rest string, returnErr error) {
82-
audit := NewQueryAudit(c, "mysql", query)
84+
if err := catalog.RejectSensitiveSQL(query); err != nil {
85+
return query, err
86+
}
87+
ctx = mycontext.WithFrontendQuery(ctx)
88+
audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
8389
defer func() {
8490
audit.Complete(returnErr)
8591
}()
@@ -121,7 +127,11 @@ func (h *MyHandler) ComQuery(
121127
query string,
122128
callback mysql.ResultSpoolFn,
123129
) (returnErr error) {
124-
audit := NewQueryAudit(c, "mysql", query)
130+
if err := catalog.RejectSensitiveSQL(query); err != nil {
131+
return err
132+
}
133+
ctx = mycontext.WithFrontendQuery(ctx)
134+
audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
125135
defer func() {
126136
audit.Complete(returnErr)
127137
}()
@@ -160,7 +170,14 @@ func (h *MyHandler) ComStmtExecute(ctx context.Context, c *mysql.Conn, prepare *
160170
if prepare != nil {
161171
query = prepare.PrepareStmt
162172
}
163-
audit := NewQueryAudit(c, "mysql", query)
173+
if err := catalog.RejectSensitiveSQL(query); err != nil {
174+
return err
175+
}
176+
ctx = mycontext.WithFrontendQuery(ctx)
177+
if returnErr = h.rejectReadOnly(ctx, c, query); returnErr != nil {
178+
return returnErr
179+
}
180+
audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
164181
defer func() {
165182
audit.Complete(returnErr)
166183
}()
@@ -185,28 +202,44 @@ func (h *MyHandler) ComStmtExecute(ctx context.Context, c *mysql.Conn, prepare *
185202
}
186203

187204
func (h *MyHandler) ComPrepare(ctx context.Context, c *mysql.Conn, query string, prepare *mysql.PrepareData) ([]*querypb.Field, error) {
205+
if err := catalog.RejectSensitiveSQL(query); err != nil {
206+
return nil, err
207+
}
208+
ctx = mycontext.WithFrontendQuery(ctx)
188209
if err := h.rejectReadOnly(ctx, c, query); err != nil {
189210
return nil, err
190211
}
191212
return h.Handler.ComPrepare(ctx, c, query, prepare)
192213
}
193214

194215
func (h *MyHandler) ComPrepareParsed(ctx context.Context, c *mysql.Conn, query string, parsed sqlparser.Statement, prepare *mysql.PrepareData) (mysql.ParsedQuery, []*querypb.Field, error) {
216+
if err := catalog.RejectSensitiveSQL(query); err != nil {
217+
return nil, nil, err
218+
}
219+
ctx = mycontext.WithFrontendQuery(ctx)
195220
if err := h.rejectReadOnly(ctx, c, query); err != nil {
196221
return nil, nil, err
197222
}
198223
return h.Handler.ComPrepareParsed(ctx, c, query, parsed, prepare)
199224
}
200225

201226
func (h *MyHandler) ComBind(ctx context.Context, c *mysql.Conn, query string, parsedQuery mysql.ParsedQuery, prepare *mysql.PrepareData) (mysql.BoundQuery, []*querypb.Field, error) {
227+
if err := catalog.RejectSensitiveSQL(query); err != nil {
228+
return nil, nil, err
229+
}
230+
ctx = mycontext.WithFrontendQuery(ctx)
202231
if err := h.rejectReadOnly(ctx, c, query); err != nil {
203232
return nil, nil, err
204233
}
205234
return h.Handler.ComBind(ctx, c, query, parsedQuery, prepare)
206235
}
207236

208237
func (h *MyHandler) ComExecuteBound(ctx context.Context, c *mysql.Conn, query string, boundQuery mysql.BoundQuery, callback mysql.ResultSpoolFn) (returnErr error) {
209-
audit := NewQueryAudit(c, "mysql", query)
238+
if err := catalog.RejectSensitiveSQL(query); err != nil {
239+
return err
240+
}
241+
ctx = mycontext.WithFrontendQuery(ctx)
242+
audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
210243
defer func() {
211244
audit.Complete(returnErr)
212245
}()
@@ -227,7 +260,11 @@ func (h *MyHandler) ComExecuteBound(ctx context.Context, c *mysql.Conn, query st
227260
}
228261

229262
func (h *MyHandler) ComParsedQuery(ctx context.Context, c *mysql.Conn, query string, parsed sqlparser.Statement, callback mysql.ResultSpoolFn) (returnErr error) {
230-
audit := NewQueryAudit(c, "mysql", query)
263+
if err := catalog.RejectSensitiveSQL(query); err != nil {
264+
return err
265+
}
266+
ctx = mycontext.WithFrontendQuery(ctx)
267+
audit := NewQueryAudit(c, "mysql", catalog.RedactSensitiveSQL(query))
231268
defer func() {
232269
audit.Complete(returnErr)
233270
}()

‎backend/query_audit.go‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,21 @@
11
package backend
22

33
import (
4+
"github.com/apecloud/myduckserver/catalog"
45
"github.com/dolthub/vitess/go/mysql"
56
"github.com/sirupsen/logrus"
67
)
78

89
// QueryAudit records the outcome of one ordinary protocol query.
910
type QueryAudit struct {
10-
entry *logrus.Entry
11-
rows uint64
11+
entry *logrus.Entry
12+
rows uint64
13+
redactErrors bool
1214
}
1315

1416
// NewQueryAudit starts an audit record for a user-facing protocol query.
1517
func NewQueryAudit(conn *mysql.Conn, protocol, query string) *QueryAudit {
18+
query = catalog.RedactSensitiveSQL(query)
1619
fields := logrus.Fields{
1720
"audit": "query",
1821
"protocol": protocol,
@@ -22,7 +25,10 @@ func NewQueryAudit(conn *mysql.Conn, protocol, query string) *QueryAudit {
2225
if conn != nil {
2326
fields["user"] = conn.User
2427
}
25-
return &QueryAudit{entry: logrus.WithFields(fields)}
28+
return &QueryAudit{
29+
entry: logrus.WithFields(fields),
30+
redactErrors: query == catalog.RedactedSensitiveSQL,
31+
}
2632
}
2733

2834
// AddRows records rows successfully handed to the protocol callback.
@@ -34,7 +40,15 @@ func (audit *QueryAudit) AddRows(rows int) {
3440
func (audit *QueryAudit) Complete(err error) {
3541
entry := audit.entry.WithField("rows", audit.rows)
3642
if err != nil {
37-
entry = entry.WithError(err)
43+
if audit.redactErrors {
44+
// Legacy BACKUP/RESTORE keeps its historical storage path, but its
45+
// downstream errors can echo endpoint or credential material. Keep the
46+
// audit useful without retaining the raw error alongside the redacted
47+
// query.
48+
entry = entry.WithField("error", "service-managed operation failed")
49+
} else {
50+
entry = entry.WithError(err)
51+
}
3852
}
3953
entry.Info("query audit")
4054
}

‎backend/query_audit_test.go‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package backend
2+
3+
import (
4+
"bytes"
5+
"errors"
6+
"testing"
7+
8+
"github.com/apecloud/myduckserver/catalog"
9+
"github.com/dolthub/vitess/go/mysql"
10+
"github.com/sirupsen/logrus"
11+
"github.com/stretchr/testify/require"
12+
)
13+
14+
func TestLegacyObjectStorageAuditRedactsQueryAndCompletionError(t *testing.T) {
15+
queries := []string{
16+
"BACKUP DATABASE app TO 's3://bucket/app/' ENDPOINT='https://s3.example.test:9443' ACCESS_KEY_ID='access-key-71' SECRET_ACCESS_KEY='secret-71'",
17+
"RESTORE DATABASE app FROM 's3://bucket/app/' ENDPOINT='https://s3.example.test:9443' ACCESS_KEY_ID='access-key-71' SECRET_ACCESS_KEY='secret-71'",
18+
"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",
19+
"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'",
20+
}
21+
forbidden := []string{"s3.example.test", "access-key-71", "secret-71", "s3://bucket/app/"}
22+
23+
logger := logrus.StandardLogger()
24+
oldOut, oldLevel, oldFormatter := logger.Out, logger.Level, logger.Formatter
25+
defer func() {
26+
logger.Out = oldOut
27+
logger.Level = oldLevel
28+
logger.Formatter = oldFormatter
29+
}()
30+
logger.SetLevel(logrus.InfoLevel)
31+
logger.SetFormatter(&logrus.TextFormatter{DisableTimestamp: true})
32+
33+
for _, query := range queries {
34+
var output bytes.Buffer
35+
logger.SetOutput(&output)
36+
require.Equal(t, catalog.RedactedSensitiveSQL, catalog.RedactSensitiveSQL(query))
37+
38+
// Pass the raw protocol text to exercise the audit boundary itself; callers
39+
// should not need to remember a separate redaction step.
40+
audit := NewQueryAudit(&mysql.Conn{}, "mysql", query)
41+
audit.Complete(errors.New("upload failed endpoint=https://s3.example.test:9443 access=access-key-71 secret=secret-71"))
42+
43+
logged := output.String()
44+
require.Contains(t, logged, catalog.RedactedSensitiveSQL)
45+
require.Contains(t, logged, "service-managed operation failed")
46+
for _, value := range forbidden {
47+
require.NotContains(t, logged, value)
48+
}
49+
}
50+
}

‎catalog/connpool.go‎

Lines changed: 68 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ type ConnectionPool struct {
3434
conns sync.Map // concurrent-safe map[uint32]*stdsql.Conn
3535
txns sync.Map // concurrent-safe map[uint32]*stdsql.Tx
3636
closedConns sync.Map // connection IDs that completed their lifecycle
37+
initializerMu sync.RWMutex
38+
connectionInitializer func(context.Context, *stdsql.Conn) error
3739
registerMySQLUDFsOnce sync.Once
3840
registerMySQLUDFsErr error
3941
}
@@ -50,6 +52,32 @@ func (p *ConnectionPool) Connector() *duckdb.Connector {
5052
return p.connector
5153
}
5254

55+
// SetConnectionInitializer installs a hook that runs whenever a logical
56+
// session acquires a connection outside an active session transaction. The
57+
// hook receives the acquisition context, including its query-origin
58+
// classification. Running it for both new and reused logical connections
59+
// prevents a connection that was previously used by one origin from carrying
60+
// session settings into another origin; GetTxn runs it before BeginTx and then
61+
// keeps transaction-scoped state stable until that transaction closes.
62+
func (p *ConnectionPool) SetConnectionInitializer(initializer func(context.Context, *stdsql.Conn) error) {
63+
p.initializerMu.Lock()
64+
p.connectionInitializer = initializer
65+
p.initializerMu.Unlock()
66+
}
67+
68+
func (p *ConnectionPool) initializeConnection(ctx context.Context, conn *stdsql.Conn) error {
69+
p.initializerMu.RLock()
70+
initializer := p.connectionInitializer
71+
p.initializerMu.RUnlock()
72+
if initializer == nil {
73+
return nil
74+
}
75+
if ctx == nil {
76+
ctx = context.Background()
77+
}
78+
return initializer(ctx, conn)
79+
}
80+
5381
// CurrentSchema retrieves the current schema of the connection.
5482
// Returns an empty string if the connection is not established
5583
// or the schema cannot be retrieved.
@@ -88,13 +116,22 @@ func (p *ConnectionPool) CurrentCatalog(id uint32) string {
88116
}
89117

90118
func (p *ConnectionPool) GetConn(ctx context.Context, id uint32) (*stdsql.Conn, error) {
119+
if ctx == nil {
120+
ctx = context.Background()
121+
}
91122
var conn *stdsql.Conn
92123
entry, ok := p.conns.Load(id)
93124
if !ok {
94125
c, err := p.DB.Conn(ctx)
95126
if err != nil {
96127
return nil, err
97128
}
129+
if _, transactionActive := p.txns.Load(id); !transactionActive {
130+
if err := p.initializeConnection(ctx, c); err != nil {
131+
_ = c.Close()
132+
return nil, err
133+
}
134+
}
98135
if err := p.registerMySQLUDFs(c); err != nil {
99136
_ = c.Close()
100137
return nil, err
@@ -104,6 +141,22 @@ func (p *ConnectionPool) GetConn(ctx context.Context, id uint32) (*stdsql.Conn,
104141
conn = c
105142
} else {
106143
conn = entry.(*stdsql.Conn)
144+
// A session transaction owns this connection's transaction-scoped
145+
// state. Re-running the initializer here would execute LOAD/CREATE
146+
// SECRET inside that transaction and could alter or roll back with user
147+
// work. GetTxn initializes before BeginTx; keep the state stable until
148+
// the transaction is closed.
149+
if _, transactionActive := p.txns.Load(id); !transactionActive {
150+
if err := p.initializeConnection(ctx, conn); err != nil {
151+
// Do not leave a failed or partially initialized connection available
152+
// to a later request. CompareAndDelete avoids removing a replacement
153+
// installed by a concurrent recovery path.
154+
p.conns.CompareAndDelete(id, conn)
155+
p.closedConns.Store(id, struct{}{})
156+
_ = conn.Close()
157+
return nil, err
158+
}
159+
}
107160
}
108161
return conn, nil
109162
}
@@ -127,18 +180,25 @@ func (p *ConnectionPool) registerMySQLUDFs(conn *stdsql.Conn) error {
127180
}
128181

129182
func (p *ConnectionPool) GetConnForSchema(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, error) {
183+
if ctx == nil {
184+
ctx = context.Background()
185+
}
130186
conn, err := p.GetConn(ctx, id)
131187
if err != nil {
132188
return nil, err
133189
}
134190

135191
if schemaName != "" {
192+
// Schema selection is session state, but it should retain the origin
193+
// value while avoiding cancellation of a request that is already being
194+
// serviced.
195+
schemaCtx := context.WithoutCancel(ctx)
136196
var currentSchema string
137-
if err := conn.QueryRowContext(context.Background(), "SELECT CURRENT_SCHEMA").Scan(&currentSchema); err != nil {
197+
if err := conn.QueryRowContext(schemaCtx, "SELECT CURRENT_SCHEMA").Scan(&currentSchema); err != nil {
138198
logrus.WithError(err).Error("Failed to get current schema")
139199
return nil, err
140200
} else if currentSchema != schemaName {
141-
if _, err := conn.ExecContext(context.Background(), "USE "+FullSchemaName(p.CurrentCatalog(id), schemaName)); err != nil {
201+
if _, err := conn.ExecContext(schemaCtx, "USE "+FullSchemaName(p.CurrentCatalog(id), schemaName)); err != nil {
142202
if IsDuckDBSetSchemaNotFoundError(err) {
143203
return nil, sql.ErrDatabaseNotFound.New(schemaName)
144204
}
@@ -184,6 +244,9 @@ func (p *ConnectionPool) CloseConn(id uint32) error {
184244
}
185245

186246
func (p *ConnectionPool) GetTxn(ctx context.Context, id uint32, schemaName string, options *stdsql.TxOptions) (*stdsql.Tx, error) {
247+
if ctx == nil {
248+
ctx = context.Background()
249+
}
187250
var tx *stdsql.Tx
188251
entry, ok := p.txns.Load(id)
189252
if !ok {
@@ -242,6 +305,9 @@ func (p *ConnectionPool) Close() error {
242305
lastErr = err
243306
}
244307
}
308+
p.conns.Clear()
309+
p.txns.Clear()
310+
p.closedConns.Clear()
245311
return errors.Join(lastErr, p.DB.Close())
246312
}
247313

0 commit comments

Comments
 (0)