super-graph/psql/tables.go

247 lines
5.4 KiB
Go
Raw Normal View History

2019-03-24 14:57:29 +01:00
package psql
import (
2019-09-26 06:35:31 +02:00
"context"
2019-03-24 14:57:29 +01:00
"fmt"
"strconv"
2019-03-24 14:57:29 +01:00
"strings"
2019-09-30 07:48:29 +02:00
"github.com/jackc/pgtype"
2019-09-26 06:35:31 +02:00
"github.com/jackc/pgx/v4/pgxpool"
2019-03-24 14:57:29 +01:00
)
type DBInfo struct {
Version int
Tables []DBTable
Columns [][]DBColumn
colmap map[string]map[string]*DBColumn
}
func GetDBInfo(db *pgxpool.Pool) (*DBInfo, error) {
di := &DBInfo{}
dbc, err := db.Acquire(context.Background())
if err != nil {
return nil, fmt.Errorf("error acquiring connection from pool: %w", err)
}
defer dbc.Release()
var version string
err = dbc.QueryRow(context.Background(), `SHOW server_version_num`).Scan(&version)
if err != nil {
return nil, fmt.Errorf("error fetching version: %w", err)
}
di.Version, err = strconv.Atoi(version)
if err != nil {
return nil, err
}
di.Tables, err = GetTables(dbc)
if err != nil {
return nil, err
}
di.colmap = make(map[string]map[string]*DBColumn, len(di.Tables))
for i, t := range di.Tables {
cols, err := GetColumns(dbc, "public", t.Name)
if err != nil {
return nil, err
}
di.Columns = append(di.Columns, cols)
di.colmap[t.Key] = make(map[string]*DBColumn, len(cols))
for n, c := range di.Columns[i] {
di.colmap[t.Key][c.Key] = &di.Columns[i][n]
}
}
return di, nil
}
func (di *DBInfo) AddTable(t DBTable, cols []DBColumn) {
t.ID = di.Tables[len(di.Tables)-1].ID
di.Tables = append(di.Tables, t)
di.colmap[t.Key] = make(map[string]*DBColumn, len(cols))
for i := range cols {
cols[i].ID = int16(i)
c := &cols[i]
di.colmap[t.Key][c.Key] = c
}
di.Columns = append(di.Columns, cols)
}
func (di *DBInfo) GetColumn(table, column string) (*DBColumn, bool) {
v, ok := di.colmap[strings.ToLower(table)][strings.ToLower(column)]
return v, ok
}
type DBTable struct {
ID int
2019-09-26 06:35:31 +02:00
Name string
Key string
2019-09-26 06:35:31 +02:00
Type string
}
func GetTables(dbc *pgxpool.Conn) ([]DBTable, error) {
sqlStmt := `
2019-09-29 21:20:59 +02:00
SELECT
c.relname as "name",
CASE c.relkind WHEN 'r' THEN 'table'
2019-09-26 06:35:31 +02:00
WHEN 'v' THEN 'view'
WHEN 'm' THEN 'materialized view'
WHEN 'f' THEN 'foreign table'
END as "type"
FROM pg_catalog.pg_class c
2019-09-29 21:20:59 +02:00
LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind IN ('r','v','m','f','')
2019-10-01 06:00:15 +02:00
AND n.nspname <> ('pg_catalog')
AND n.nspname <> ('information_schema')
AND n.nspname !~ ('^pg_toast')
2019-09-29 21:20:59 +02:00
AND pg_catalog.pg_table_is_visible(c.oid);`
var tables []DBTable
2019-09-26 06:35:31 +02:00
rows, err := dbc.Query(context.Background(), sqlStmt)
if err != nil {
return nil, fmt.Errorf("Error fetching tables: %s", err)
}
2019-09-26 06:35:31 +02:00
defer rows.Close()
for i := 0; rows.Next(); i++ {
t := DBTable{ID: i}
2019-09-26 06:35:31 +02:00
err = rows.Scan(&t.Name, &t.Type)
if err != nil {
return nil, err
}
t.Key = strings.ToLower(t.Name)
2019-12-25 07:24:30 +01:00
if t.Key != "schema_migrations" && t.Key != "ar_internal_metadata" {
tables = append(tables, t)
}
2019-09-26 06:35:31 +02:00
}
return tables, nil
}
type DBColumn struct {
2019-10-01 06:00:15 +02:00
ID int16
2019-09-26 06:35:31 +02:00
Name string
Key string
2019-09-26 06:35:31 +02:00
Type string
Array bool
2019-09-26 06:35:31 +02:00
NotNull bool
PrimaryKey bool
UniqueKey bool
FKeyTable string
2019-10-01 06:00:15 +02:00
FKeyColID []int16
2019-09-26 06:35:31 +02:00
fKeyColID pgtype.Int2Array
}
func GetColumns(dbc *pgxpool.Conn, schema, table string) ([]DBColumn, error) {
sqlStmt := `
2019-09-29 21:20:59 +02:00
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 f.attndims != 0 THEN true
WHEN right(pg_catalog.format_type(f.atttypid,f.atttypmod), 2) = '[]' THEN true
ELSE false
END AS array,
2019-09-29 21:20:59 +02:00
CASE
WHEN p.contype = ('p'::char) THEN true
ELSE false
END AS primarykey,
CASE
WHEN p.contype = ('u'::char) THEN true
ELSE false
END AS uniquekey,
CASE
WHEN p.contype = ('f'::char) THEN g.relname
ELSE ''::text
END AS foreignkey,
CASE
WHEN p.contype = ('f'::char) THEN p.confkey
ELSE ARRAY[]::int2[]
END AS foreignkey_fieldnum
FROM pg_attribute f
JOIN pg_class c ON c.oid = f.attrelid
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 IN ('r', 'v', 'm', 'f')
2019-09-29 21:20:59 +02:00
AND n.nspname = $1 -- Replace with Schema name
AND c.relname = $2 -- Replace with table name
AND f.attnum > 0
AND f.attisdropped = false
ORDER BY id;`
2019-09-26 06:35:31 +02:00
rows, err := dbc.Query(context.Background(), sqlStmt, schema, table)
if err != nil {
return nil, fmt.Errorf("error fetching columns: %s", err)
}
2019-09-26 06:35:31 +02:00
defer rows.Close()
cmap := make(map[int16]DBColumn)
2019-09-29 21:20:59 +02:00
2019-09-26 06:35:31 +02:00
for rows.Next() {
c := DBColumn{}
2019-12-12 06:47:56 +01:00
err = rows.Scan(&c.ID, &c.Name, &c.NotNull, &c.Type, &c.Array, &c.PrimaryKey, &c.UniqueKey, &c.FKeyTable, &c.fKeyColID)
2019-09-26 06:35:31 +02:00
if err != nil {
return nil, err
}
2019-09-29 21:20:59 +02:00
if v, ok := cmap[c.ID]; ok {
if c.PrimaryKey {
v.PrimaryKey = true
2019-12-25 07:24:30 +01:00
v.UniqueKey = true
2019-09-29 21:20:59 +02:00
}
if c.NotNull {
v.NotNull = true
}
if c.UniqueKey {
v.UniqueKey = true
}
if c.Array {
v.Array = true
}
2019-12-12 06:47:56 +01:00
if len(c.FKeyTable) != 0 {
v.FKeyTable = c.FKeyTable
}
if c.fKeyColID.Elements != nil {
v.fKeyColID = c.fKeyColID
err := v.fKeyColID.AssignTo(&v.FKeyColID)
if err != nil {
return nil, err
}
}
cmap[c.ID] = v
2019-09-29 21:20:59 +02:00
} else {
2019-10-01 06:00:15 +02:00
err := c.fKeyColID.AssignTo(&c.FKeyColID)
if err != nil {
return nil, err
}
c.Key = strings.ToLower(c.Name)
2019-12-25 07:24:30 +01:00
if c.PrimaryKey {
c.UniqueKey = true
}
cmap[c.ID] = c
2019-09-29 21:20:59 +02:00
}
}
cols := make([]DBColumn, 0, len(cmap))
2019-12-12 06:47:56 +01:00
for i := range cmap {
cols = append(cols, cmap[i])
}
2019-09-26 06:35:31 +02:00
return cols, nil
}