Auto-inject LIMIT 1 in bob.One

CHANGELOG.md

@@ -10,13 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Added PostgreSQL `MERGE` statement SQL parser used for sql-to-code generation (thanks @atzedus)
+- Added `bob.First`. It is identical to how `bob.One` used to work, as it does not inject a `LIMIT 1` clause. (thanks @jacobmolby)
 - Added `As(alias)` method to `bob.BaseQuery`, allowing queries (e.g. `mysql.Select(...)`) to be aliased directly when used as subqueries in a column list. (thanks @jacobmolby)
 
 ### Changed
 
 - **BREAKING:** Renamed the generated `ThenLoadCount` variable to `SelectThenLoadCount` for naming consistency with `SelectThenLoad`/`InsertThenLoad`/`UpdateThenLoad`. Update call sites from `models.ThenLoadCount.X.Y` to `models.SelectThenLoadCount.X.Y`. (thanks @jacobmolby)
 - **BREAKING:** `mm.SetCol()` now returns `mods.Set[*mm.UpdateAction]` instead of a custom `SetChain` type. `.ToExpr(val)` is replaced by `.To(val)`, and `.ToDefault()` is replaced by `.To(psql.Raw("DEFAULT"))`. `.To()` and `.ToArg()` work as before. (thanks @atzedus)
 - **BREAKING:** `mm.Recursive()` has been removed. PostgreSQL does not support `WITH RECURSIVE` in MERGE statements. (thanks @atzedus)
+- `bob.One` (and the generated `.One()` methods that delegate to it) now automatically adds `LIMIT 1` to SELECT queries when no row-limit is explicitly set. (thanks @jacobmolby)
 
 ### Fixed
 

dialect/mysql/dialect/select.go

@@ -41,6 +41,19 @@ func (s *SelectQuery) SetInto(i any) {
 	s.into = i
 }
 
+// SetLimitIfUnset sets a default LIMIT only when no limit is currently configured.
+func (s *SelectQuery) SetLimitIfUnset(limit any) {
+	if len(s.Combines.Queries) > 0 {
+		if s.CombinedLimit.Count == nil {
+			s.CombinedLimit.SetLimit(limit)
+		}
+		return
+	}
+	if s.Limit.Count == nil {
+		s.Limit.SetLimit(limit)
+	}
+}
+
 func (s SelectQuery) WriteSQL(ctx context.Context, w io.StringWriter, d bob.Dialect, start int) ([]any, error) {
 	var args []any
 	var err error

dialect/mysql/select_limit_test.go

@@ -0,0 +1,66 @@
+package mysql_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/stephenafamo/bob"
+	"github.com/stephenafamo/bob/dialect/mysql"
+	"github.com/stephenafamo/bob/dialect/mysql/sm"
+)
+
+func TestSelectQuerySetLimitIfUnset(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("injects limit when unset", func(t *testing.T) {
+		q := mysql.Select(sm.Columns("id"), sm.From("users"))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if !strings.Contains(sql, "LIMIT 1") {
+			t.Fatalf("expected LIMIT 1 in SQL, got: %s", sql)
+		}
+	})
+
+	t.Run("preserves existing limit", func(t *testing.T) {
+		q := mysql.Select(sm.Columns("id"), sm.From("users"), sm.Limit(5))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if !strings.Contains(sql, "LIMIT 5") {
+			t.Fatalf("expected LIMIT 5 (preserved), got: %s", sql)
+		}
+		if strings.Contains(sql, "LIMIT 1") {
+			t.Fatalf("did not expect LIMIT 1, got: %s", sql)
+		}
+	})
+
+	t.Run("union sets CombinedLimit", func(t *testing.T) {
+		inner := mysql.Select(sm.Columns("id"), sm.From("orders"))
+		q := mysql.Select(sm.Columns("id"), sm.From("users"), sm.Union(inner))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		idx := strings.LastIndex(sql, "LIMIT 1")
+		unionIdx := strings.Index(sql, "UNION")
+		if idx == -1 {
+			t.Fatalf("expected LIMIT 1, got: %s", sql)
+		}
+		if unionIdx == -1 || idx < unionIdx {
+			t.Fatalf("expected LIMIT 1 after UNION, got: %s", sql)
+		}
+	})
+}

dialect/psql/dialect/select.go

@@ -36,6 +36,20 @@ type SelectQuery struct {
 	CombinedOffset clause.Offset
 }
 
+// SetLimitIfUnset sets a default LIMIT only when no limit is currently configured.
+// PostgreSQL accepts either LIMIT or FETCH FIRST, so both are inspected.
+func (s *SelectQuery) SetLimitIfUnset(limit any) {
+	if len(s.Combines.Queries) > 0 {
+		if s.CombinedLimit.Count == nil && s.CombinedFetch.Count == nil {
+			s.CombinedLimit.SetLimit(limit)
+		}
+		return
+	}
+	if s.Limit.Count == nil && s.Fetch.Count == nil {
+		s.Limit.SetLimit(limit)
+	}
+}
+
 func (s SelectQuery) WriteSQL(ctx context.Context, w io.StringWriter, d bob.Dialect, start int) ([]any, error) {
 	var err error
 	var args []any

dialect/psql/select_limit_test.go

@@ -0,0 +1,81 @@
+package psql_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/stephenafamo/bob"
+	"github.com/stephenafamo/bob/dialect/psql"
+	"github.com/stephenafamo/bob/dialect/psql/sm"
+)
+
+func TestSelectQuerySetLimitIfUnset(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("injects limit when unset", func(t *testing.T) {
+		q := psql.Select(sm.Columns("id"), sm.From("users"))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if !strings.Contains(sql, "LIMIT 1") {
+			t.Fatalf("expected LIMIT 1 in SQL, got: %s", sql)
+		}
+	})
+
+	t.Run("preserves existing LIMIT", func(t *testing.T) {
+		q := psql.Select(sm.Columns("id"), sm.From("users"), sm.Limit(5))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if !strings.Contains(sql, "LIMIT 5") {
+			t.Fatalf("expected LIMIT 5 (preserved), got: %s", sql)
+		}
+		if strings.Contains(sql, "LIMIT 1") {
+			t.Fatalf("did not expect LIMIT 1, got: %s", sql)
+		}
+	})
+
+	t.Run("preserves existing FETCH", func(t *testing.T) {
+		q := psql.Select(sm.Columns("id"), sm.From("users"), sm.Fetch(5, false))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if strings.Contains(sql, "LIMIT") {
+			t.Fatalf("did not expect LIMIT (FETCH already set), got: %s", sql)
+		}
+	})
+
+	t.Run("union sets CombinedLimit", func(t *testing.T) {
+		inner := psql.Select(sm.Columns("id"), sm.From("orders"))
+		q := psql.Select(sm.Columns("id"), sm.From("users"), sm.Union(inner))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		// LIMIT must appear after the UNION, not inside the first SELECT's parens.
+		idx := strings.LastIndex(sql, "LIMIT 1")
+		unionIdx := strings.Index(sql, "UNION")
+		if idx == -1 {
+			t.Fatalf("expected LIMIT 1, got: %s", sql)
+		}
+		if unionIdx == -1 || idx < unionIdx {
+			t.Fatalf("expected LIMIT 1 after UNION, got: %s", sql)
+		}
+	})
+}

dialect/sqlite/dialect/select.go

@@ -29,6 +29,13 @@ type SelectQuery struct {
 	bob.ContextualModdable[*SelectQuery]
 }
 
+// SetLimitIfUnset sets a default LIMIT only when no limit is currently configured.
+func (s *SelectQuery) SetLimitIfUnset(limit any) {
+	if s.Limit.Count == nil {
+		s.Limit.SetLimit(limit)
+	}
+}
+
 func (s SelectQuery) WriteSQL(ctx context.Context, w io.StringWriter, d bob.Dialect, start int) ([]any, error) {
 	var args []any
 	var err error

dialect/sqlite/select_limit_test.go

@@ -0,0 +1,46 @@
+package sqlite_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/stephenafamo/bob"
+	"github.com/stephenafamo/bob/dialect/sqlite"
+	"github.com/stephenafamo/bob/dialect/sqlite/sm"
+)
+
+func TestSelectQuerySetLimitIfUnset(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("injects limit when unset", func(t *testing.T) {
+		q := sqlite.Select(sm.Columns("id"), sm.From("users"))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if !strings.Contains(sql, "LIMIT 1") {
+			t.Fatalf("expected LIMIT 1 in SQL, got: %s", sql)
+		}
+	})
+
+	t.Run("preserves existing limit", func(t *testing.T) {
+		q := sqlite.Select(sm.Columns("id"), sm.From("users"), sm.Limit(5))
+		var l bob.Limiter = q
+		l.SetLimitIfUnset(1)
+
+		sql, _, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("build: %v", err)
+		}
+		if !strings.Contains(sql, "LIMIT 5") {
+			t.Fatalf("expected LIMIT 5 (preserved), got: %s", sql)
+		}
+		if strings.Contains(sql, "LIMIT 1") {
+			t.Fatalf("did not expect LIMIT 1, got: %s", sql)
+		}
+	})
+}

exec.go

@@ -24,6 +24,12 @@ type (
 	HookableType interface {
 		AfterQueryHook(context.Context, Executor, QueryType) error
 	}
+
+	// Limiter is implemented by queries that can have a row-output limit.
+	// The limit is only set if no limit is currently configured.
+	Limiter interface {
+		SetLimitIfUnset(limit any)
+	}
 )
 
 type Executor interface {
@@ -73,7 +79,19 @@ func Exec(ctx context.Context, exec Executor, q Query) (sql.Result, error) {
 	return result, nil
 }
 
+// One executes the query and scans a single row into T.
+//
+// If the query satisfies Limiter and no row-limit is currently set on it,
+// One adds LIMIT 1 to the underlying SQL.
 func One[T any](ctx context.Context, exec Executor, q Query, m scan.Mapper[T]) (T, error) {
+	if l, ok := q.(Limiter); ok {
+		l.SetLimitIfUnset(1)
+	}
+	return First(ctx, exec, q, m)
+}
+
+// First executes the query and scans the first row into T.
+func First[T any](ctx context.Context, exec Executor, q Query, m scan.Mapper[T]) (T, error) {
 	var t T
 	var err error
 

query.go

@@ -107,6 +107,12 @@ func (b BaseQuery[E]) GetLoaders() []Loader {
 	return nil
 }
 
+func (b BaseQuery[E]) SetLimitIfUnset(limit any) {
+	if l, ok := any(b.Expression).(Limiter); ok {
+		l.SetLimitIfUnset(limit)
+	}
+}
+
 func (b BaseQuery[E]) GetMapperMods() []scan.MapperMod {
 	if l, ok := any(b.Expression).(MapperModder); ok {
 		return l.GetMapperMods()