super-graph/psql/tables.go

package psql

import (
	"fmt"
	"strings"

	"github.com/go-pg/pg"
	"github.com/gobuffalo/flect"
)

type DBTable struct {
	Name string `sql:"name"`
	Type string `sql:"type"`
}

func GetTables(db *pg.DB) ([]*DBTable, error) {
	sqlStmt := `
	SELECT
  c.relname as "name",
  CASE c.relkind WHEN 'r' THEN 'table'
  WHEN 'v' THEN 'view'
  WHEN 'm' THEN 'materialized view'
  WHEN 'f' THEN 'foreign table' END as "type"
FROM pg_catalog.pg_class c
     LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind IN ('r','v','m','f','')
      AND n.nspname <> 'pg_catalog'
      AND n.nspname <> 'information_schema'
      AND n.nspname !~ '^pg_toast'
  AND pg_catalog.pg_table_is_visible(c.oid);
	`

	var t []*DBTable
	_, err := db.Query(&t, sqlStmt)

	if err != nil {
		return nil, fmt.Errorf("Error fetching tables: %s", err)
	}

	return t, nil
}

type DBColumn struct {
	ID         int    `sql:"id"`
	Name       string `sql:"name"`
	Type       string `sql:"type"`
	NotNull    bool   `sql:"notnull"`
	PrimaryKey bool   `sql:"primarykey"`
	Uniquekey  bool   `sql:"uniquekey"`
	FKeyTable  string `sql:"foreignkey"`
	FKeyColID  []int  `sql:"foreignkey_fieldnum,array"`
}

func GetColumns(db *pg.DB, schema, table string) ([]*DBColumn, error) {
	sqlStmt := `
	SELECT  
    f.attnum AS id,  
    f.attname AS name,  
    f.attnotnull AS notnull,  
    pg_catalog.format_type(f.atttypid,f.atttypmod) AS type,  
    CASE  
        WHEN p.contype = 'p' THEN 't'  
        ELSE 'f'  
    END AS primarykey,  
    CASE  
        WHEN p.contype = 'u' THEN 't'  
        ELSE 'f'
    END AS uniquekey,
    CASE
        WHEN p.contype = 'f' THEN g.relname
    END AS foreignkey,
    CASE
        WHEN p.contype = 'f' THEN p.confkey
    END AS foreignkey_fieldnum
FROM pg_attribute f  
    JOIN pg_class c ON c.oid = f.attrelid  
    JOIN pg_type t ON t.oid = f.atttypid  
    LEFT JOIN pg_attrdef d ON d.adrelid = c.oid AND d.adnum = f.attnum  
    LEFT JOIN pg_namespace n ON n.oid = c.relnamespace  
    LEFT JOIN pg_constraint p ON p.conrelid = c.oid AND f.attnum = ANY (p.conkey)  
    LEFT JOIN pg_class AS g ON p.confrelid = g.oid  
WHERE c.relkind = 'r'::char  
    AND n.nspname = $1  -- Replace with Schema name  
    AND c.relname = $2  -- Replace with table name  
    AND f.attnum > 0 ORDER BY id;
	`

	stmt, err := db.Prepare(sqlStmt)
	if err != nil {
		return nil, fmt.Errorf("error fetching columns: %s", err)
	}

	var t []*DBColumn
	_, err = stmt.Query(&t, schema, table)
	if err != nil {
		return nil, fmt.Errorf("error fetching columns: %s", err)
	}

	return t, nil
}

type DBSchema struct {
	t  map[string]*DBTableInfo
	rm map[string]map[string]*DBRel
	al map[string]struct{}
}

type DBTableInfo struct {
	Name        string
	Singular    bool
	PrimaryCol  string
	TSVCol      string
	Columns     map[string]*DBColumn
	ColumnNames []string
}

type RelType int

const (
	RelBelongTo RelType = iota + 1
	RelOneToMany
	RelOneToManyThrough
	RelRemote
)

type DBRel struct {
	Type    RelType
	Through string
	ColT    string
	Col1    string
	Col2    string
}

func NewDBSchema(db *pg.DB, aliases map[string][]string) (*DBSchema, error) {
	schema := &DBSchema{
		t:  make(map[string]*DBTableInfo),
		rm: make(map[string]map[string]*DBRel),
		al: make(map[string]struct{}),
	}

	tables, err := GetTables(db)
	if err != nil {
		return nil, err
	}

	for _, t := range tables {
		cols, err := GetColumns(db, "public", t.Name)
		if err != nil {
			return nil, err
		}

		schema.updateSchema(t, cols, aliases)
	}

	return schema, nil
}

func (s *DBSchema) updateSchema(
	t *DBTable,
	cols []*DBColumn,
	aliases map[string][]string) {

	// Foreign key columns in current table
	colByID := make(map[int]*DBColumn)
	columns := make(map[string]*DBColumn, len(cols))
	colNames := make([]string, len(cols))

	for i := range cols {
		c := cols[i]
		name := strings.ToLower(c.Name)
		columns[name] = cols[i]
		colNames = append(colNames, name)
		colByID[c.ID] = cols[i]
	}

	singular := strings.ToLower(flect.Singularize(t.Name))
	s.t[singular] = &DBTableInfo{
		Name:        t.Name,
		Singular:    true,
		Columns:     columns,
		ColumnNames: colNames,
	}

	plural := strings.ToLower(flect.Pluralize(t.Name))
	s.t[plural] = &DBTableInfo{
		Name:        t.Name,
		Singular:    false,
		Columns:     columns,
		ColumnNames: colNames,
	}

	ct := strings.ToLower(t.Name)

	if al, ok := aliases[ct]; ok {
		for i := range al {
			k1 := flect.Singularize(al[i])
			s.t[k1] = s.t[singular]

			k2 := flect.Pluralize(al[i])
			s.t[k2] = s.t[plural]

			s.al[k1] = struct{}{}
			s.al[k2] = struct{}{}
		}
	}

	jcols := make([]*DBColumn, 0, len(cols))

	for _, c := range cols {
		switch {
		case c.Type == "tsvector":
			s.t[singular].TSVCol = c.Name
			s.t[plural].TSVCol = c.Name

		case c.PrimaryKey:
			s.t[singular].PrimaryCol = c.Name
			s.t[plural].PrimaryCol = c.Name

		case len(c.FKeyTable) != 0:
			if len(c.FKeyColID) == 0 {
				continue
			}

			// Foreign key column name
			ft := strings.ToLower(c.FKeyTable)
			fc, ok := colByID[c.FKeyColID[0]]
			if !ok {
				continue
			}

			// Belongs-to relation between current table and the
			// table in the foreign key
			rel1 := &DBRel{RelBelongTo, "", "", c.Name, fc.Name}
			s.SetRel(ct, ft, rel1)

			// One-to-many relation between the foreign key table and the
			// the current table
			rel2 := &DBRel{RelOneToMany, "", "", fc.Name, c.Name}
			s.SetRel(ft, ct, rel2)

			jcols = append(jcols, c)
		}
	}

	// If table contains multiple foreign key columns it's a possible
	// join table for many-to-many relationships or multiple one-to-many
	// relations

	// Below one-to-many relations use the current table as the
	// join table aka through table.
	if len(jcols) > 1 {
		for i := range jcols {
			for n := range jcols {
				if n != i {
					s.updateSchemaOTMT(ct, jcols[i], jcols[n], colByID)
				}
			}
		}
	}
}

func (s *DBSchema) updateSchemaOTMT(
	ct string,
	col1, col2 *DBColumn,
	colByID map[int]*DBColumn) {

	t1 := strings.ToLower(col1.FKeyTable)
	t2 := strings.ToLower(col2.FKeyTable)

	fc1, ok := colByID[col1.FKeyColID[0]]
	if !ok {
		return
	}
	fc2, ok := colByID[col2.FKeyColID[0]]
	if !ok {
		return
	}

	// One-to-many-through relation between 1nd foreign key table and the
	// 2nd foreign key table
	//rel1 := &DBRel{RelOneToManyThrough, ct, fc1.Name, col1.Name}
	rel1 := &DBRel{RelOneToManyThrough, ct, col2.Name, fc2.Name, col1.Name}
	s.SetRel(t1, t2, rel1)

	// One-to-many-through relation between 2nd foreign key table and the
	// 1nd foreign key table
	//rel2 := &DBRel{RelOneToManyThrough, ct, col2.Name, fc2.Name}
	rel2 := &DBRel{RelOneToManyThrough, ct, col1.Name, fc1.Name, col2.Name}
	s.SetRel(t2, t1, rel2)
}

func (s *DBSchema) GetTable(table string) (*DBTableInfo, error) {
	t, ok := s.t[table]
	if !ok {
		return nil, fmt.Errorf("unknown table '%s'", table)
	}
	return t, nil
}

func (s *DBSchema) SetRel(child, parent string, rel *DBRel) error {
	sc := strings.ToLower(flect.Singularize(child))
	pc := strings.ToLower(flect.Pluralize(child))

	if _, ok := s.rm[sc]; !ok {
		s.rm[sc] = make(map[string]*DBRel)
	}

	if _, ok := s.rm[pc]; !ok {
		s.rm[pc] = make(map[string]*DBRel)
	}

	sp := strings.ToLower(flect.Singularize(parent))
	pp := strings.ToLower(flect.Pluralize(parent))

	s.rm[sc][sp] = rel
	s.rm[sc][pp] = rel
	s.rm[pc][sp] = rel
	s.rm[pc][pp] = rel

	return nil
}

func (s *DBSchema) GetRel(child, parent string) (*DBRel, error) {
	rel, ok := s.rm[child][parent]
	if !ok {
		return nil, fmt.Errorf("unknown relationship '%s' -> '%s'",
			child, parent)
	}
	return rel, nil
}

func (s *DBSchema) IsAlias(name string) bool {
	_, ok := s.al[name]
	return ok
}