-
Notifications
You must be signed in to change notification settings - Fork 3
/
data.go
154 lines (136 loc) · 3.89 KB
/
data.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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
package pgdump
import (
"database/sql"
"fmt"
"strings"
)
// options for dumping selective tables.
type TableOptions struct {
TableSuffix string
TablePrefix string
Schema string
}
// returns a slice of table names matching options, if left blank will default to :
//
// -> no prefix or suffix
// -> public schema
func getTables(db *sql.DB, opts *TableOptions) ([]string, error) {
var (
query string
)
if opts != nil {
if opts.Schema == "" {
opts.Schema = "public"
}
query = fmt.Sprintf("SELECT table_name FROM information_schema.tables WHERE table_schema = '%s' AND table_name LIKE '%s'", opts.Schema, (opts.TablePrefix + "%%" + opts.TableSuffix))
} else {
query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"
}
rows, err := db.Query(query)
if err != nil {
return nil, err
}
defer rows.Close()
var tables []string
for rows.Next() {
var tableName string
if err := rows.Scan(&tableName); err != nil {
return nil, err
}
if opts.Schema != "public" {
tables = append(tables, opts.Schema+"."+tableName)
} else {
tables = append(tables, tableName)
}
}
return tables, nil
}
// generates the SQL for creating a table, including column definitions.
func getCreateTableStatement(db *sql.DB, tableName string) (string, error) {
query := fmt.Sprintf("SELECT column_name, data_type, character_maximum_length FROM information_schema.columns WHERE table_name = '%s'", tableName)
rows, err := db.Query(query)
if err != nil {
return "", err
}
defer rows.Close()
var columns []string
for rows.Next() {
var columnName, dataType string
var charMaxLength *int
if err := rows.Scan(&columnName, &dataType, &charMaxLength); err != nil {
return "", err
}
columnDef := fmt.Sprintf("%s %s", columnName, dataType)
if charMaxLength != nil {
columnDef += fmt.Sprintf("(%d)", *charMaxLength)
}
columns = append(columns, columnDef)
}
return fmt.Sprintf("CREATE TABLE %s (\n %s\n);", tableName, strings.Join(columns, ",\n ")), nil
}
// generates the COPY command to import data for a table.
func getTableDataCopyFormat(db *sql.DB, tableName string) (string, error) {
query := fmt.Sprintf("SELECT * FROM %s", tableName)
rows, err := db.Query(query)
if err != nil {
return "", err
}
defer rows.Close()
columns, err := rows.Columns()
if err != nil {
return "", err
}
values := make([]sql.RawBytes, len(columns))
scanArgs := make([]interface{}, len(values))
for i := range values {
scanArgs[i] = &values[i]
}
var output strings.Builder
output.WriteString(fmt.Sprintf("COPY %s (%s) FROM stdin;\n", tableName, strings.Join(columns, ", ")))
for rows.Next() {
err := rows.Scan(scanArgs...)
if err != nil {
return "", err
}
var valueStrings []string
for _, value := range values {
valueStrings = append(valueStrings, string(value))
}
output.WriteString(strings.Join(valueStrings, "\t") + "\n")
}
output.WriteString("\\.\n")
return output.String(), nil
}
func getTableDataAsCSV(db *sql.DB, tableName string) ([][]string, error) {
query := fmt.Sprintf("SELECT * FROM %s", tableName)
rows, err := db.Query(query)
if err != nil {
return nil, err
}
defer rows.Close()
columns, err := rows.Columns()
if err != nil {
return nil, err
}
output := [][]string{columns}
values := make([]sql.RawBytes, len(columns))
scanArgs := make([]interface{}, len(values))
for i := range values {
scanArgs[i] = &values[i]
}
for rows.Next() {
if err := rows.Scan(scanArgs...); err != nil {
return nil, err
}
var valueStrings []string
for _, value := range values {
if value == nil {
valueStrings = append(valueStrings, "NULL")
} else {
valueStrings = append(valueStrings, string(value))
}
}
output = append(output, valueStrings)
}
return output, nil
}