super-graph/psql/schema.go

package psql

import (
	"fmt"
	"strings"

	"github.com/gobuffalo/flect"
	"github.com/jackc/pgx/v4/pgxpool"
)

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

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

type RelType int

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

type DBRel struct {
	Type    RelType
	Through string
	ColT    string
	Left    struct {
		col   *DBColumn
		Table string
		Col   string
		Array bool
	}
	Right struct {
		col   *DBColumn
		Table string
		Col   string
		Array bool
	}
}

func NewDBSchema(db *pgxpool.Pool,
	info *DBInfo, aliases map[string][]string) (*DBSchema, error) {

	schema := &DBSchema{
		t:  make(map[string]*DBTableInfo),
		rm: make(map[string]map[string]*DBRel),
	}

	for i, t := range info.Tables {
		err := schema.addTable(t, info.Columns[i], aliases)
		if err != nil {
			return nil, err
		}
	}

	for i, t := range info.Tables {
		err := schema.updateRelationships(t, info.Columns[i])
		if err != nil {
			return nil, err
		}
	}

	return schema, nil
}

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

	colmap := make(map[string]*DBColumn, len(cols))
	colidmap := make(map[int16]*DBColumn, len(cols))

	singular := flect.Singularize(t.Key)
	s.t[singular] = &DBTableInfo{
		Name:     t.Name,
		Singular: true,
		Columns:  cols,
		ColMap:   colmap,
		ColIDMap: colidmap,
	}

	plural := flect.Pluralize(t.Key)
	s.t[plural] = &DBTableInfo{
		Name:     t.Name,
		Singular: false,
		Columns:  cols,
		ColMap:   colmap,
		ColIDMap: colidmap,
	}

	if al, ok := aliases[t.Key]; 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]
		}
	}

	for i := range cols {
		c := &cols[i]

		switch {
		case c.Type == "tsvector":
			s.t[singular].TSVCol = c
			s.t[plural].TSVCol = c

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

		colmap[c.Key] = c
		colidmap[c.ID] = c
	}

	return nil
}

func (s *DBSchema) updateRelationships(t DBTable, cols []DBColumn) error {
	jcols := make([]DBColumn, 0, len(cols))
	ct := t.Key
	cti, ok := s.t[ct]
	if !ok {
		return fmt.Errorf("invalid foreign key table '%s'", ct)
	}

	for _, c := range cols {
		if len(c.FKeyTable) == 0 || len(c.FKeyColID) == 0 {
			continue
		}

		// Foreign key column name
		ft := strings.ToLower(c.FKeyTable)
		fcid := c.FKeyColID[0]

		ti, ok := s.t[ft]
		if !ok {
			return fmt.Errorf("invalid foreign key table '%s'", ft)
		}

		fc, ok := ti.ColIDMap[fcid]
		if !ok {
			return fmt.Errorf("invalid foreign key column id '%d' for table '%s'",
				fcid, ti.Name)
		}

		var rel1, rel2 *DBRel

		// One-to-many relation between current table and the
		// table in the foreign key
		if fc.UniqueKey {
			rel1 = &DBRel{Type: RelOneToOne}
		} else {
			rel1 = &DBRel{Type: RelOneToMany}
		}

		rel1.Left.col = &c
		rel1.Left.Table = t.Name
		rel1.Left.Col = c.Name
		rel1.Left.Array = c.Array

		rel1.Right.col = fc
		rel1.Right.Table = c.FKeyTable
		rel1.Right.Col = fc.Name
		rel1.Right.Array = fc.Array

		if err := s.SetRel(ct, ft, rel1); err != nil {
			return err
		}

		// One-to-many reverse relation between the foreign key table and the
		// the current table
		if c.UniqueKey {
			rel2 = &DBRel{Type: RelOneToOne}
		} else {
			rel2 = &DBRel{Type: RelOneToMany}
		}

		rel2.Left.Table = c.FKeyTable
		rel2.Left.Col = fc.Name
		rel2.Left.Array = fc.Array

		rel2.Right.Table = t.Name
		rel2.Right.Col = c.Name
		rel2.Right.Array = c.Array

		if err := s.SetRel(ft, ct, rel2); err != nil {
			return err
		}

		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 {
					continue
				}
				err := s.updateSchemaOTMT(cti, jcols[i], jcols[n])
				if err != nil {
					return err
				}
			}
		}
	}

	return nil
}

func (s *DBSchema) updateSchemaOTMT(
	ti *DBTableInfo, col1, col2 DBColumn) error {

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

	fc1, ok := s.t[t1].ColIDMap[col1.FKeyColID[0]]
	if !ok {
		return fmt.Errorf("invalid foreign key column id '%d' for table '%s'",
			col1.FKeyColID[0], ti.Name)
	}
	fc2, ok := s.t[t2].ColIDMap[col2.FKeyColID[0]]
	if !ok {
		return fmt.Errorf("invalid foreign key column id '%d' for table '%s'",
			col2.FKeyColID[0], ti.Name)
	}

	// One-to-many-through relation between 1nd foreign key table and the
	// 2nd foreign key table
	rel1 := &DBRel{Type: RelOneToManyThrough}
	rel1.Through = ti.Name
	rel1.ColT = col2.Name

	rel1.Left.Table = col2.FKeyTable
	rel1.Left.Col = fc2.Name

	rel1.Right.Table = ti.Name
	rel1.Right.Col = col1.Name

	if err := s.SetRel(t1, t2, rel1); err != nil {
		return err
	}

	// One-to-many-through relation between 2nd foreign key table and the
	// 1nd foreign key table
	rel2 := &DBRel{Type: RelOneToManyThrough}
	rel2.Through = ti.Name
	rel2.ColT = col1.Name

	rel2.Left.Table = col1.FKeyTable
	rel2.Left.Col = fc1.Name

	rel2.Right.Table = ti.Name
	rel2.Right.Col = col2.Name

	if err := s.SetRel(t2, t1, rel2); err != nil {
		return err
	}

	return nil
}

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
}