lemma/server/internal/db/query.go

package db

import (
	"fmt"
	"strings"
)

type JoinType string

const (
	InnerJoin JoinType = "INNER JOIN"
	LeftJoin  JoinType = "LEFT JOIN"
	RightJoin JoinType = "RIGHT JOIN"
)

// Query represents a SQL query with its parameters
type Query struct {
	builder     strings.Builder
	args        []interface{}
	dbType      DBType
	pos         int // tracks the current placeholder position
	hasSelect   bool
	hasFrom     bool
	hasWhere    bool
	hasOrderBy  bool
	hasGroupBy  bool
	hasLimit    bool
	hasOffset   bool
	isInParens  bool
	parensDepth int
}

// NewQuery creates a new Query instance
func NewQuery(dbType DBType) *Query {
	return &Query{
		dbType: dbType,
		args:   make([]interface{}, 0),
	}
}

// Select adds a SELECT clause
func (q *Query) Select(columns ...string) *Query {
	if !q.hasSelect {
		q.Write("SELECT ")
		q.Write(strings.Join(columns, ", "))
		q.hasSelect = true
	}
	return q
}

// From adds a FROM clause
func (q *Query) From(table string) *Query {
	if !q.hasFrom {
		q.Write(" FROM ")
		q.Write(table)
		q.hasFrom = true
	}
	return q
}

// Where adds a WHERE clause
func (q *Query) Where(condition string) *Query {
	if !q.hasWhere {
		q.Write(" WHERE ")
		q.hasWhere = true
	} else {
		q.Write(" AND ")
	}
	q.Write(condition)
	return q
}

// WhereIn adds a WHERE IN clause
func (q *Query) WhereIn(column string, count int) *Query {
	if !q.hasWhere {
		q.Write(" WHERE ")
		q.hasWhere = true
	} else {
		q.Write(" AND ")
	}
	q.Write(column)
	q.Write(" IN (")
	q.Placeholders(count)
	q.Write(")")
	return q
}

// And adds an AND condition
func (q *Query) And(condition string) *Query {
	q.Write(" AND ")
	q.Write(condition)
	return q
}

// Or adds an OR condition
func (q *Query) Or(condition string) *Query {
	q.Write(" OR ")
	q.Write(condition)
	return q
}

// Join adds a JOIN clause
func (q *Query) Join(joinType JoinType, table, condition string) *Query {
	q.Write(" ")
	q.Write(string(joinType))
	q.Write(" ")
	q.Write(table)
	q.Write(" ON ")
	q.Write(condition)
	return q
}

// OrderBy adds an ORDER BY clause
func (q *Query) OrderBy(columns ...string) *Query {
	if !q.hasOrderBy {
		q.Write(" ORDER BY ")
		q.Write(strings.Join(columns, ", "))
		q.hasOrderBy = true
	}
	return q
}

// GroupBy adds a GROUP BY clause
func (q *Query) GroupBy(columns ...string) *Query {
	if !q.hasGroupBy {
		q.Write(" GROUP BY ")
		q.Write(strings.Join(columns, ", "))
		q.hasGroupBy = true
	}
	return q
}

// Limit adds a LIMIT clause
func (q *Query) Limit(limit int) *Query {
	if !q.hasLimit {
		q.Write(" LIMIT ")
		q.Write(fmt.Sprintf("%d", limit))
		q.hasLimit = true
	}
	return q
}

// Offset adds an OFFSET clause
func (q *Query) Offset(offset int) *Query {
	if !q.hasOffset {
		q.Write(" OFFSET ")
		q.Write(fmt.Sprintf("%d", offset))
		q.hasOffset = true
	}
	return q
}

// Insert starts an INSERT statement
func (q *Query) Insert(table string, columns ...string) *Query {
	q.Write("INSERT INTO ")
	q.Write(table)
	q.Write(" (")
	q.Write(strings.Join(columns, ", "))
	q.Write(") VALUES ")
	return q
}

// Values adds a VALUES clause
func (q *Query) Values(count int) *Query {
	q.Write("(")
	q.Placeholders(count)
	q.Write(")")
	return q
}

// Update starts an UPDATE statement
func (q *Query) Update(table string) *Query {
	q.Write("UPDATE ")
	q.Write(table)
	q.Write(" SET ")
	return q
}

// Set adds a SET clause for updates
func (q *Query) Set(column string) *Query {
	if strings.Contains(q.builder.String(), "SET ") &&
		!strings.HasSuffix(q.builder.String(), "SET ") {
		q.Write(", ")
	}
	q.Write(column)
	q.Write(" = ")
	return q
}

// Delete starts a DELETE statement
func (q *Query) Delete() *Query {
	q.Write("DELETE")
	return q
}

// StartGroup starts a parenthetical group
func (q *Query) StartGroup() *Query {
	q.Write(" (")
	q.parensDepth++
	return q
}

// EndGroup ends a parenthetical group
func (q *Query) EndGroup() *Query {
	if q.parensDepth > 0 {
		q.Write(")")
		q.parensDepth--
	}
	return q
}

// Write adds a string to the query
func (q *Query) Write(s string) *Query {
	q.builder.WriteString(s)
	return q
}

// Placeholder adds a placeholder for a single argument
func (q *Query) Placeholder(arg interface{}) *Query {
	q.pos++
	q.args = append(q.args, arg)

	if q.dbType == DBTypePostgres {
		q.builder.WriteString(fmt.Sprintf("$%d", q.pos))
	} else {
		q.builder.WriteString("?")
	}

	return q
}

// Placeholders adds n placeholders separated by commas
func (q *Query) Placeholders(n int) *Query {
	placeholders := make([]string, n)

	for i := 0; i < n; i++ {
		q.pos++
		if q.dbType == DBTypePostgres {
			placeholders[i] = fmt.Sprintf("$%d", q.pos)
		} else {
			placeholders[i] = "?"
		}
	}

	q.builder.WriteString(strings.Join(placeholders, ", "))
	return q
}

// AddArgs adds arguments to the query
func (q *Query) AddArgs(args ...interface{}) *Query {
	q.args = append(q.args, args...)
	return q
}

// String returns the formatted query string
func (q *Query) String() string {
	return q.builder.String()
}

// Args returns the query arguments
func (q *Query) Args() []interface{} {
	return q.args
}