package psql import ( "context" "fmt" "strings" "github.com/gobuffalo/flect" "github.com/jackc/pgx/pgtype" "github.com/jackc/pgx/v4/pgxpool" ) type DBTable struct { Name string Type string } func GetTables(dbc *pgxpool.Conn) ([]*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 tables []*DBTable rows, err := dbc.Query(context.Background(), sqlStmt) if err != nil { return nil, fmt.Errorf("Error fetching tables: %s", err) } defer rows.Close() for rows.Next() { t := DBTable{} err = rows.Scan(&t.Name, &t.Type) if err != nil { return nil, err } tables = append(tables, &t) } return tables, nil } type DBColumn struct { ID int Name string Type string NotNull bool PrimaryKey bool UniqueKey bool FKeyTable string FKeyColID []int fKeyColID pgtype.Int2Array } func GetColumns(dbc *pgxpool.Conn, 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 true ELSE false END AS primarykey, CASE WHEN p.contype = 'u' THEN true ELSE false END AS uniquekey, CASE WHEN p.contype = 'f' THEN g.relname ELSE ''::text END AS foreignkey, CASE WHEN p.contype = 'f' THEN p.confkey ELSE ARRAY[]::int2[] 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; ` var cols []*DBColumn rows, err := dbc.Query(context.Background(), sqlStmt, schema, table) if err != nil { return nil, fmt.Errorf("error fetching columns: %s", err) } defer rows.Close() for rows.Next() { c := DBColumn{} err = rows.Scan(&c.ID, &c.Name, &c.NotNull, &c.Type, &c.PrimaryKey, &c.UniqueKey, &c.FKeyTable, &c.fKeyColID) if err != nil { return nil, err } c.fKeyColID.AssignTo(&c.FKeyColID) cols = append(cols, &c) } return cols, 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 *pgxpool.Pool, 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{}), } dbc, err := db.Acquire(context.Background()) if err != nil { return nil, fmt.Errorf("error acquiring connection from pool") } defer dbc.Release() tables, err := GetTables(dbc) if err != nil { return nil, err } for _, t := range tables { cols, err := GetColumns(dbc, "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 }