github.com/octohelm/storage@v0.0.0-20240516030302-1ac2cc1ea347/pkg/sqlbuilder/stmt_with.go (about)

     1  package sqlbuilder
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  )
     7  
     8  type BuildSubQuery func(table Table) SqlExpr
     9  
    10  func WithRecursive(t Table, build BuildSubQuery) *WithStmt {
    11  	return With(t, build, "RECURSIVE")
    12  }
    13  
    14  func With(t Table, build BuildSubQuery, modifiers ...string) *WithStmt {
    15  	return (&WithStmt{modifiers: modifiers}).With(t, build)
    16  }
    17  
    18  type WithStmt struct {
    19  	modifiers []string
    20  	tables    []Table
    21  	asList    []BuildSubQuery
    22  	statement func(tables ...Table) SqlExpr
    23  }
    24  
    25  func (w *WithStmt) IsNil() bool {
    26  	return w == nil || len(w.tables) == 0 || len(w.asList) == 0 || w.statement == nil
    27  }
    28  
    29  func (w WithStmt) With(t Table, build BuildSubQuery) *WithStmt {
    30  	w.tables = append(w.tables, t)
    31  	w.asList = append(w.asList, build)
    32  	return &w
    33  }
    34  
    35  func (w WithStmt) Exec(statement func(tables ...Table) SqlExpr) *WithStmt {
    36  	w.statement = statement
    37  	return &w
    38  }
    39  
    40  func (w *WithStmt) Ex(ctx context.Context) *Ex {
    41  	e := Expr("WITH ")
    42  
    43  	if len(w.modifiers) > 0 {
    44  		e.WriteQuery(strings.Join(w.modifiers, " "))
    45  		e.WriteQuery(" ")
    46  	}
    47  
    48  	for i := range w.tables {
    49  		if i > 0 {
    50  			e.WriteQuery(", ")
    51  		}
    52  
    53  		table := w.tables[i]
    54  
    55  		e.WriteExpr(table)
    56  		e.WriteGroup(func(e *Ex) {
    57  			e.WriteExpr(table.Cols())
    58  		})
    59  
    60  		e.WriteQuery(" AS ")
    61  
    62  		build := w.asList[i]
    63  
    64  		e.WriteGroup(func(e *Ex) {
    65  			e.WriteQueryByte('\n')
    66  			e.WriteExpr(build(table))
    67  			e.WriteQueryByte('\n')
    68  		})
    69  	}
    70  
    71  	e.WriteQueryByte('\n')
    72  	e.WriteExpr(w.statement(w.tables...))
    73  
    74  	return e.Ex(ctx)
    75  }