Files
spark-mcp/internal/storage/sqlite.go
T
tao.chenandClaude 2c39091706 Phase 3+4: 共享 HTTP 客户端 + admin API + audit log
Phase 3 (httpclient):
- shared Client (Timeout + MaxResponseBytes 截断)
- SSRF: 16 段私网 CIDR (含 169.254/::ffff:0:0/96 IPv4-mapped)
  + hostname allowlist (精确/后缀/*. 通配)
  + scheme 校验 (仅 http/https)
- auth 适配器: none / simple (YARN user.name) / basic (SetBasicAuth)
- DoWithRedirect: 跨主机跳转保留 Authorization, max 5 默认
- 19 个测试全绿 (httptest 模拟)

Phase 4 (admin + audit):
- internal/audit: Entry + Repo.Insert/List + MarshalDetails, snake_case JSON
- internal/middleware: AdminAuth/AgentAuth (constant-time 比对)
- internal/admin: 6 端点 (GET/POST/PUT/DELETE /admin/clusters, GET /admin/audit)
  + 写操作触发 audit_log (含 before/after diff)
  + AuthPassword 空=保留旧密码
  + 校验: 必填字段 + auth_type 枚举 + rate_limit>=0
- main.go: 挂载 storage.Open + admin.Mount
- L fix: audit.Entry 加 json tag (smoke test 发现大写 key 不规范)

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-10 13:54:24 +08:00

89 lines
2.4 KiB
Go

package storage
import (
"database/sql"
"fmt"
"os"
"path/filepath"
_ "modernc.org/sqlite"
)
// Schema defines the SQLite schema for clusters and audit_log.
const Schema = `CREATE TABLE IF NOT EXISTS clusters (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
rm_url TEXT NOT NULL,
shs_url TEXT NOT NULL,
spark_submit_execute_bin TEXT NOT NULL,
is_active INTEGER NOT NULL DEFAULT 1,
auth_type TEXT NOT NULL DEFAULT 'none',
auth_username TEXT,
auth_password_enc BLOB,
ssl_verify INTEGER NOT NULL DEFAULT 1,
ssl_ca_bundle TEXT,
url_allowlist TEXT,
default_submit_args TEXT,
rate_limit_per_min INTEGER NOT NULL DEFAULT 10,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS audit_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp INTEGER NOT NULL,
actor TEXT NOT NULL,
action TEXT NOT NULL,
cluster_id TEXT,
details TEXT
);
CREATE INDEX IF NOT EXISTS idx_audit_ts ON audit_log(timestamp DESC);
`
// DB is the storage handle.
type DB struct {
Path string
sqlDB *sql.DB
}
// Open opens the SQLite database at path, creating the parent directory if
// needed. It enables WAL mode and foreign keys, then applies the schema.
func Open(path string) (*DB, error) {
dir := filepath.Dir(path)
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, fmt.Errorf("storage: mkdir %s: %w", dir, err)
}
sqlDB, err := sql.Open("sqlite", path)
if err != nil {
return nil, fmt.Errorf("storage: open sqlite %s: %w", path, err)
}
if _, err := sqlDB.Exec("PRAGMA journal_mode=WAL"); err != nil {
_ = sqlDB.Close()
return nil, fmt.Errorf("storage: set WAL mode: %w", err)
}
if _, err := sqlDB.Exec("PRAGMA foreign_keys=ON"); err != nil {
_ = sqlDB.Close()
return nil, fmt.Errorf("storage: enable foreign keys: %w", err)
}
d := &DB{Path: path, sqlDB: sqlDB}
if err := d.Migrate(); err != nil {
_ = sqlDB.Close()
return nil, err
}
return d, nil
}
// Close releases the underlying database connection.
func (d *DB) Close() error {
if d == nil || d.sqlDB == nil {
return nil
}
return d.sqlDB.Close()
}
// SQLDB returns the underlying *sql.DB handle.
func (d *DB) SQLDB() *sql.DB { return d.sqlDB }