313 lines
9.8 KiB
Go
313 lines
9.8 KiB
Go
package runtime
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"browser.local/run/protocol"
|
|
)
|
|
|
|
const (
|
|
defaultSQLiteQueryMaxRows = 500
|
|
maxSQLiteQueryRows = 500
|
|
maxSQLiteQueryColumns = 128
|
|
maxSQLiteQuerySQLBytes = 64 * 1024
|
|
maxSQLiteQueryResultBytes = 64 * 1024
|
|
maxSQLiteQueryTimeout = 60 * time.Second
|
|
)
|
|
|
|
var sqliteNamedParameterPattern = regexp.MustCompile(`[:@$]([A-Za-z_][A-Za-z0-9_]*)`)
|
|
|
|
// SQLiteQueryExecutor executes a package-provided, read-only SQL asset against
|
|
// a declared SQLite snapshot inside the server's scoped Run workspace.
|
|
type SQLiteQueryExecutor struct{ resolver WorkspaceResolver }
|
|
|
|
func NewSQLiteQueryExecutor(workspaceRoot string) *SQLiteQueryExecutor {
|
|
return &SQLiteQueryExecutor{resolver: NewWorkspaceResolver(workspaceRoot)}
|
|
}
|
|
|
|
func (executor *SQLiteQueryExecutor) Execute(ctx context.Context, assignment protocol.RunJobAssignment) LifecycleExecutionResult {
|
|
if err := validateSQLiteQueryAssignment(assignment); err != nil {
|
|
return sqliteQueryFailure("invalid_request", false)
|
|
}
|
|
queryAsset, maxRows, err := sqliteQueryRequest(assignment.ExecutionInput.Inputs)
|
|
if err != nil {
|
|
return sqliteQueryFailure("invalid_request", false)
|
|
}
|
|
scope, err := executor.resolver.Scope(assignment.ServerInstanceID, assignment.ExecutionInput.WorkspaceScope)
|
|
if err != nil {
|
|
return sqliteQueryFailure("target_unavailable", false)
|
|
}
|
|
databasePath, err := executor.resolver.ExistingTarget(scope, assignment.TargetKey)
|
|
if err != nil {
|
|
return sqliteQueryFailure("target_unavailable", false)
|
|
}
|
|
sqlPath, err := executor.resolver.ExistingTarget(scope, queryAsset)
|
|
if err != nil {
|
|
return sqliteQueryFailure("sql_asset_unavailable", false)
|
|
}
|
|
statement, err := readSQLiteQueryAsset(sqlPath)
|
|
if err != nil {
|
|
return sqliteQueryFailure(sqliteQueryAssetErrorCode(err), false)
|
|
}
|
|
if !isReadOnlySQLiteStatement(statement) {
|
|
return sqliteQueryFailure("sql_asset_not_read_only", false)
|
|
}
|
|
|
|
queryCtx, cancel := sqliteQueryContext(ctx, assignment.ExecutionInput.TimeoutSeconds)
|
|
defer cancel()
|
|
database, err := sql.Open(sqliteSchemaProbeDriver, "file:"+databasePath+"?mode=ro")
|
|
if err != nil {
|
|
return sqliteQueryFailure("sqlite_open_failed", true)
|
|
}
|
|
defer database.Close()
|
|
database.SetMaxOpenConns(1)
|
|
if _, err := database.ExecContext(queryCtx, "PRAGMA query_only = ON"); err != nil {
|
|
return sqliteQueryFailure(sqliteQueryErrorCode(queryCtx, err), true)
|
|
}
|
|
rows, err := database.QueryContext(queryCtx, statement, sqliteQueryArguments(statement, maxRows)...)
|
|
if err != nil {
|
|
return sqliteQueryFailure(sqliteQueryErrorCode(queryCtx, err), true)
|
|
}
|
|
defer rows.Close()
|
|
content, err := collectSQLiteQueryRows(rows, maxRows)
|
|
if err != nil {
|
|
return sqliteQueryFailure(sqliteQueryErrorCode(queryCtx, err), false)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return sqliteQueryFailure(sqliteQueryErrorCode(queryCtx, err), true)
|
|
}
|
|
return LifecycleExecutionResult{
|
|
State: lifecycleResultStateSucceeded,
|
|
Progress: protocol.RunJobProgressReport{Percent: 100, Message: "SQLite query completed"},
|
|
Message: "SQLite query completed",
|
|
ExecutionResult: protocol.RunJobExecutionResult{
|
|
Kind: "sqlite.query",
|
|
Checksum: checksumForText(content),
|
|
SizeBytes: int64(len([]byte(content))),
|
|
Summary: "bounded query-only SQLite rows",
|
|
Content: content,
|
|
},
|
|
}
|
|
}
|
|
|
|
func validateSQLiteQueryAssignment(assignment protocol.RunJobAssignment) error {
|
|
if assignment.Capability != protocol.RunCapabilityRemoteRunDBSQLiteQuery || assignment.ServerInstanceID == "" || assignment.ExecutionInput.WorkspaceScope == "" {
|
|
return fmt.Errorf("SQLite query requires a scoped database job")
|
|
}
|
|
if err := protocol.ValidateRunJobAssignment(assignment); err != nil {
|
|
return err
|
|
}
|
|
if assignment.ExecutionInput.RemoteAdapterKind != "" && assignment.ExecutionInput.RemoteAdapterKind != "database" {
|
|
return fmt.Errorf("SQLite query requires the database adapter")
|
|
}
|
|
if assignment.ExecutionInput.RemoteAdapterKey != "" && !protocol.ValidLogicalFileKey(assignment.ExecutionInput.RemoteAdapterKey) {
|
|
return fmt.Errorf("SQLite query adapter key is invalid")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func sqliteQueryRequest(inputs map[string]string) (string, int, error) {
|
|
if len(inputs) == 0 {
|
|
return "", 0, fmt.Errorf("SQLite query inputs are missing")
|
|
}
|
|
queryAsset := strings.TrimSpace(inputs["sqlRef"])
|
|
if queryAsset == "" || queryAsset != inputs["sqlRef"] || !protocol.ValidLogicalFileKey(queryAsset) || !strings.HasPrefix(queryAsset, "sql/") || !strings.HasSuffix(strings.ToLower(queryAsset), ".sql") {
|
|
return "", 0, fmt.Errorf("SQLite query SQL asset is invalid")
|
|
}
|
|
maxRows := defaultSQLiteQueryMaxRows
|
|
if rawMaxRows, exists := inputs["maxRows"]; exists {
|
|
parsed, err := strconv.Atoi(rawMaxRows)
|
|
if err != nil || parsed < 1 || parsed > maxSQLiteQueryRows {
|
|
return "", 0, fmt.Errorf("SQLite query row limit is invalid")
|
|
}
|
|
maxRows = parsed
|
|
}
|
|
return queryAsset, maxRows, nil
|
|
}
|
|
|
|
func readSQLiteQueryAsset(path string) (string, error) {
|
|
info, err := os.Stat(path)
|
|
if err != nil || !info.Mode().IsRegular() || info.Size() < 1 || info.Size() > maxSQLiteQuerySQLBytes {
|
|
return "", errSQLiteQueryAssetInvalid
|
|
}
|
|
body, err := os.ReadFile(path)
|
|
if err != nil || len(body) == 0 || len(body) > maxSQLiteQuerySQLBytes {
|
|
return "", errSQLiteQueryAssetInvalid
|
|
}
|
|
statement := strings.TrimSpace(string(body))
|
|
if statement == "" || strings.ContainsRune(statement, '\x00') {
|
|
return "", errSQLiteQueryAssetInvalid
|
|
}
|
|
return statement, nil
|
|
}
|
|
|
|
var errSQLiteQueryAssetInvalid = errors.New("SQLite query asset is invalid")
|
|
|
|
func sqliteQueryAssetErrorCode(err error) string {
|
|
if errors.Is(err, errSQLiteQueryAssetInvalid) {
|
|
return "sql_asset_invalid"
|
|
}
|
|
return "sql_asset_unavailable"
|
|
}
|
|
|
|
func isReadOnlySQLiteStatement(statement string) bool {
|
|
trimmed := strings.TrimSpace(statement)
|
|
trimmed = strings.TrimSuffix(trimmed, ";")
|
|
if trimmed == "" || strings.Contains(trimmed, ";") {
|
|
return false
|
|
}
|
|
parts := strings.Fields(strings.ToUpper(strings.TrimSpace(stripSQLiteLeadingComments(trimmed))))
|
|
if len(parts) == 0 {
|
|
return false
|
|
}
|
|
switch parts[0] {
|
|
case "SELECT", "WITH":
|
|
return true
|
|
case "EXPLAIN":
|
|
return len(parts) > 1 && (parts[1] == "SELECT" || len(parts) > 2 && parts[1] == "QUERY" && parts[2] == "PLAN")
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func stripSQLiteLeadingComments(statement string) string {
|
|
for {
|
|
statement = strings.TrimSpace(statement)
|
|
switch {
|
|
case strings.HasPrefix(statement, "--"):
|
|
if lineEnd := strings.IndexByte(statement, '\n'); lineEnd >= 0 {
|
|
statement = statement[lineEnd+1:]
|
|
continue
|
|
}
|
|
return ""
|
|
case strings.HasPrefix(statement, "/*"):
|
|
commentEnd := strings.Index(statement[2:], "*/")
|
|
if commentEnd < 0 {
|
|
return ""
|
|
}
|
|
statement = statement[commentEnd+4:]
|
|
continue
|
|
}
|
|
return statement
|
|
}
|
|
}
|
|
|
|
func sqliteQueryContext(ctx context.Context, timeoutSeconds int) (context.Context, context.CancelFunc) {
|
|
timeout := maxSQLiteQueryTimeout
|
|
if timeoutSeconds > 0 && time.Duration(timeoutSeconds)*time.Second < timeout {
|
|
timeout = time.Duration(timeoutSeconds) * time.Second
|
|
}
|
|
return context.WithTimeout(ctx, timeout)
|
|
}
|
|
|
|
func sqliteQueryArguments(statement string, maxRows int) []any {
|
|
seen := map[string]struct{}{}
|
|
arguments := make([]any, 0)
|
|
for _, match := range sqliteNamedParameterPattern.FindAllStringSubmatch(statement, -1) {
|
|
name := match[1]
|
|
if _, exists := seen[name]; exists {
|
|
continue
|
|
}
|
|
seen[name] = struct{}{}
|
|
value := any(nil)
|
|
if name == "limit" {
|
|
value = maxRows
|
|
}
|
|
arguments = append(arguments, sql.Named(name, value))
|
|
}
|
|
return arguments
|
|
}
|
|
|
|
func collectSQLiteQueryRows(rows *sql.Rows, maxRows int) (string, error) {
|
|
columns, err := rows.Columns()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if len(columns) == 0 || len(columns) > maxSQLiteQueryColumns || !uniqueSQLiteQueryColumns(columns) {
|
|
return "", errSQLiteQueryResultLimit
|
|
}
|
|
result := struct {
|
|
Rows []map[string]any `json:"rows"`
|
|
}{Rows: make([]map[string]any, 0, maxRows)}
|
|
for rows.Next() {
|
|
if len(result.Rows) >= maxRows {
|
|
return "", errSQLiteQueryResultLimit
|
|
}
|
|
values := make([]any, len(columns))
|
|
pointers := make([]any, len(values))
|
|
for index := range values {
|
|
pointers[index] = &values[index]
|
|
}
|
|
if err := rows.Scan(pointers...); err != nil {
|
|
return "", err
|
|
}
|
|
row := make(map[string]any, len(columns))
|
|
for index, column := range columns {
|
|
row[column] = sqliteQueryValue(values[index])
|
|
}
|
|
result.Rows = append(result.Rows, row)
|
|
payload, err := json.Marshal(result)
|
|
if err != nil || len(payload) > maxSQLiteQueryResultBytes {
|
|
return "", errSQLiteQueryResultLimit
|
|
}
|
|
}
|
|
payload, err := json.Marshal(result)
|
|
if err != nil || len(payload) > maxSQLiteQueryResultBytes {
|
|
return "", errSQLiteQueryResultLimit
|
|
}
|
|
return string(payload), nil
|
|
}
|
|
|
|
var errSQLiteQueryResultLimit = errors.New("SQLite query result exceeds limit")
|
|
|
|
func uniqueSQLiteQueryColumns(columns []string) bool {
|
|
seen := map[string]struct{}{}
|
|
for _, column := range columns {
|
|
if strings.TrimSpace(column) == "" {
|
|
return false
|
|
}
|
|
if _, exists := seen[column]; exists {
|
|
return false
|
|
}
|
|
seen[column] = struct{}{}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func sqliteQueryValue(value any) any {
|
|
if bytes, ok := value.([]byte); ok {
|
|
return append([]byte(nil), bytes...)
|
|
}
|
|
return value
|
|
}
|
|
|
|
func sqliteQueryErrorCode(ctx context.Context, err error) string {
|
|
if errors.Is(err, errSQLiteQueryResultLimit) {
|
|
return "result_limit_exceeded"
|
|
}
|
|
if errors.Is(ctx.Err(), context.Canceled) {
|
|
return "cancelled"
|
|
}
|
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
|
return "timeout"
|
|
}
|
|
lower := strings.ToLower(fmt.Sprint(err))
|
|
if strings.Contains(lower, "locked") || strings.Contains(lower, "busy") {
|
|
return "database_busy"
|
|
}
|
|
return "sqlite_read_failed"
|
|
}
|
|
|
|
func sqliteQueryFailure(code string, retryable bool) LifecycleExecutionResult {
|
|
return lifecycleExecutionFailure(code, "SQLite query failed", retryable)
|
|
}
|