@@ -26,6 +26,7 @@ package drivermgr
2626// #define ADBC_EXPORTING
2727// #endif
2828// #include "arrow-adbc/adbc.h"
29+ // #include "arrow-adbc/adbc_driver_manager.h"
2930// #include <stdlib.h>
3031// #include <string.h>
3132//
@@ -50,6 +51,7 @@ package drivermgr
5051import "C"
5152import (
5253 "context"
54+ "strconv"
5355 "sync"
5456 "unsafe"
5557
@@ -59,6 +61,18 @@ import (
5961 "github.com/apache/arrow-go/v18/arrow/cdata"
6062)
6163
64+ const (
65+ LoadFlagsSearchEnv = 1 << iota
66+ LoadFlagsSearchPath
67+ LoadFlagsSearchSystem
68+ LoadFlagsAllowRelativePaths
69+
70+ LoadFlagsDefault = LoadFlagsSearchEnv | LoadFlagsSearchPath | LoadFlagsSearchSystem | LoadFlagsAllowRelativePaths
71+ // LoadFlagsOptionKey is the key to use for an option to set specific
72+ // load flags for the database to decide where to look for driver manifests.
73+ LoadFlagsOptionKey = "load_flags"
74+ )
75+
6276type option struct {
6377 key , val * C.char
6478}
@@ -93,18 +107,54 @@ func (d Driver) NewDatabaseWithContext(_ context.Context, opts map[string]string
93107 options : make (map [string ]option ),
94108 }
95109
110+ defer func () {
111+ if db .db == nil { // cleanup options if we failed to create the database
112+ for _ , o := range dbOptions {
113+ C .free (unsafe .Pointer (o .key ))
114+ C .free (unsafe .Pointer (o .val ))
115+ }
116+ }
117+ }()
118+
96119 var err C.struct_AdbcError
97120 db .db = (* C .struct_AdbcDatabase )(unsafe .Pointer (C .calloc (C .sizeof_struct_AdbcDatabase , C .size_t (1 ))))
98121 if code := adbc .Status (C .AdbcDatabaseNew (db .db , & err )); code != adbc .StatusOK {
99122 return nil , toAdbcError (code , & err )
100123 }
101124
102- for _ , o := range dbOptions {
103- if code := adbc .Status (C .AdbcDatabaseSetOption (db .db , o .key , o .val , & err )); code != adbc .StatusOK {
104- errOut := toAdbcError (code , & err )
105- C .AdbcDatabaseRelease (db .db , & err )
106- db .db = nil
107- return nil , errOut
125+ if code := adbc .Status (C .AdbcDriverManagerDatabaseSetLoadFlags (db .db , C .AdbcLoadFlags (LoadFlagsDefault ), & err )); code != adbc .StatusOK {
126+ errOut := toAdbcError (code , & err )
127+ C .AdbcDatabaseRelease (db .db , & err )
128+ db .db = nil
129+ return nil , errOut
130+ }
131+
132+ for k , o := range dbOptions {
133+ switch k {
134+ case LoadFlagsOptionKey :
135+ f , errOut := strconv .Atoi (C .GoString (o .val ))
136+ if errOut != nil {
137+ C .AdbcDatabaseRelease (db .db , & err )
138+ db .db = nil
139+ return nil , adbc.Error {
140+ Code : adbc .StatusInvalidArgument ,
141+ Msg : "invalid load flags value: " + C .GoString (o .val ),
142+ }
143+ }
144+
145+ if code := adbc .Status (C .AdbcDriverManagerDatabaseSetLoadFlags (db .db , C .AdbcLoadFlags (f ), & err )); code != adbc .StatusOK {
146+ errOut := toAdbcError (code , & err )
147+ C .AdbcDatabaseRelease (db .db , & err )
148+ db .db = nil
149+ return nil , errOut
150+ }
151+ default :
152+ if code := adbc .Status (C .AdbcDatabaseSetOption (db .db , o .key , o .val , & err )); code != adbc .StatusOK {
153+ errOut := toAdbcError (code , & err )
154+ C .AdbcDatabaseRelease (db .db , & err )
155+ db .db = nil
156+ return nil , errOut
157+ }
108158 }
109159 }
110160
0 commit comments