private/dbutil/pgutil: faster cockroach constraint finding

Change-Id: Ia100b9ef7d2d59dfad0389feb8f2e7c47c2c4c9b

private/dbutil/cockroachutil/db.go

@@ -5,8 +5,8 @@ package cockroachutil
 
 import (
 	"context"
+	"crypto/rand"
 	"encoding/hex"
-	"math/rand"
 	"net/url"
 	"strings"
 
@@ -23,8 +23,6 @@ var mon = monkit.Package()
 // CreateRandomTestingSchemaName creates a random schema name string.
 func CreateRandomTestingSchemaName(n int) string {
 	data := make([]byte, n)
-
-	// math/rand.Read() always returns a nil error so there's no need to handle the error.
 	_, _ = rand.Read(data)
 	return hex.EncodeToString(data)
 }

private/dbutil/dbschema/schema.go

@@ -14,6 +14,9 @@ import (
 type Queryer interface {
 	// QueryContext executes a query that returns rows, typically a SELECT.
 	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
+
+	// QueryRowContext executes a query that returns a single row.
+	QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
 }
 
 // Schema is the database structure.

private/dbutil/pgutil/query.go

@@ -19,6 +19,13 @@ import (
 func QuerySchema(ctx context.Context, db dbschema.Queryer) (*dbschema.Schema, error) {
 	schema := &dbschema.Schema{}
 
+	// get version string to do efficient queries
+	var version string
+	row := db.QueryRowContext(ctx, `SELECT version()`)
+	if err := row.Scan(&version); err != nil {
+		return nil, errs.Wrap(err)
+	}
+
 	// find tables
 	err := func() error {
 		rows, err := db.QueryContext(ctx, `
@@ -54,6 +61,14 @@ func QuerySchema(ctx context.Context, db dbschema.Queryer) (*dbschema.Schema, er
 
 	// find constraints
 	err = func() error {
+		// cockroach has a .condef field and it's way faster than the function call
+		var definitionClause string
+		if strings.Contains(version, "CockroachDB") {
+			definitionClause = `pg_constraint.condef AS definition`
+		} else {
+			definitionClause = `pg_get_constraintdef(pg_constraint.oid) AS definition`
+		}
+
 		rows, err := db.QueryContext(ctx, `
 			SELECT
 				pg_class.relname      AS table_name,
@@ -67,8 +82,7 @@ func QuerySchema(ctx context.Context, db dbschema.Queryer) (*dbschema.Schema, er
 						JOIN UNNEST(pg_constraint.conkey) WITH ORDINALITY AS u(attnum, pos) ON u.attnum = pg_attribute.attnum
 					WHERE
 						pg_attribute.attrelid = pg_class.oid
-				) AS columns,
-				pg_get_constraintdef(pg_constraint.oid) AS definition
+				) AS columns, `+definitionClause+`
 			FROM
 				pg_constraint
 				JOIN pg_class ON pg_class.oid = pg_constraint.conrelid