feat: Generate DB unique constraints as enums

mafredri · mafredri · commit d32e9e8c9751 · 2022-08-26T12:36:38.000+03:00
This fixes a TODO from #3409.
diff --git a/coderd/database/errors.go b/coderd/database/errors.go
@@ -6,15 +6,6 @@ import (
 	"github.com/lib/pq"
 )
 
-// UniqueConstraint represents a named unique constraint on a table.
-type UniqueConstraint string
-
-// UniqueConstraint enums.
-// TODO(mafredri): Generate these from the database schema.
-const (
-	UniqueWorkspacesOwnerIDLowerIdx UniqueConstraint = "workspaces_owner_id_lower_idx"
-)
-
 // IsUniqueViolation checks if the error is due to a unique violation.
 // If one or more specific unique constraints are given as arguments,
 // the error must be caused by one of them. If no constraints are given,
diff --git a/coderd/database/gen/enum/main.go b/coderd/database/gen/enum/main.go
@@ -0,0 +1,120 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+
+	"golang.org/x/xerrors"
+)
+
+const header = `// Code generated by gen/enum. DO NOT EDIT.
+package database
+`
+
+func main() {
+	if err := run(); err != nil {
+		panic(err)
+	}
+}
+
+func run() error {
+	dump, err := os.Open("dump.sql")
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "error: %s must be run in the database directory with dump.sql present\n", os.Args[0])
+		return err
+	}
+	defer dump.Close()
+
+	var uniqueConstraints []string
+
+	s := bufio.NewScanner(dump)
+	query := ""
+	for s.Scan() {
+		line := s.Text()
+		switch {
+		case strings.HasPrefix(line, "--"):
+		case line == "":
+		case strings.HasSuffix(line, ";"):
+			query += strings.TrimSpace(line)
+			if isUniqueConstraint(query) {
+				uniqueConstraints = append(uniqueConstraints, query)
+			}
+			query = ""
+		default:
+			query += strings.TrimSpace(line) + " "
+		}
+	}
+	if err = s.Err(); err != nil {
+		return err
+	}
+
+	return writeContents("unique_constraint.go", uniqueConstraints, generateUniqueConstraints)
+}
+
+func isUniqueConstraint(query string) bool {
+	return strings.Contains(query, "UNIQUE")
+}
+
+func generateUniqueConstraints(queries []string) ([]byte, error) {
+	s := &bytes.Buffer{}
+
+	_, _ = fmt.Fprint(s, header)
+	_, _ = fmt.Fprint(s, `
+// UniqueConstraint represents a named unique constraint on a table.
+type UniqueConstraint string
+
+// UniqueConstraint enums.
+const (
+`)
+	for _, query := range queries {
+		name := ""
+		switch {
+		case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
+			name = strings.Split(query, " ")[6]
+		case strings.Contains(query, "CREATE UNIQUE INDEX"):
+			name = strings.Split(query, " ")[3]
+		default:
+			return nil, xerrors.Errorf("unknown unique constraint format: %s", query)
+		}
+		_, _ = fmt.Fprintf(s, "\tUnique%s UniqueConstraint = %q // %s\n", nameFromSnakeCase(name), name, query)
+	}
+	_, _ = fmt.Fprint(s, ")\n")
+
+	return s.Bytes(), nil
+}
+
+func writeContents[T any](dest string, arg T, fn func(T) ([]byte, error)) error {
+	b, err := fn(arg)
+	if err != nil {
+		return err
+	}
+	err = os.WriteFile(dest, b, 0o600)
+	if err != nil {
+		return err
+	}
+	cmd := exec.Command("goimports", "-w", dest)
+	return cmd.Run()
+}
+
+func nameFromSnakeCase(s string) string {
+	var ret string
+	for _, ss := range strings.Split(s, "_") {
+		switch ss {
+		case "id":
+			ret += "ID"
+		case "ids":
+			ret += "IDs"
+		case "jwt":
+			ret += "JWT"
+		case "idx":
+			ret += "Index"
+		default:
+			ret += strings.Title(ss)
+		}
+	}
+	return ret
+}
diff --git a/coderd/database/generate.sh b/coderd/database/generate.sh
@@ -49,4 +49,7 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")
 	# suggestions.
 	go mod download
 	goimports -w queries.sql.go
+
+	# Generate enums (e.g. unique constraints).
+	go run gen/enum/main.go
 )
diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go

Original file line number	Diff line number	Diff line change
`@@ -49,4 +49,7 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")`
`49`	`49`	`# suggestions.`
`50`	`50`	`go mod download`
`51`	`51`	`goimports -w queries.sql.go`
	`52`	`+`
	`53`	`+ # Generate enums (e.g. unique constraints).`
	`54`	`+ go run gen/enum/main.go`
`52`	`55`	`)`