Add runtime SQLite query targets

2026-09-01 17:15:09 +08:00
parent 65da73eff0
commit 163fbda394
17 changed files with 708 additions and 114 deletions
@@ -0,0 +1,312 @@
+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)
+}