coder · johnstcn · Sep 13, 2023 · Sep 13, 2023 · Sep 13, 2023 · johnstcn
diff --git a/coderd/apikey.go b/coderd/apikey.go
@@ -91,7 +91,7 @@ func (api *API) postToken(rw http.ResponseWriter, r *http.Request) {
 		TokenName:        tokenName,
 	})
 	if err != nil {
-		if database.IsUniqueViolation(err, database.UniqueIndexApiKeyName) {
+		if database.IsUniqueViolation(err, database.UniqueIndexAPIKeyName) {
 			httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
 				Message: fmt.Sprintf("A token with name %q already exists.", tokenName),
 				Validations: []codersdk.ValidationError{{

diff --git a/coderd/database/errors.go b/coderd/database/errors.go
@@ -37,6 +37,28 @@ func IsUniqueViolation(err error, uniqueConstraints ...UniqueConstraint) bool {
 	return false
 }
 
+// IsForeignKeyViolation checks if the error is due to a foreign key violation.
+// If one or more specific foreign key constraints are given as arguments,
+// the error must be caused by one of them. If no constraints are given,
+// this function returns true for any foreign key violation.
+func IsForeignKeyViolation(err error, foreignKeyConstraints ...ForeignKeyConstraint) bool {
+	var pqErr *pq.Error
+	if errors.As(err, &pqErr) {
+		if pqErr.Code.Name() == "foreign_key_violation" {
+			if len(foreignKeyConstraints) == 0 {
+				return true
+			}
+			for _, fc := range foreignKeyConstraints {
+				if pqErr.Constraint == string(fc) {
+					return true
+				}
+			}
+		}
+	}
+
+	return false
+}
+
 // IsQueryCanceledError checks if the error is due to a query being canceled.
 func IsQueryCanceledError(err error) bool {
 	var pqErr *pq.Error

diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go
diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go
diff --git a/scripts/dbgen/main.go b/scripts/dbgen/main.go
@@ -84,6 +84,11 @@ return %s
 		return xerrors.Errorf("generate unique constraints: %w", err)
 	}
 
+	err = generateForeignKeyConstraints()
+	if err != nil {
+		return xerrors.Errorf("generate foreign key constraints: %w", err)
+	}
+
 	return nil
 }
 
@@ -125,7 +130,7 @@ func generateUniqueConstraints() error {
 
 	s := &bytes.Buffer{}
 
-	_, _ = fmt.Fprint(s, `// Code generated by gen/enum. DO NOT EDIT.
+	_, _ = fmt.Fprint(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT.
 package database
 `)
 	_, _ = fmt.Fprint(s, `
@@ -160,6 +165,78 @@ const (
 	return os.WriteFile(outputPath, data, 0o600)
 }
 
+// generateForeignKeyConstraints generates the ForeignKeyConstraint enum.
+func generateForeignKeyConstraints() error {
+	localPath, err := localFilePath()
+	if err != nil {
+		return err
+	}
+	databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database")
+
+	dump, err := os.Open(filepath.Join(databasePath, "dump.sql"))
+	if err != nil {
+		return err
+	}
+	defer dump.Close()
+
+	var foreignKeyConstraints []string
+	dumpScanner := bufio.NewScanner(dump)
+	query := ""
+	for dumpScanner.Scan() {
+		line := strings.TrimSpace(dumpScanner.Text())
+		switch {
+		case strings.HasPrefix(line, "--"):
+		case line == "":
+		case strings.HasSuffix(line, ";"):
+			query += line
+			if strings.Contains(query, "FOREIGN KEY") {
+				foreignKeyConstraints = append(foreignKeyConstraints, query)
+			}
+			query = ""
+		default:
+			query += line + " "
+		}
+	}
+
+	if err := dumpScanner.Err(); err != nil {
+		return err
+	}
+
+	s := &bytes.Buffer{}
+
+	_, _ = fmt.Fprint(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT.
+package database
+`)
+	_, _ = fmt.Fprint(s, `
+// ForeignKeyConstraint represents a named foreign key constraint on a table.
+type ForeignKeyConstraint string
+
+// ForeignKeyConstraint enums.
+const (
+`)
+	for _, query := range foreignKeyConstraints {
+		name := ""
+		switch {
+		case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
+			name = strings.Split(query, " ")[6]
+		default:
+			return xerrors.Errorf("unknown foreign key constraint format: %s", query)
+		}
+		_, _ = fmt.Fprintf(s, "\tForeignKey%s ForeignKeyConstraint = %q // %s\n", nameFromSnakeCase(name), name, query)
+	}
+	_, _ = fmt.Fprint(s, ")\n")
+
+	outputPath := filepath.Join(databasePath, "foreign_key_constraint.go")
+
+	data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{
+		Comments: true,
+	})
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(outputPath, data, 0o600)
+}
+
 type stubParams struct {
 	FuncName   string
 	Parameters string
@@ -560,6 +637,14 @@ func nameFromSnakeCase(s string) string {
 			ret += "JWT"
 		case "idx":
 			ret += "Index"
+		case "api":
+			ret += "API"
+		case "uuid":
+			ret += "UUID"
+		case "gitsshkeys":
+			ret += "GitSSHKeys"
+		case "fkey":
+			// ignore
 		default:
 			ret += strings.Title(ss)
 		}