111 lines
3.2 KiB
Go
111 lines
3.2 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);
|
|
|
|
CREATE TABLE IF NOT EXISTS upload_files (
|
|
file_id TEXT PRIMARY KEY,
|
|
name TEXT NOT NULL,
|
|
size INTEGER NOT NULL,
|
|
sha256 TEXT NOT NULL,
|
|
uploaded_at INTEGER NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_uploads_ts ON upload_files(uploaded_at DESC);
|
|
CREATE TABLE IF NOT EXISTS spark_submissions (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
file_id TEXT NOT NULL,
|
|
app_id TEXT NOT NULL,
|
|
tracking_url TEXT,
|
|
cluster_id TEXT NOT NULL,
|
|
exit_code INTEGER NOT NULL,
|
|
duration_ms INTEGER NOT NULL,
|
|
submitted_at INTEGER NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_submissions_app_id ON spark_submissions(app_id);
|
|
CREATE INDEX IF NOT EXISTS idx_submissions_submitted_at ON spark_submissions(submitted_at DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_submissions_file_id ON spark_submissions(file_id);
|
|
`
|
|
|
|
// 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 }
|