storj/private/dbutil/pgutil/query.go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.

package pgutil

import (
	"context"
	"fmt"
	"regexp"
	"strings"

	"github.com/lib/pq"
	"github.com/zeebo/errs"

	"storj.io/storj/private/dbutil/dbschema"
)

// QuerySchema loads the schema from postgres database.
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() (err error) {
		rows, err := db.QueryContext(ctx, `
			SELECT table_name, column_name, is_nullable, coalesce(column_default, ''), data_type
			FROM  information_schema.columns
			WHERE table_schema = CURRENT_SCHEMA
		`)
		if err != nil {
			return err
		}
		defer func() { err = errs.Combine(err, rows.Close()) }()

		for rows.Next() {
			var tableName, columnName, isNullable, columnDefault, dataType string
			err := rows.Scan(&tableName, &columnName, &isNullable, &columnDefault, &dataType)
			if err != nil {
				return err
			}

			table := schema.EnsureTable(tableName)
			table.AddColumn(&dbschema.Column{
				Name:       columnName,
				Type:       dataType,
				IsNullable: isNullable == "YES",
				Default:    parseColumnDefault(columnDefault),
			})
		}

		return rows.Err()
	}()
	if err != nil {
		return nil, err
	}

	// find constraints
	err = func() (err 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,
				pg_constraint.conname AS constraint_name,
				pg_constraint.contype AS constraint_type,
				(
					SELECT
						ARRAY_AGG(pg_attribute.attname ORDER BY u.pos)
					FROM
						pg_attribute
						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, `+definitionClause+`
			FROM
				pg_constraint
				JOIN pg_class ON pg_class.oid = pg_constraint.conrelid
				JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace
			WHERE pg_namespace.nspname = CURRENT_SCHEMA
		`)
		if err != nil {
			return err
		}
		defer func() { err = errs.Combine(err, rows.Close()) }()

		for rows.Next() {
			var tableName, constraintName, constraintType string
			var columns pq.StringArray
			var definition string

			err := rows.Scan(&tableName, &constraintName, &constraintType, &columns, &definition)
			if err != nil {
				return err
			}

			switch constraintType {
			case "p": // primary key
				table := schema.EnsureTable(tableName)
				table.PrimaryKey = ([]string)(columns)
			case "f": // foreign key
				if len(columns) != 1 {
					return fmt.Errorf("expected one column, got: %q", columns)
				}

				table := schema.EnsureTable(tableName)
				column, ok := table.FindColumn(columns[0])
				if !ok {
					return fmt.Errorf("did not find column %q", columns[0])
				}

				matches := rxPostgresForeignKey.FindStringSubmatch(definition)
				if len(matches) == 0 {
					return fmt.Errorf("unable to parse constraint %q", definition)
				}

				column.Reference = &dbschema.Reference{
					Table:    matches[1],
					Column:   matches[2],
					OnUpdate: matches[3],
					OnDelete: matches[4],
				}
			case "u": // unique
				table := schema.EnsureTable(tableName)
				table.Unique = append(table.Unique, columns)
			default:
				return fmt.Errorf("unhandled constraint type %q", constraintType)
			}
		}
		return rows.Err()
	}()
	if err != nil {
		return nil, err
	}

	// find indexes
	err = func() (err error) {
		rows, err := db.QueryContext(ctx, `SELECT indexdef FROM pg_indexes WHERE schemaname = CURRENT_SCHEMA`)
		if err != nil {
			return errs.Wrap(err)
		}
		defer func() { err = errs.Combine(err, rows.Close()) }()

		for rows.Next() {
			var indexdef string
			err := rows.Scan(&indexdef)
			if err != nil {
				return errs.Wrap(err)
			}

			index, err := parseIndexDefinition(indexdef)
			if err != nil {
				return errs.Wrap(err)
			}
			schema.Indexes = append(schema.Indexes, index)
		}

		return errs.Wrap(rows.Err())
	}()
	if err != nil {
		return nil, err
	}

	schema.Sort()
	return schema, nil
}

// matches FOREIGN KEY (project_id) REFERENCES projects(id) ON UPDATE CASCADE ON DELETE CASCADE
var rxPostgresForeignKey = regexp.MustCompile(
	`^FOREIGN KEY \([[:word:]]+\) ` +
		`REFERENCES ([[:word:]]+)\(([[:word:]]+)\)` +
		`(?:\s*ON UPDATE (CASCADE|RESTRICT|SET NULL|SET DEFAULT|NO ACTION))?` +
		`(?:\s*ON DELETE (CASCADE|RESTRICT|SET NULL|SET DEFAULT|NO ACTION))?$`,
)

// UnquoteIdentifier is the analog of pq.QuoteIdentifier.
func UnquoteIdentifier(quotedIdent string) string {
	if len(quotedIdent) >= 2 && quotedIdent[0] == '"' && quotedIdent[len(quotedIdent)-1] == '"' {
		quotedIdent = strings.ReplaceAll(quotedIdent[1:len(quotedIdent)-1], "\"\"", "\"")
	}
	return quotedIdent
}

var (
	rxIndex        = regexp.MustCompile(`^CREATE( UNIQUE)? INDEX (.*) ON .*\.(.*) USING btree \((.*)\)$`)
	indexDirRemove = strings.NewReplacer(" ASC", "", " DESC", "")
)

func parseColumnDefault(columnDefault string) string {
	// hackity hack: See the comments in parseIndexDefinition for why we do this.
	if columnDefault == "nextval('storagenode_storage_tallies_id_seq'::regclass)" {
		return "nextval('accounting_raws_id_seq'::regclass)"
	}

	// hackity hack: cockroach sometimes adds type descriptors to the default. ignore em!
	if idx := strings.Index(columnDefault, ":::"); idx >= 0 {
		columnDefault = columnDefault[:idx]
	}

	return columnDefault
}

func parseIndexDefinition(indexdef string) (*dbschema.Index, error) {
	matches := rxIndex.FindStringSubmatch(indexdef)
	if matches == nil {
		return nil, errs.New("unable to parse index (you should go make the parser better): %q", indexdef)
	}

	// hackity hack: cockroach returns all primary key index names as `"primary"`, but sometimes
	// our migrations create them with explicit names. so let's change all of them.
	name := matches[2]
	if name == `"primary"` {
		name = matches[3] + "_pkey"
	}

	// hackity hack: sometimes they end with _pk, sometimes they end with _pkey. let's make them
	// all end with _pkey.
	if strings.HasSuffix(name, "_pk") {
		name = name[:len(name)-3] + "_pkey"
	}

	// biggest hackity hack of all: we apparently did
	//
	//     CREATE TABLE accounting_raws ( ... )
	//     ALTER TABLE accounting_raws RENAME TO storagenode_storage_tallies
	//
	// but that means the primary key index is still named accounting_raws_pkey and not
	// the expected storagenode_storage_tallies_pkey.
	//
	// "No big deal", you might say, "just add an ALTER INDEX". Ah, but recall: cockroach
	// does not name their primary key indexes. They are instead all named `"primary"`.
	// Now, at this point, a clever person might suggest ALTER INDEX IF EXISTS so that
	// it renames it on postgres but not cockroach. Surely if the index does not exist
	// it will happily succeed. You'd like to think that, wouldn't you! Alas, cockroach
	// will error on ALTER INDEX IF EXISTS even if the index does not exist. Basic
	// conditionals are apparently too hard for it.
	//
	// Undaunted, I searched their bug tracker and found this comment within this issue:
	//     https://github.com/cockroachdb/cockroach/issues/42399#issuecomment-558377915
	// It turns out, you apparently need to specify the index with some sort of `@` sigil
	// or it just errors with an unhelpful message. But only a great fool would think that
	// the query would remain valid for postgres!
	//
	// In summary, because cockroach errors even if the index does not exist, I can clearly
	// not use cockroach. But because postgres will error if the sigil is included, I can
	// clearly not use postgres.
	//
	// As a last resort, one may suggest changing the postgres.N.sql file to ALSO create
	// the wrong table name and rename it. Truly, they have a dizzying intellect. But prepare
	// yourself for the final killing blow: if we do that, then the final output does not match
	// the dbx schema that is autogenerated, and the test still fails.
	//
	// The lesson? Never go in against a database when death is on the line. HA HA HA HA...
	//
	// Bleh.
	if name == "accounting_raws_pkey" {
		name = "storagenode_storage_tallies_pkey"
	}

	return &dbschema.Index{
		Name:    name,
		Table:   matches[3],
		Unique:  matches[1] != "",
		Columns: strings.Split(indexDirRemove.Replace(matches[4]), ", "),
	}, nil
}