forked from marcboeker/go-duckdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
duckdb.go
144 lines (118 loc) · 3.69 KB
/
duckdb.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
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
// Package duckdb implements a database/sql driver for the DuckDB database.
package duckdb
/*
#include <duckdb.h>
*/
import "C"
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"fmt"
"net/url"
"strings"
"unsafe"
)
func init() {
sql.Register("duckdb", Driver{})
}
type Driver struct{}
func (d Driver) Open(dataSourceName string) (driver.Conn, error) {
connector, err := d.OpenConnector(dataSourceName)
if err != nil {
return nil, err
}
return connector.Connect(context.Background())
}
func (Driver) OpenConnector(dataSourceName string) (driver.Connector, error) {
return createConnector(dataSourceName, func(execerContext driver.ExecerContext) error { return nil })
}
// NewConnector creates a new Connector for the DuckDB database.
func NewConnector(dsn string, connInitFn func(execer driver.ExecerContext) error) (driver.Connector, error) {
return createConnector(dsn, connInitFn)
}
func createConnector(dataSourceName string, connInitFn func(execer driver.ExecerContext) error) (driver.Connector, error) {
var db C.duckdb_database
parsedDSN, err := url.Parse(dataSourceName)
if err != nil {
return nil, fmt.Errorf("%w: %s", parseConfigError, err.Error())
}
connectionString := C.CString(extractConnectionString(dataSourceName))
defer C.free(unsafe.Pointer(connectionString))
// Check for config options.
if len(parsedDSN.RawQuery) == 0 {
errMsg := C.CString("")
defer C.duckdb_free(unsafe.Pointer(errMsg))
if state := C.duckdb_open_ext(connectionString, &db, nil, &errMsg); state == C.DuckDBError {
return nil, fmt.Errorf("%w: %s", openError, C.GoString(errMsg))
}
} else {
config, err := prepareConfig(parsedDSN.Query())
if err != nil {
return nil, err
}
errMsg := C.CString("")
defer C.duckdb_free(unsafe.Pointer(errMsg))
if state := C.duckdb_open_ext(connectionString, &db, config, &errMsg); state == C.DuckDBError {
return nil, fmt.Errorf("%w: %s", openError, C.GoString(errMsg))
}
}
return &connector{db: &db, connInitFn: connInitFn}, nil
}
type connector struct {
db *C.duckdb_database
connInitFn func(execer driver.ExecerContext) error
}
func (c *connector) Driver() driver.Driver {
return Driver{}
}
func (c *connector) Connect(context.Context) (driver.Conn, error) {
var con C.duckdb_connection
if state := C.duckdb_connect(*c.db, &con); state == C.DuckDBError {
return nil, openError
}
conn := &conn{con: &con}
// Call the connection init function if defined
if c.connInitFn != nil {
if err := c.connInitFn(conn); err != nil {
return nil, err
}
}
return conn, nil
}
func (c *connector) Close() error {
C.duckdb_close(c.db)
c.db = nil
return nil
}
func extractConnectionString(dataSourceName string) string {
var queryIndex = strings.Index(dataSourceName, "?")
if queryIndex < 0 {
queryIndex = len(dataSourceName)
}
return dataSourceName[0:queryIndex]
}
func prepareConfig(options map[string][]string) (C.duckdb_config, error) {
var config C.duckdb_config
if state := C.duckdb_create_config(&config); state == C.DuckDBError {
return nil, createConfigError
}
for k, v := range options {
if len(v) > 0 {
state := C.duckdb_set_config(config, C.CString(k), C.CString(v[0]))
if state == C.DuckDBError {
return nil, fmt.Errorf("%w: affected config option %s=%s", prepareConfigError, k, v[0])
}
}
}
return config, nil
}
var (
openError = errors.New("could not open database")
parseConfigError = errors.New("could not parse config for database")
createConfigError = errors.New("could not create config for database")
prepareConfigError = errors.New("could not set config for database")
)