-
Notifications
You must be signed in to change notification settings - Fork 5
/
con_batch.go
115 lines (101 loc) · 2.58 KB
/
con_batch.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
package godatabend
import (
"bufio"
"context"
"database/sql/driver"
"encoding/csv"
"fmt"
"os"
"path/filepath"
"regexp"
"time"
"github.com/google/uuid"
"github.com/pkg/errors"
)
// \x60 represents a backtick
var httpInsertRe = regexp.MustCompile(`(?i)^INSERT INTO\s+\x60?([\w.^\(]+)\x60?\s*(\([^\)]*\))? VALUES`)
type Batch interface {
AppendToFile(v []driver.Value) error
BatchInsert() error
}
func (dc *DatabendConn) prepareBatch(ctx context.Context, query string) (Batch, error) {
matches := httpInsertRe.FindStringSubmatch(query)
if len(matches) < 2 {
return nil, errors.New("cannot get table name from query")
}
csvFileName := fmt.Sprintf("%s/%s.csv", os.TempDir(), uuid.NewString())
csvFile, err := os.OpenFile(csvFileName, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
if err != nil {
return nil, err
}
defer csvFile.Close()
writer := csv.NewWriter(csvFile)
writer.Flush()
return &httpBatch{
query: query,
ctx: ctx,
conn: dc,
batchFile: csvFileName,
}, nil
}
type httpBatch struct {
query string
ctx context.Context
conn *DatabendConn
batchFile string
err error
}
func (b *httpBatch) BatchInsert() error {
defer func() {
err := os.RemoveAll(b.batchFile)
if err != nil {
b.conn.log("delete batch insert file failed: ", err)
}
}()
stage, err := b.UploadToStage(context.Background())
if err != nil {
return errors.Wrap(err, "upload to stage failed")
}
_, err = b.conn.rest.InsertWithStage(b.ctx, b.query, stage, nil, nil)
if err != nil {
return errors.Wrap(err, "insert with stage failed")
}
return nil
}
func (b *httpBatch) AppendToFile(v []driver.Value) error {
csvFile, err := os.OpenFile(b.batchFile, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
if err != nil {
return err
}
defer csvFile.Close()
lineData := make([]string, 0, len(v))
for i := range v {
lineData = append(lineData, fmt.Sprintf("%v", v[i]))
}
writer := csv.NewWriter(csvFile)
err = writer.Write(lineData)
if err != nil {
return err
}
writer.Flush()
return nil
}
func (b *httpBatch) UploadToStage(ctx context.Context) (*StageLocation, error) {
ctx = checkQueryID(ctx)
fi, err := os.Stat(b.batchFile)
if err != nil {
return nil, errors.Wrap(err, "get batch file size failed")
}
size := fi.Size()
f, err := os.Open(b.batchFile)
if err != nil {
return nil, errors.Wrap(err, "open batch file failed")
}
defer f.Close()
input := bufio.NewReader(f)
stage := &StageLocation{
Name: "~",
Path: fmt.Sprintf("batch/%d-%s", time.Now().Unix(), filepath.Base(b.batchFile)),
}
return stage, b.conn.rest.UploadToStage(ctx, stage, input, size)
}