Improve generated CountIntent.

1. Move count_intent to option.
2. Only SELECT query by default will have count intents.

Author: Stumble

internal/codegen/golang/option.go

@@ -7,13 +7,15 @@ import (
 )
 
 const (
-	WPgxOptionKeyCache      = "cache"
-	WPgxOptionKeyInvalidate = "invalidate"
+	WPgxOptionKeyCache       = "cache"
+	WPgxOptionKeyInvalidate  = "invalidate"
+	WpgxOptionKeyCountIntent = "count_intent"
 )
 
 type WPgxOption struct {
 	Cache       time.Duration
 	Invalidates []string
+	CountIntent bool
 }
 
 func parseOption(options map[string]string, queryNames map[string]bool) (rv WPgxOption, err error) {
@@ -24,7 +26,7 @@ func parseOption(options map[string]string, queryNames map[string]bool) (rv WPgx
 			if err != nil {
 				return
 			}
-			if rv.Cache < 1 * time.Millisecond {
+			if rv.Cache < 1*time.Millisecond {
 				return rv, fmt.Errorf("cache duration too short: %s", v)
 			}
 		case WPgxOptionKeyInvalidate:
@@ -37,6 +39,16 @@ func parseOption(options map[string]string, queryNames map[string]bool) (rv WPgx
 				}
 				rv.Invalidates = append(rv.Invalidates, queryName)
 			}
+		case WpgxOptionKeyCountIntent:
+			if v == "true" {
+				rv.CountIntent = true
+			} else if v == "false" {
+				rv.CountIntent = false
+			} else {
+				return rv, fmt.Errorf("Unknown count_intent value: %s", v)
+			}
+		default:
+			return rv, fmt.Errorf("Unknown option: %s", k)
 		}
 	}
 	return

internal/codegen/golang/query.go

@@ -320,6 +320,11 @@ func (q Query) TableIdentifierForMySQL() string {
 	return strings.Join(escapedNames, ".")
 }
 
+// CountIntent is used by WPgx only.
+func (q Query) CountIntent() bool {
+	return q.Option.CountIntent
+}
+
 // CacheKey is used by WPgx only.
 func (q Query) CacheKey() string {
 	return genCacheKeyWithArgName(q, q.Arg.Name)

internal/codegen/golang/templates/wpgx/queryCode.tmpl

@@ -36,7 +36,9 @@ type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
 {{range .Comments}}//{{.}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.InvalidateArgs}}) (*{{.Ret.Type}}, error) {
+{{- if .CountIntent }}
     q.db.CountIntent("{{.UniqueLabel}}")
+{{- end}}
 {{- if eq .Option.Cache.Milliseconds 0}}
     row := q.db.WQueryRow(ctx, "{{.UniqueLabel}}", {{.ConstantName}}, {{.Arg.Params}})
     var {{.Ret.Name}} *{{.Ret.Type}} = new({{.Ret.Type}})
@@ -99,7 +101,9 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.Invalida
 {{range .Comments}}//{{.}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.InvalidateArgs}}) ([]{{.Ret.Type}}, error) {
+{{- if .CountIntent }}
     q.db.CountIntent("{{.UniqueLabel}}")
+{{- end}}
 {{- if eq .Option.Cache.Milliseconds 0}}
     rows, err := q.db.WQuery(ctx, "{{.UniqueLabel}}", {{.ConstantName}}, {{.Arg.Params}})
     if err != nil {
@@ -179,7 +183,6 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.Invalida
 {{range .Comments}}//{{.}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.InvalidateArgs}}) error {
-    q.db.CountIntent("{{.UniqueLabel}}")
     _, err := q.db.WExec(ctx, "{{.UniqueLabel}}", {{.ConstantName}}, {{.Arg.Params}})
     if err != nil {
         return err
@@ -214,7 +217,6 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.Invalida
 {{range .Comments}}//{{.}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.InvalidateArgs}}) (int64, error) {
-    q.db.CountIntent("{{.UniqueLabel}}")
     result, err := q.db.WExec(ctx, "{{.UniqueLabel}}", {{.ConstantName}}, {{.Arg.Params}})
     if err != nil {
         return 0, err
@@ -249,7 +251,6 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.Invalida
 {{range .Comments}}//{{.}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}} {{.InvalidateArgs}}) (pgconn.CommandTag, error) {
-    q.db.CountIntent("{{.UniqueLabel}}")
     rv, err := q.db.WExec(ctx, "{{.UniqueLabel}}", {{.ConstantName}}, {{.Arg.Params}})
     if err != nil {
         return rv, err

internal/compiler/parse.go

@@ -15,6 +15,8 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/sql/astutils"
 	"github.com/sqlc-dev/sqlc/internal/sql/rewrite"
 	"github.com/sqlc-dev/sqlc/internal/sql/validate"
+
+	"github.com/sqlc-dev/sqlc/internal/codegen/golang"
 )
 
 var ErrUnsupportedStatementType = errors.New("parseQuery: unsupported statement type")
@@ -79,6 +81,11 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		raw.Stmt, queryConfig.Name, queryConfig.Cmd, queryConfig.Options); err != nil {
 		return nil, err
 	}
+	err = validateAndSetDefaultOptions(
+		raw.Stmt, queryConfig.Name, queryConfig.Cmd, queryConfig.Options)
+	if err != nil {
+		return nil, err
+	}
 	rvs := rangeVars(raw.Stmt)
 	refs, err := findParameters(raw.Stmt)
 	if err != nil {
@@ -183,3 +190,29 @@ func uniqueParamRefs(in []paramRef, dollar bool) []paramRef {
 	}
 	return o
 }
+
+// wicked-sqlc specific function
+func validateAndSetDefaultOptions(n ast.Node, name, cmd string, options map[string]string) error {
+	countIntent, hasCountIntentDefaultVal := options[golang.WpgxOptionKeyCountIntent]
+	if !(cmd == metadata.CmdMany || cmd == metadata.CmdOne) {
+		if countIntent == "true" {
+			return fmt.Errorf(
+				"query %q uses count_intent option but cmd is neither 'one' or 'many'", name)
+		}
+		return nil
+	}
+	_, isSelect := n.(*ast.SelectStmt)
+	if !hasCountIntentDefaultVal {
+		if isSelect {
+			options[golang.WpgxOptionKeyCountIntent] = "true"
+		} else {
+			options[golang.WpgxOptionKeyCountIntent] = "false"
+		}
+	}
+	if !isSelect {
+		if _, ok := options[golang.WPgxOptionKeyCache]; ok {
+			return fmt.Errorf("query %q uses cache option but is not a SELECT", name)
+		}
+	}
+	return nil
+}

internal/sql/validate/cmd.go

@@ -6,8 +6,6 @@ import (
 
 	"github.com/sqlc-dev/sqlc/internal/metadata"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
-
-	"github.com/sqlc-dev/sqlc/internal/codegen/golang"
 )
 
 func validateCopyfrom(n ast.Node) error {
@@ -80,9 +78,5 @@ func Cmd(n ast.Node, name, cmd string, options map[string]string) error {
 	if list == nil || len(list.Items) == 0 {
 		return fmt.Errorf("query %q specifies parameter %q without containing a RETURNING clause", name, cmd)
 	}
-	// XXX(yumin): dirty hack here to use the wpgx options here.
-	if _, ok := options[golang.WPgxOptionKeyCache]; ok {
-		return fmt.Errorf("query %q uses cache option but is not a SELECT", name)
-	}
 	return nil
 }