Reservation support

README.md

@@ -36,6 +36,14 @@ func main() {
 }
 ```
 
+You can attach a reservation to the DSN via a query parameter:
+
+```
+bigquery://projectid/dataset?reservation=projects/<project>/locations/<location>/reservations/<reservation>
+```
+
+When a reservation is set, the driver prefixes each query with a `SET @@reservation` statement and uses multi-statement scripting.
+
 ## Gorm Usage
 
 Opening a Gorm bigquery db

driver/driver.go

@@ -24,6 +24,7 @@ type bigQueryConfig struct {
 	disableAuth    bool
 	credentialFile string
 	credentialJSON []byte
+	reservation    string
 }
 
 func (b bigQueryDriver) Open(uri string) (driver.Conn, error) {
@@ -90,6 +91,7 @@ func configFromUri(uri string) (*bigQueryConfig, error) {
 		endpoint:       u.Query().Get("endpoint"),
 		disableAuth:    u.Query().Get("disable_auth") == "true",
 		credentialFile: u.Query().Get("credential_file"),
+		reservation:    getReservation(u.Query()),
 	}
 
 	if u.Query().Get("credential_json") != "" {
@@ -116,6 +118,13 @@ func getScopes(query url.Values) []string {
 	return strings.Split(q, ",")
 }
 
+func getReservation(query url.Values) string {
+	if reservation := query.Get("reservation"); reservation != "" {
+		return reservation
+	}
+	return query.Get("reservation_id")
+}
+
 func invalidConnectionStringError(uri string) error {
 	return fmt.Errorf("invalid connection string: %s", uri)
 }

driver/statement.go

@@ -5,6 +5,9 @@ import (
 	"context"
 	"database/sql/driver"
 	"errors"
+	"fmt"
+	"strings"
+
 	"github.com/sirupsen/logrus"
 	"gorm.io/driver/bigquery/adaptor"
 )
@@ -144,7 +147,8 @@ func (statement bigQueryStatement) Query(args []driver.Value) (driver.Rows, erro
 
 func (statement bigQueryStatement) buildQuery(args []driver.Value) (*bigquery.Query, error) {
 
-	query, err := statement.connection.query(statement.query)
+	queryText := applyReservation(statement.query, statement.connection.config.reservation)
+	query, err := statement.connection.query(queryText)
 	if err != nil {
 		return nil, err
 	}
@@ -157,6 +161,19 @@ func (statement bigQueryStatement) buildQuery(args []driver.Value) (*bigquery.Qu
 	return query, err
 }
 
+func applyReservation(query, reservation string) string {
+	if reservation == "" {
+		return query
+	}
+	trimmed := strings.TrimSpace(query)
+	upper := strings.ToUpper(trimmed)
+	if strings.HasPrefix(upper, "SET @@RESERVATION=") || strings.HasPrefix(upper, "SET @@RESERVATION ") {
+		return query
+	}
+	escaped := strings.ReplaceAll(reservation, "'", "''")
+	return fmt.Sprintf("SET @@reservation='%s';\n%s", escaped, query)
+}
+
 func (statement bigQueryStatement) buildParameters(args []driver.Value) ([]bigquery.QueryParameter, error) {
 	if args == nil {
 		return nil, nil

test/common_text.go

@@ -6,6 +6,7 @@ import (
 	"gorm.io/driver/bigquery"
 	"gorm.io/gorm"
 	"log"
+	"os"
 )
 
 type GormTestSuite struct {
@@ -18,7 +19,11 @@ func (suite *GormTestSuite) SetupSuite() {
 	logrus.SetLevel(logrus.DebugLevel)
 
 	var err error
-	suite.db, err = gorm.Open(bigquery.Open("bigquery://go-bigquery-driver/playground"), &gorm.Config{})
+	dsn := os.Getenv("BIGQUERY_TEST_DSN")
+	if dsn == "" {
+		dsn = "bigquery://go-bigquery-driver/playground"
+	}
+	suite.db, err = gorm.Open(bigquery.Open(dsn), &gorm.Config{})
 	if err != nil {
 		log.Fatal(err)
 	}