forked from maiadegraaf/go-duckdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
duckdb.go
140 lines (115 loc) · 3.61 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
// 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", errParseConfig, err.Error())
}
config, err := prepareConfig(parsedDSN)
if err != nil {
return nil, err
}
defer C.duckdb_destroy_config(&config)
connectionString := C.CString(extractConnectionString(dataSourceName))
defer C.free(unsafe.Pointer(connectionString))
var errMsg *C.char
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", errOpen, 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, errOpen
}
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(parsedDSN *url.URL) (C.duckdb_config, error) {
var config C.duckdb_config
if state := C.duckdb_create_config(&config); state == C.DuckDBError {
return nil, errCreateConfig
}
if state := C.duckdb_set_config(config, C.CString("duckdb_api"), C.CString("go")); state == C.DuckDBError {
return nil, fmt.Errorf("%w: failed to set duckdb_api", errPrepareConfig)
}
if len(parsedDSN.RawQuery) > 0 {
for k, v := range parsedDSN.Query() {
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", errPrepareConfig, k, v[0])
}
}
}
}
return config, nil
}
var (
errOpen = errors.New("could not open database")
errParseConfig = errors.New("could not parse config for database")
errCreateConfig = errors.New("could not create config for database")
errPrepareConfig = errors.New("could not set config for database")
)