-
Notifications
You must be signed in to change notification settings - Fork 1
/
postgres.go
131 lines (120 loc) · 3 KB
/
postgres.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package database
import (
"database/sql"
"fmt"
_ "github.com/lib/pq"
"github.com/whaios/showdocdb/log"
)
func init() {
postgreFactory = newPostgreSQL
}
func newPostgreSQL(host, user, pwd, dbName, schema string) Query {
// "postgres://username:password@localhost:5432/database_name?sslmode=disable"
dsn := fmt.Sprintf("postgres://%s:%s@%s/%s?sslmode=disable", user, pwd, host, dbName)
if schema == "" {
schema = "public"
}
return &postgreSQL{
dsn: dsn,
dbName: dbName,
schema: schema,
}
}
type postgreSQL struct {
dsn string
dbName string
schema string // 模式,多个表的集合,默认 public
db *sql.DB
}
func (d *postgreSQL) Open() error {
db, err := sql.Open(PostgreSQL, d.dsn)
d.db = db
return err
}
func (d *postgreSQL) Close() error {
if d.db != nil {
return d.db.Close()
}
return nil
}
func (d *postgreSQL) Query() ([]*Table, error) {
tbs, err := d.queryTable()
if err != nil {
return nil, err
}
cols, err := d.queryColumn()
if err != nil {
return nil, err
}
tbmap := make(map[string]*Table, len(tbs))
{
for _, tb := range tbs {
tbmap[tb.Name] = tb
log.Debug("table: %s", tb.Name)
}
}
for _, col := range cols {
tb, ok := tbmap[col.TableName]
if !ok {
continue
}
tb.Columns = append(tb.Columns, col)
log.Debug("column: %s.%s", col.TableName, col.Name)
}
return tbs, err
}
func (d *postgreSQL) queryTable() ([]*Table, error) {
querySql := `
SELECT T.table_name, COALESCE(obj_description(relfilenode, 'pg_class'), '') AS comment
FROM information_schema.TABLES AS T, pg_class C
WHERE T.table_name = C.relname
AND TABLE_SCHEMA = $1
ORDER BY T.table_name`
rows, err := d.db.Query(querySql, d.schema)
if err != nil {
return nil, err
}
defer rows.Close()
tbs := make([]*Table, 0)
for rows.Next() {
var name, comment string
if err = rows.Scan(&name, &comment); err != nil {
return tbs, err
}
tbs = append(tbs, &Table{Name: name, Comment: comment})
}
return tbs, nil
}
func (d *postgreSQL) queryColumn() ([]*Column, error) {
querySql := `
SELECT COL.table_name, COL.column_name, COL.column_default, COL.is_nullable, COL.udt_name,
COALESCE(col_description(ATT.attrelid, ATT.attnum), '') AS comment
FROM information_schema.COLUMNS COL, pg_attribute ATT, pg_class CLS
WHERE COL.TABLE_SCHEMA = $1
AND COL.table_name = CLS.relname
AND COL.column_name = ATT.attname
AND ATT.attrelid = CLS.oid
ORDER BY COL.table_name, COL.ordinal_position`
rows, err := d.db.Query(querySql, d.schema)
if err != nil {
return nil, err
}
defer rows.Close()
cols := make([]*Column, 0)
for rows.Next() {
var tableName, colName, colDefault sql.NullString
var isNull, colType, comment string
if err = rows.Scan(&tableName, &colName, &colDefault, &isNull, &colType, &comment); err != nil {
return cols, err
}
cols = append(cols, &Column{
TableName: tableName.String,
Name: colName.String,
Default: colDefault.String,
IsNull: isNull,
Type: colType,
Comment: comment,
})
}
return cols, nil
}