From e10debba6c15633bafb8d23c07a02998cf43f507 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Mon, 20 Feb 2017 01:18:34 +0000
Subject: [PATCH] Enable query params for USE NS / USE DB in queries

---
 sql/tokens.go |  4 ++++
 sql/use.go    | 42 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/sql/tokens.go b/sql/tokens.go
index b2d1260c..45e36cc9 100644
--- a/sql/tokens.go
+++ b/sql/tokens.go
@@ -114,6 +114,7 @@ const (
 	CONTENT
 	CREATE
 	DATABASE
+	DB
 	DEFAULT
 	DEFINE
 	DELETE
@@ -152,6 +153,7 @@ const (
 	NOT
 	NOTNULL
 	NOW
+	NS
 	NULL
 	OFFSET
 	ON
@@ -279,6 +281,7 @@ var tokens = [...]string{
 	CONTENT:         "CONTENT",
 	CREATE:          "CREATE",
 	DATABASE:        "DATABASE",
+	DB:              "DB",
 	DEFAULT:         "DEFAULT",
 	DEFINE:          "DEFINE",
 	DELETE:          "DELETE",
@@ -317,6 +320,7 @@ var tokens = [...]string{
 	NOT:             "NOT",
 	NOTNULL:         "NOTNULL",
 	NOW:             "NOW",
+	NS:              "NS",
 	NULL:            "NULL",
 	ON:              "ON",
 	OR:              "OR",
diff --git a/sql/use.go b/sql/use.go
index 384ae2ad..147f362a 100644
--- a/sql/use.go
+++ b/sql/use.go
@@ -21,18 +21,33 @@ func (p *parser) parseUseStatement() (stmt *UseStatement, err error) {
 	var tok Token
 	var exi bool
 
-	tok, _, err = p.shouldBe(NAMESPACE, DATABASE)
+	tok, _, err = p.shouldBe(NAMESPACE, NS, DATABASE, DB)
 	if err != nil {
 		return nil, err
 	}
 
 	for {
 
-		if p.is(tok, NAMESPACE) {
+		var ok bool
+		var lit string
+		var val interface{}
 
-			_, stmt.NS, err = p.shouldBe(IDENT, STRING)
+		if p.is(tok, NAMESPACE, NS) {
+
+			tok, lit, err = p.shouldBe(IDENT, STRING, PARAM)
 			if err != nil {
-				return nil, &ParseError{Found: stmt.NS, Expected: []string{"namespace name"}}
+				return nil, &ParseError{Found: lit, Expected: []string{"namespace name"}}
+			}
+
+			switch tok {
+			default:
+				val, err = p.declare(STRING, lit)
+			case PARAM:
+				val, err = p.declare(PARAM, lit)
+			}
+
+			if stmt.NS, ok = val.(string); !ok || err != nil {
+				return nil, &ParseError{Found: lit, Expected: []string{"namespace name as a STRING"}}
 			}
 
 			if err = p.o.ns(stmt.NS); err != nil {
@@ -41,11 +56,22 @@ func (p *parser) parseUseStatement() (stmt *UseStatement, err error) {
 
 		}
 
-		if p.is(tok, DATABASE) {
+		if p.is(tok, DATABASE, DB) {
 
-			_, stmt.DB, err = p.shouldBe(IDENT, DATE, TIME, STRING, NUMBER, DOUBLE)
+			tok, lit, err = p.shouldBe(IDENT, DATE, TIME, STRING, NUMBER, DOUBLE, PARAM)
 			if err != nil {
-				return nil, &ParseError{Found: stmt.DB, Expected: []string{"database name"}}
+				return nil, &ParseError{Found: lit, Expected: []string{"database name"}}
+			}
+
+			switch tok {
+			default:
+				val, err = p.declare(STRING, lit)
+			case PARAM:
+				val, err = p.declare(PARAM, lit)
+			}
+
+			if stmt.DB, ok = val.(string); !ok || err != nil {
+				return nil, &ParseError{Found: lit, Expected: []string{"database name as a STRING"}}
 			}
 
 			if err = p.o.db(stmt.DB); err != nil {
@@ -54,7 +80,7 @@ func (p *parser) parseUseStatement() (stmt *UseStatement, err error) {
 
 		}
 
-		tok, _, exi = p.mightBe(NAMESPACE, DATABASE)
+		tok, _, exi = p.mightBe(NAMESPACE, NS, DATABASE, DB)
 		if !exi {
 			break
 		}