super-graph/core/internal/psql/tables.go

323 lines
6.8 KiB
Go
Raw Normal View History

2019-03-24 14:57:29 +01:00
package psql
import (
"database/sql"
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-03-24 14:57:29 +01:00
)
type DBInfo struct {
Version int
Tables []DBTable
Columns [][]DBColumn
Functions []DBFunction
colMap map[string]map[string]*DBColumn
}
func GetDBInfo(db *sql.DB) (*DBInfo, error) {
di := &DBInfo{}
var version string
err := db.QueryRow(`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(db)
if err != nil {
return nil, err
}
for _, t := range di.Tables {
cols, err := GetColumns(db, "public", t.Name)
if err != nil {
return nil, err
}
di.Columns = append(di.Columns, cols)
}
di.colMap = newColMap(di.Tables, di.Columns)
di.Functions, err = GetFunctions(db)
if err != nil {
return nil, err
}
return di, nil
}
func newColMap(tables []DBTable, columns [][]DBColumn) map[string]map[string]*DBColumn {
cm := make(map[string]map[string]*DBColumn, len(tables))
for i, t := range tables {
cols := columns[i]
cm[t.Key] = make(map[string]*DBColumn, len(cols))
for n, c := range cols {
cm[t.Key][c.Key] = &columns[i][n]
}
}
return cm
}
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(db *sql.DB) ([]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
rows, err := db.Query(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(db *sql.DB, 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
2020-02-07 07:12:14 +01:00
WHEN p.contype = ('f'::char) THEN p.confkey::int2[]
2019-09-29 21:20:59 +02:00
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 := db.Query(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
}
2020-02-10 07:45:37 +01:00
type DBFunction struct {
Name string
Params []DBFuncParam
}
type DBFuncParam struct {
ID int
Name string
Type string
}
func GetFunctions(db *sql.DB) ([]DBFunction, error) {
sqlStmt := `
SELECT
routines.routine_name,
parameters.specific_name,
parameters.data_type,
parameters.parameter_name,
parameters.ordinal_position
FROM
information_schema.routines
RIGHT JOIN
information_schema.parameters
ON (routines.specific_name = parameters.specific_name and parameters.ordinal_position IS NOT NULL)
WHERE
routines.specific_schema = 'public'
ORDER BY
routines.routine_name, parameters.ordinal_position;`
rows, err := db.Query(sqlStmt)
if err != nil {
return nil, fmt.Errorf("Error fetching functions: %s", err)
}
defer rows.Close()
var funcs []DBFunction
fm := make(map[string]int)
for rows.Next() {
var fn, fid string
fp := DBFuncParam{}
err = rows.Scan(&fn, &fid, &fp.Type, &fp.Name, &fp.ID)
if err != nil {
return nil, err
}
if i, ok := fm[fid]; ok {
funcs[i].Params = append(funcs[i].Params, fp)
} else {
funcs = append(funcs, DBFunction{Name: fn, Params: []DBFuncParam{fp}})
fm[fid] = len(funcs) - 1
}
}
return funcs, nil
}
2020-02-10 07:45:37 +01:00
// func GetValType(type string) qcode.ValType {
// switch {
// case "bigint", "integer", "smallint", "numeric", "bigserial":
// return qcode.ValInt
// case "double precision", "real":
// return qcode.ValFloat
// case ""
// }
// }