Support user specified batch in DSN

driver.go

@@ -64,7 +64,7 @@ func (d drv) Open(dsn string) (driver.Conn, error) {
 		return nil, err
 	}
 
-	options := hiveOptions{PollIntervalSeconds: 5, BatchSize: 100000}
+	options := hiveOptions{PollIntervalSeconds: 5, BatchSize: int64(cfg.Batch)}
 	conn := &hiveConnection{
 		thrift:  client,
 		session: session.SessionHandle,

dsn.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/url"
 	"regexp"
+	"strconv"
 	"strings"
 )
 
@@ -13,6 +14,7 @@ type Config struct {
 	Addr       string
 	DBName     string
 	Auth       string
+	Batch      int
 	SessionCfg map[string]string
 }
 
@@ -22,7 +24,13 @@ var (
 	reUserPasswd = regexp.MustCompile(`([^:@]+)(:[^:@]+)?@`)
 )
 
-const sessionConfPrefix = "session."
+const (
+	sessionConfPrefix = "session."
+	authConfName      = "auth"
+	defaultAuth       = "NOSASL"
+	batchSizeName     = "batch"
+	defaultBatchSize  = 10000
+)
 
 // ParseDSN requires DSN names in the format [user[:password]@]addr/dbname.
 func ParseDSN(dsn string) (*Config, error) {
@@ -50,13 +58,23 @@ func ParseDSN(dsn string) (*Config, error) {
 		}
 	}
 
-	auth := "NOSASL"
+	auth := defaultAuth
+	batch := defaultBatchSize
 	sc := make(map[string]string)
 	if len(sub[3]) > 0 && sub[3][0] == '?' {
 		qry, _ := url.ParseQuery(sub[3][1:])
-		if v, found := qry["auth"]; found {
+
+		if v, found := qry[authConfName]; found {
 			auth = v[0]
 		}
+		if v, found := qry[batchSizeName]; found {
+			bch, err := strconv.Atoi(v[0])
+			if err != nil {
+				return nil, err
+			}
+			batch = bch
+		}
+
 		for k, v := range qry {
 			if strings.HasPrefix(k, sessionConfPrefix) {
 				sc[k[len(sessionConfPrefix):]] = v[0]
@@ -70,6 +88,7 @@ func ParseDSN(dsn string) (*Config, error) {
 		Addr:       addr,
 		DBName:     dbname,
 		Auth:       auth,
+		Batch:      batch,
 		SessionCfg: sc,
 	}, nil
 }
@@ -80,22 +99,13 @@ func (cfg *Config) FormatDSN() string {
 	if len(cfg.DBName) > 0 {
 		dsn = fmt.Sprintf("%s/%s", dsn, cfg.DBName)
 	}
-	queryExisted := false
+	dsn += fmt.Sprintf("?batch=%d", cfg.Batch)
 	if len(cfg.Auth) > 0 {
-		dsn = fmt.Sprintf("%s?auth=%s", dsn, cfg.Auth)
-		queryExisted = true
+		dsn += fmt.Sprintf("&auth=%s", cfg.Auth)
 	}
 	if len(cfg.SessionCfg) > 0 {
-		if !queryExisted {
-			dsn += "?"
-		}
 		for k, v := range cfg.SessionCfg {
-			if !queryExisted {
-				dsn += fmt.Sprintf("%s%s=%s", sessionConfPrefix, k, v)
-			} else {
-				dsn += fmt.Sprintf("&%s%s=%s", sessionConfPrefix, k, v)
-			}
-			queryExisted = true
+			dsn += fmt.Sprintf("&%s%s=%s", sessionConfPrefix, k, v)
 		}
 	}
 	return dsn

dsn_test.go

@@ -15,10 +15,11 @@ func TestParseDSNWithSessionConf(t *testing.T) {
 		Addr:       "hiveserver",
 		DBName:     "mydb",
 		Auth:       "PLAIN",
+		Batch:      200,
 		SessionCfg: sc,
 	}
 	dsn := cfg.FormatDSN()
-	assert.Equal(t, dsn, "usr:pswd@hiveserver/mydb?auth=PLAIN&session.mapreduce_job_quenename=mr")
+	assert.Equal(t, dsn, "usr:pswd@hiveserver/mydb?batch=200&auth=PLAIN&session.mapreduce_job_quenename=mr")
 
 	cfg2, e := ParseDSN(dsn)
 	assert.Nil(t, e)
@@ -27,6 +28,7 @@ func TestParseDSNWithSessionConf(t *testing.T) {
 	assert.Equal(t, cfg.Addr, cfg2.Addr)
 	assert.Equal(t, cfg.DBName, cfg2.DBName)
 	assert.Equal(t, cfg.Auth, cfg2.Auth)
+	assert.Equal(t, cfg.Batch, cfg2.Batch)
 	sc, sc2 := cfg.SessionCfg, cfg2.SessionCfg
 	assert.Equal(t, len(sc), len(sc2))
 	for k, v := range sc {
@@ -44,6 +46,7 @@ func TestParseDSNWithAuth(t *testing.T) {
 	assert.Equal(t, cfg.Addr, "127.0.0.1")
 	assert.Equal(t, cfg.DBName, "mnist")
 	assert.Equal(t, cfg.Auth, "PLAIN")
+	assert.Equal(t, cfg.Batch, 10000)
 
 	cfg, e = ParseDSN("root@127.0.0.1/mnist")
 	assert.Nil(t, e)
@@ -98,7 +101,7 @@ func TestParseDSNWithoutDBName(t *testing.T) {
 }
 
 func TestFormatDSNWithDBName(t *testing.T) {
-	ds := "user:passwd@127.0.0.1/mnist?auth=NOSASL"
+	ds := "user:passwd@127.0.0.1/mnist?batch=100000&auth=NOSASL"
 	cfg, e := ParseDSN(ds)
 	assert.Nil(t, e)
 
@@ -107,7 +110,7 @@ func TestFormatDSNWithDBName(t *testing.T) {
 }
 
 func TestFormatDSNWithoutDBName(t *testing.T) {
-	ds := "user:passwd@127.0.0.1?auth=NOSASL"
+	ds := "user:passwd@127.0.0.1?batch=100&auth=NOSASL"
 	cfg, e := ParseDSN(ds)
 	assert.Nil(t, e)
 

rows.go

@@ -51,7 +51,9 @@ func (r *rowSet) Next(dest []driver.Value) error {
 	// First execution or reach the end of the current result set.
 	if r.resultSet == nil || r.offset >= len(r.resultSet[0]) {
 		r.offset = 0
-		r.batchFetch()
+		if err := r.batchFetch(); err != nil {
+			return err
+		}
 	}
 
 	if len(r.resultSet) <= 0 {