LEM/pkg/lem/query.go

package lem

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"
)

// QueryOpts holds configuration for the query command.
type QueryOpts struct {
	DB   string // DuckDB database path (defaults to LEM_DB env)
	JSON bool   // Output as JSON instead of table
}

// RunQuery is the CLI entry point for the query command.
// Runs ad-hoc SQL against the DuckDB database.
// The args slice contains the SQL query as positional arguments.
func RunQuery(cfg QueryOpts, args []string) error {
	if cfg.DB == "" {
		cfg.DB = os.Getenv("LEM_DB")
	}
	if cfg.DB == "" {
		return fmt.Errorf("--db or LEM_DB required")
	}

	sql := strings.Join(args, " ")
	if sql == "" {
		return fmt.Errorf("SQL query required as positional argument\n  lem query --db path.duckdb \"SELECT * FROM golden_set LIMIT 5\"\n  lem query --db path.duckdb \"domain = 'ethics'\"  (auto-wraps as WHERE clause)")
	}

	// Auto-wrap non-SELECT queries as WHERE clauses.
	trimmed := strings.TrimSpace(strings.ToUpper(sql))
	if !strings.HasPrefix(trimmed, "SELECT") && !strings.HasPrefix(trimmed, "SHOW") &&
		!strings.HasPrefix(trimmed, "DESCRIBE") && !strings.HasPrefix(trimmed, "EXPLAIN") {
		sql = "SELECT * FROM golden_set WHERE " + sql + " LIMIT 20"
	}

	db, err := OpenDB(cfg.DB)
	if err != nil {
		return fmt.Errorf("open db: %w", err)
	}
	defer db.Close()

	rows, err := db.conn.Query(sql)
	if err != nil {
		return fmt.Errorf("query: %w", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return fmt.Errorf("columns: %w", err)
	}

	var results []map[string]any

	for rows.Next() {
		values := make([]any, len(cols))
		ptrs := make([]any, len(cols))
		for i := range values {
			ptrs[i] = &values[i]
		}

		if err := rows.Scan(ptrs...); err != nil {
			return fmt.Errorf("scan: %w", err)
		}

		row := make(map[string]any)
		for i, col := range cols {
			v := values[i]
			// Convert []byte to string for readability.
			if b, ok := v.([]byte); ok {
				v = string(b)
			}
			row[col] = v
		}
		results = append(results, row)
	}

	if cfg.JSON {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		return enc.Encode(results)
	}

	// Table output.
	if len(results) == 0 {
		fmt.Println("(no results)")
		return nil
	}

	// Calculate column widths.
	widths := make(map[string]int)
	for _, col := range cols {
		widths[col] = len(col)
	}
	for _, row := range results {
		for _, col := range cols {
			s := fmt.Sprintf("%v", row[col])
			if len(s) > 60 {
				s = s[:57] + "..."
			}
			if len(s) > widths[col] {
				widths[col] = len(s)
			}
		}
	}

	// Print header.
	for i, col := range cols {
		if i > 0 {
			fmt.Print("  ")
		}
		fmt.Printf("%-*s", widths[col], col)
	}
	fmt.Println()

	// Print separator.
	for i, col := range cols {
		if i > 0 {
			fmt.Print("  ")
		}
		fmt.Print(strings.Repeat("─", widths[col]))
	}
	fmt.Println()

	// Print rows.
	for _, row := range results {
		for i, col := range cols {
			if i > 0 {
				fmt.Print("  ")
			}
			s := fmt.Sprintf("%v", row[col])
			if len(s) > 60 {
				s = s[:57] + "..."
			}
			fmt.Printf("%-*s", widths[col], s)
		}
		fmt.Println()
	}

	fmt.Printf("\n(%d rows)\n", len(results))
	return nil
}