From 29b4d12922878f1e3fa4b145846b8c96a7585e98 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sun, 14 Jun 2026 22:53:03 +0000 Subject: [PATCH 1/2] feat(rain): implement First, TxOptions, and relation column selection Co-authored-by: cungminh2710 <8063319+cungminh2710@users.noreply.github.com> --- pkg/rain/prepared_select.go | 6 ++++++ pkg/rain/query_select.go | 7 +++++++ pkg/rain/rain.go | 14 ++++++++++++-- pkg/rain/relation_loading.go | 28 ++++++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 2 deletions(-) diff --git a/pkg/rain/prepared_select.go b/pkg/rain/prepared_select.go index 1c7dcf6..cd2a149 100644 --- a/pkg/rain/prepared_select.go +++ b/pkg/rain/prepared_select.go @@ -87,6 +87,12 @@ func (q *SelectQuery) Prepare(ctx context.Context) (*PreparedSelectQuery, error) return prepared, nil } +// First executes the prepared SELECT query and scans the first row into dest. +// Returns sql.ErrNoRows if no result is found. +func (p *PreparedSelectQuery) First(ctx context.Context, args PreparedArgs, dest any) error { + return p.Scan(ctx, args, dest) +} + // Scan executes the prepared SELECT query and scans results into dest. func (p *PreparedSelectQuery) Scan(ctx context.Context, args PreparedArgs, dest any) error { bound, err := p.selectQuery.bind(args) diff --git a/pkg/rain/query_select.go b/pkg/rain/query_select.go index 0318ccf..8072645 100644 --- a/pkg/rain/query_select.go +++ b/pkg/rain/query_select.go @@ -184,6 +184,7 @@ func (q *SelectQuery) WithRelations(names ...string) *SelectQuery { type RelationConfig struct { Where schema.Predicate OrderBy []schema.OrderExpr + Columns []schema.Expression } // Relation configures filters and ordering for a named relation. @@ -629,6 +630,12 @@ func (q *SelectQuery) writeJoins(ctx *compileContext) error { return nil } +// First executes the SELECT query with an implicit LIMIT 1 and scans the first row into dest. +// Returns sql.ErrNoRows if no result is found. +func (q *SelectQuery) First(ctx context.Context, dest any) error { + return q.clone().Limit(1).Scan(ctx, dest) +} + // Scan executes the SELECT query and scans results into dest. func (q *SelectQuery) Scan(ctx context.Context, dest any) error { if q.runner == nil { diff --git a/pkg/rain/rain.go b/pkg/rain/rain.go index 93db036..2535093 100644 --- a/pkg/rain/rain.go +++ b/pkg/rain/rain.go @@ -279,12 +279,17 @@ func (db *DB) QueryRow(ctx context.Context, query string, args ...any) *sql.Row // Begin starts a new transaction. func (db *DB) Begin(ctx context.Context) (*Tx, error) { + return db.BeginTx(ctx, nil) +} + +// BeginTx starts a new transaction with the provided options. +func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*Tx, error) { primary := db.primaryHandle() if primary.db == nil { return nil, ErrNoConnection } - tx, err := primary.db.BeginTx(ctx, nil) + tx, err := primary.db.BeginTx(ctx, opts) if err != nil { return nil, err } @@ -294,7 +299,12 @@ func (db *DB) Begin(ctx context.Context) (*Tx, error) { // RunInTx executes fn in a transaction, rolling back on error and committing on success. func (db *DB) RunInTx(ctx context.Context, fn func(*Tx) error) error { - tx, err := db.Begin(ctx) + return db.RunInTxOpts(ctx, nil, fn) +} + +// RunInTxOpts executes fn in a transaction with the provided options, rolling back on error and committing on success. +func (db *DB) RunInTxOpts(ctx context.Context, opts *sql.TxOptions, fn func(*Tx) error) error { + tx, err := db.BeginTx(ctx, opts) if err != nil { return err } diff --git a/pkg/rain/relation_loading.go b/pkg/rain/relation_loading.go index 0ab7978..17944bf 100644 --- a/pkg/rain/relation_loading.go +++ b/pkg/rain/relation_loading.go @@ -295,6 +295,20 @@ func (q *SelectQuery) loadRelatedManyToManyRows( batchDest := reflect.New(reflect.SliceOf(relatedElemType)) targetQuery := &SelectQuery{runner: q.runner, dialect: q.dialect, table: tableDefSource{table: relation.TargetTable}} + if len(config.Columns) > 0 { + targetQuery.Column(config.Columns...) + // Ensure mapping target column is selected. + found := false + for _, col := range config.Columns { + if cr, ok := col.(schema.ColumnReference); ok && cr.ColumnDef() == relation.TargetColumn { + found = true + break + } + } + if !found { + targetQuery.Column(schema.Ref(relation.TargetColumn)) + } + } if config.Where != nil { targetQuery.Where(config.Where) } @@ -376,6 +390,20 @@ func (q *SelectQuery) loadRelatedRows( end := min(start+relationBatchSize, len(sourceKeys)) batchDest := reflect.New(reflect.SliceOf(relatedElemType)) query := &SelectQuery{runner: q.runner, dialect: q.dialect, table: tableDefSource{table: relation.TargetTable}} + if len(config.Columns) > 0 { + query.Column(config.Columns...) + // Ensure mapping target column is selected. + found := false + for _, col := range config.Columns { + if cr, ok := col.(schema.ColumnReference); ok && cr.ColumnDef() == relation.TargetColumn { + found = true + break + } + } + if !found { + query.Column(schema.Ref(relation.TargetColumn)) + } + } if config.Where != nil { query.Where(config.Where) } From a8cbcb14a970babcecea3afe848a1e39e2a817f7 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sun, 14 Jun 2026 23:54:44 +0000 Subject: [PATCH 2/2] feat(rain): implement First, TxOptions, and relation column selection - Add .First() to SelectQuery for ergonomic single-row fetching with implicit LIMIT 1. - Add BeginTx and RunInTxOpts to DB for sql.TxOptions support. - Add Columns selection to RelationConfig for restricted relational fetching. - Refactor relation loading to use ensureTargetColumnSelected helper. Co-authored-by: cungminh2710 <8063319+cungminh2710@users.noreply.github.com> --- pkg/rain/prepared_select.go | 6 ----- pkg/rain/relation_loading.go | 50 ++++++++++++++++-------------------- 2 files changed, 22 insertions(+), 34 deletions(-) diff --git a/pkg/rain/prepared_select.go b/pkg/rain/prepared_select.go index cd2a149..1c7dcf6 100644 --- a/pkg/rain/prepared_select.go +++ b/pkg/rain/prepared_select.go @@ -87,12 +87,6 @@ func (q *SelectQuery) Prepare(ctx context.Context) (*PreparedSelectQuery, error) return prepared, nil } -// First executes the prepared SELECT query and scans the first row into dest. -// Returns sql.ErrNoRows if no result is found. -func (p *PreparedSelectQuery) First(ctx context.Context, args PreparedArgs, dest any) error { - return p.Scan(ctx, args, dest) -} - // Scan executes the prepared SELECT query and scans results into dest. func (p *PreparedSelectQuery) Scan(ctx context.Context, args PreparedArgs, dest any) error { bound, err := p.selectQuery.bind(args) diff --git a/pkg/rain/relation_loading.go b/pkg/rain/relation_loading.go index 17944bf..d34d277 100644 --- a/pkg/rain/relation_loading.go +++ b/pkg/rain/relation_loading.go @@ -295,20 +295,7 @@ func (q *SelectQuery) loadRelatedManyToManyRows( batchDest := reflect.New(reflect.SliceOf(relatedElemType)) targetQuery := &SelectQuery{runner: q.runner, dialect: q.dialect, table: tableDefSource{table: relation.TargetTable}} - if len(config.Columns) > 0 { - targetQuery.Column(config.Columns...) - // Ensure mapping target column is selected. - found := false - for _, col := range config.Columns { - if cr, ok := col.(schema.ColumnReference); ok && cr.ColumnDef() == relation.TargetColumn { - found = true - break - } - } - if !found { - targetQuery.Column(schema.Ref(relation.TargetColumn)) - } - } + ensureTargetColumnSelected(targetQuery, config.Columns, relation.TargetColumn) if config.Where != nil { targetQuery.Where(config.Where) } @@ -390,20 +377,7 @@ func (q *SelectQuery) loadRelatedRows( end := min(start+relationBatchSize, len(sourceKeys)) batchDest := reflect.New(reflect.SliceOf(relatedElemType)) query := &SelectQuery{runner: q.runner, dialect: q.dialect, table: tableDefSource{table: relation.TargetTable}} - if len(config.Columns) > 0 { - query.Column(config.Columns...) - // Ensure mapping target column is selected. - found := false - for _, col := range config.Columns { - if cr, ok := col.(schema.ColumnReference); ok && cr.ColumnDef() == relation.TargetColumn { - found = true - break - } - } - if !found { - query.Column(schema.Ref(relation.TargetColumn)) - } - } + ensureTargetColumnSelected(query, config.Columns, relation.TargetColumn) if config.Where != nil { query.Where(config.Where) } @@ -572,6 +546,26 @@ func setRelationValue(parent reflect.Value, relationName string, relationType sc } } +func ensureTargetColumnSelected(query *SelectQuery, columns []schema.Expression, targetCol *schema.ColumnDef) { + if len(columns) == 0 { + return + } + + query.Column(columns...) + + // Ensure mapping target column is selected. + found := false + for _, col := range columns { + if cr, ok := col.(schema.ColumnReference); ok && cr.ColumnDef() == targetCol { + found = true + break + } + } + if !found { + query.Column(schema.Ref(targetCol)) + } +} + func dereferenceModelValue(value reflect.Value) reflect.Value { current := value for current.IsValid() && current.Kind() == reflect.Pointer {