Skip to content

Commit dc9b415

Browse files
mafredrimafredri
authored
feat: Generate DB unique constraints as enums (#3701)
* feat: Generate DB unique constraints as enums This fixes a TODO from #3409.
1 parent f4c5020 commit dc9b415

File tree

7 files changed

+155
-15
lines changed

7 files changed

+155
-15
lines changed

‎Makefile

+3-3
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,11 @@ build: site/out/index.html $(shell find . -not -path './vendor/*' -type f -name
5656
.PHONY: build
5757

5858
# Runs migrations to output a dump of the database.
59-
coderd/database/dump.sql: coderd/database/dump/main.go $(wildcard coderd/database/migrations/*.sql)
60-
go run coderd/database/dump/main.go
59+
coderd/database/dump.sql: coderd/database/gen/dump/main.go $(wildcard coderd/database/migrations/*.sql)
60+
go run coderd/database/gen/dump/main.go
6161

6262
# Generates Go code for querying the database.
63-
coderd/database/querier.go: coderd/database/sqlc.yaml coderd/database/dump.sql $(wildcard coderd/database/queries/*.sql)
63+
coderd/database/querier.go: coderd/database/sqlc.yaml coderd/database/dump.sql $(wildcard coderd/database/queries/*.sql) coderd/database/gen/enum/main.go
6464
coderd/database/generate.sh
6565

6666
fmt/prettier:

‎coderd/database/errors.go

-9
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,6 @@ import (
66
"github.com/lib/pq"
77
)
88

9-
// UniqueConstraint represents a named unique constraint on a table.
10-
type UniqueConstraint string
11-
12-
// UniqueConstraint enums.
13-
// TODO(mafredri): Generate these from the database schema.
14-
const (
15-
UniqueWorkspacesOwnerIDLowerIdx UniqueConstraint = "workspaces_owner_id_lower_idx"
16-
)
17-
189
// IsUniqueViolation checks if the error is due to a unique violation.
1910
// If one or more specific unique constraints are given as arguments,
2011
// the error must be caused by one of them. If no constraints are given,

‎coderd/database/dump/main.go renamed to ‎coderd/database/gen/dump/main.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ func main() {
8888
if !ok {
8989
panic("couldn't get caller path")
9090
}
91-
err = os.WriteFile(filepath.Join(mainPath, "..", "..", "dump.sql"), []byte(dump), 0600)
91+
err = os.WriteFile(filepath.Join(mainPath, "..", "..", "..", "dump.sql"), []byte(dump), 0o600)
9292
if err != nil {
9393
panic(err)
9494
}

‎coderd/database/gen/enum/main.go

+120
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
package main
2+
3+
import (
4+
"bufio"
5+
"bytes"
6+
"fmt"
7+
"os"
8+
"os/exec"
9+
"strings"
10+
11+
"golang.org/x/xerrors"
12+
)
13+
14+
const header = `// Code generated by gen/enum. DO NOT EDIT.
15+
package database
16+
`
17+
18+
func main() {
19+
if err := run(); err != nil {
20+
panic(err)
21+
}
22+
}
23+
24+
func run() error {
25+
dump, err := os.Open("dump.sql")
26+
if err != nil {
27+
_, _ = fmt.Fprintf(os.Stderr, "error: %s must be run in the database directory with dump.sql present\n", os.Args[0])
28+
return err
29+
}
30+
defer dump.Close()
31+
32+
var uniqueConstraints []string
33+
34+
s := bufio.NewScanner(dump)
35+
query := ""
36+
for s.Scan() {
37+
line := strings.TrimSpace(s.Text())
38+
switch {
39+
case strings.HasPrefix(line, "--"):
40+
case line == "":
41+
case strings.HasSuffix(line, ";"):
42+
query += line
43+
if isUniqueConstraint(query) {
44+
uniqueConstraints = append(uniqueConstraints, query)
45+
}
46+
query = ""
47+
default:
48+
query += line + " "
49+
}
50+
}
51+
if err = s.Err(); err != nil {
52+
return err
53+
}
54+
55+
return writeContents("unique_constraint.go", uniqueConstraints, generateUniqueConstraints)
56+
}
57+
58+
func isUniqueConstraint(query string) bool {
59+
return strings.Contains(query, "UNIQUE")
60+
}
61+
62+
func generateUniqueConstraints(queries []string) ([]byte, error) {
63+
s := &bytes.Buffer{}
64+
65+
_, _ = fmt.Fprint(s, header)
66+
_, _ = fmt.Fprint(s, `
67+
// UniqueConstraint represents a named unique constraint on a table.
68+
type UniqueConstraint string
69+
70+
// UniqueConstraint enums.
71+
const (
72+
`)
73+
for _, query := range queries {
74+
name := ""
75+
switch {
76+
case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
77+
name = strings.Split(query, " ")[6]
78+
case strings.Contains(query, "CREATE UNIQUE INDEX"):
79+
name = strings.Split(query, " ")[3]
80+
default:
81+
return nil, xerrors.Errorf("unknown unique constraint format: %s", query)
82+
}
83+
_, _ = fmt.Fprintf(s, "\tUnique%s UniqueConstraint = %q // %s\n", nameFromSnakeCase(name), name, query)
84+
}
85+
_, _ = fmt.Fprint(s, ")\n")
86+
87+
return s.Bytes(), nil
88+
}
89+
90+
func writeContents[T any](dest string, arg T, fn func(T) ([]byte, error)) error {
91+
b, err := fn(arg)
92+
if err != nil {
93+
return err
94+
}
95+
err = os.WriteFile(dest, b, 0o600)
96+
if err != nil {
97+
return err
98+
}
99+
cmd := exec.Command("goimports", "-w", dest)
100+
return cmd.Run()
101+
}
102+
103+
func nameFromSnakeCase(s string) string {
104+
var ret string
105+
for _, ss := range strings.Split(s, "_") {
106+
switch ss {
107+
case "id":
108+
ret += "ID"
109+
case "ids":
110+
ret += "IDs"
111+
case "jwt":
112+
ret += "JWT"
113+
case "idx":
114+
ret += "Index"
115+
default:
116+
ret += strings.Title(ss)
117+
}
118+
}
119+
return ret
120+
}

‎coderd/database/generate.sh

+4-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")
1414
cd "$SCRIPT_DIR"
1515

1616
# Dump the updated schema.
17-
go run dump/main.go
17+
go run gen/dump/main.go
1818
# The logic below depends on the exact version being correct :(
1919
go run github.com/kyleconroy/sqlc/cmd/sqlc@v1.13.0 generate
2020

@@ -49,4 +49,7 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")
4949
# suggestions.
5050
go mod download
5151
goimports -w queries.sql.go
52+
53+
# Generate enums (e.g. unique constraints).
54+
go run gen/enum/main.go
5255
)

‎coderd/database/unique_constraint.go

+26
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/workspaces.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -512,7 +512,7 @@ func (api *API) patchWorkspace(rw http.ResponseWriter, r *http.Request) {
512512
return
513513
}
514514
// Check if the name was already in use.
515-
if database.IsUniqueViolation(err, database.UniqueWorkspacesOwnerIDLowerIdx) {
515+
if database.IsUniqueViolation(err, database.UniqueWorkspacesOwnerIDLowerIndex) {
516516
httpapi.Write(rw, http.StatusConflict, codersdk.Response{
517517
Message: fmt.Sprintf("Workspace %q already exists.", req.Name),
518518
Validations: []codersdk.ValidationError{{

0 commit comments

Comments
 (0)