-
Notifications
You must be signed in to change notification settings - Fork 566
/
util.go
101 lines (92 loc) · 3.76 KB
/
util.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
package pachsql
import (
"database/sql"
"fmt"
"reflect"
"strconv"
"strings"
"time"
"github.com/pachyderm/pachyderm/v2/src/internal/errors"
)
// Placeholder returns a placeholder for the given driver
// assuming i placeholders have been provided before it. It is 0 indexed in this way.
//
// This means that the first Postgres placeholder is `$1` when i = 0.
// This is more ergonimic for contructing a list of arguments since i = len(args)
// but is perhaps unintuitive for those familiar with Postgres.
func Placeholder(driverName string, i int) string {
switch driverName {
case "pgx":
return "$" + strconv.Itoa(i+1)
case "mysql", "snowflake":
return "?"
default:
panic(driverName)
}
}
// SplitTableSchema splits the tablePath on the first . and interprets the first part
// as the schema, if the driver supports schemas.
func SplitTableSchema(driver string, tablePath string) (schemaName string, tableName string) {
if parts := strings.SplitN(tablePath, ".", 2); len(parts) == 2 {
schemaName = parts[0]
tableName = parts[1]
} else {
tableName = tablePath
switch driver {
case "mysql":
case "pgx", "snowflake":
schemaName = "public"
default:
panic(fmt.Sprintf("driver not supported: %s", driver))
}
}
return schemaName, tableName
}
// TestRow is the type of a row in the test table
// struct tag: sql:"<column_name>,<data_type>,<table_constraint>"
type TestRow struct {
Id int16 `column:"c_id" dtype:"SMALLINT" constraint:"PRIMARY KEY NOT NULL"`
Smallint int16 `column:"c_smallint" dtype:"SMALLINT" constraint:"NOT NULL"`
Int int32 `column:"c_int" dtype:"INT" constraint:"NOT NULL"`
Bigint int64 `column:"c_bigint" dtype:"BIGINT" constraint:"NOT NULL"`
Float float32 `column:"c_float" dtype:"FLOAT" constraint:"NOT NULL"`
NumericInt int64 `column:"c_numeric_int" dtype:"NUMERIC(20,0)" constraint:"NOT NULL"`
NumericFloat float64 `column:"c_numeric_float" dtype:"NUMERIC(20,19)" constraint:"NOT NULL"`
Varchar string `column:"c_varchar" dtype:"VARCHAR(100)" constraint:"NOT NULL"`
Time time.Time `column:"c_time" dtype:"TIMESTAMP" constraint:"NOT NULL"`
SmallintNull sql.NullInt16 `column:"c_smallint_null" dtype:"SMALLINT" constraint:"NULL"`
IntNull sql.NullInt32 `column:"c_int_null" dtype:"INT" constraint:"NULL"`
BigintNull sql.NullInt64 `column:"c_bigint_null" dtype:"BIGINT" constraint:"NULL"`
FloatNull sql.NullFloat64 `column:"c_float_null" dtype:"FLOAT" constraint:"NULL"`
NumericIntNull sql.NullInt64 `column:"c_numeric_int_null" dtype:"NUMERIC(20,0)" constraint:"NULL"`
NumericFloatNull sql.NullFloat64 `column:"c_numeric_float_null" dtype:"NUMERIC(20,19)" constraint:"NULL"`
VarcharNull sql.NullString `column:"c_varchar_null" dtype:"VARCHAR(100)" constraint:"NULL"`
TimeNull sql.NullTime `column:"c_time_null" dtype:"TIMESTAMP" constraint:"NULL"`
}
func (row *TestRow) SetID(id int16) {
row.Id = id
}
// CreateTestTable creates a test table at name in the database
func CreateTestTable(db *DB, name string, schema interface{}) error {
t := reflect.TypeOf(schema)
var processFields func(reflect.Type) []string
processFields = func(t reflect.Type) []string {
var cols []string
if t.Kind() == reflect.Ptr {
return processFields(t.Elem())
}
for i := 0; i < t.NumField(); i++ {
field := t.Field(i)
switch {
case field.Anonymous && field.Type.Kind() == reflect.Struct:
cols = append(cols, processFields(field.Type)...)
default:
cols = append(cols, fmt.Sprintf("%s %s %s", field.Tag.Get("column"), field.Tag.Get("dtype"), field.Tag.Get("constraint")))
}
}
return cols
}
q := fmt.Sprintf(`CREATE TABLE %s (%s)`, name, strings.Join(processFields(t), ", "))
_, err := db.Exec(q)
return errors.EnsureStack(err)
}