From 4849eb44867495fc8fc79ba06caf2e875eb8c921 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Thu, 15 Apr 2021 23:02:35 +0100
Subject: [PATCH] Enable query params to be defined on a connection

---
 db/db.go    | 10 ++++++++++
 web/rpc.go  | 15 +++++++++++++++
 web/sess.go | 29 +++++++++++++++++++++++++++++
 web/vars.go |  1 +
 web/web.go  |  4 ++++
 5 files changed, 59 insertions(+)
 create mode 100644 web/sess.go

diff --git a/db/db.go b/db/db.go
index f7099b51..2c8330d3 100644
--- a/db/db.go
+++ b/db/db.go
@@ -138,6 +138,16 @@ func Process(fib *fibre.Context, ast *sql.Query, vars map[string]interface{}) (o
 		vars = make(map[string]interface{})
 	}
 
+	// Fetch any variables which have been set
+	// on the connection. These might have been
+	// set on the WebSocket or using HTTP Headers.
+
+	if sess := fib.Get(ctxKeyVars); sess != nil {
+		for key, val := range sess.(map[string]interface{}) {
+			vars[key] = val
+		}
+	}
+
 	// Get the unique id for this connection
 	// so that we can assign it to the context
 	// and detect any websocket notifications.
diff --git a/web/rpc.go b/web/rpc.go
index 9737977a..3b817425 100644
--- a/web/rpc.go
+++ b/web/rpc.go
@@ -71,6 +71,21 @@ func (r *rpc) Live(c *fibre.Context, class string) (interface{}, error) {
 // Methods for static queries
 // --------------------------------------------------
 
+func (r *rpc) Let(c *fibre.Context, key string, val interface{}) (interface{}, error) {
+	switch val := val.(type) {
+	case *fibre.RPCNull:
+		vars := c.Get("vars").(map[string]interface{})
+		delete(vars, key)
+		c.Set("vars", vars)
+		return vars, nil
+	default:
+		vars := c.Get("vars").(map[string]interface{})
+		vars[key] = val
+		c.Set("vars", vars)
+		return vars, nil
+	}
+}
+
 func (r *rpc) Query(c *fibre.Context, sql string, vars map[string]interface{}) (interface{}, error) {
 	return db.Execute(c, sql, vars)
 }
diff --git a/web/sess.go b/web/sess.go
new file mode 100644
index 00000000..a3c6aca1
--- /dev/null
+++ b/web/sess.go
@@ -0,0 +1,29 @@
+// Copyright © 2016 Abcum Ltd
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package web
+
+import (
+	"github.com/abcum/fibre"
+)
+
+func sess() fibre.MiddlewareFunc {
+	return func(h fibre.HandlerFunc) fibre.HandlerFunc {
+		return func(c *fibre.Context) (err error) {
+			vars := make(map[string]interface{})
+			c.Set(varKeyVars, vars)
+			return h(c)
+		}
+	}
+}
diff --git a/web/vars.go b/web/vars.go
index 3d9df8a9..9382409e 100644
--- a/web/vars.go
+++ b/web/vars.go
@@ -23,6 +23,7 @@ const (
 	varKeyTb   = "TB"
 	varKeyId   = "ID"
 	varKeyAuth = "auth"
+	varKeyVars = "vars"
 	varKeyUser = "user"
 	varKeyPass = "pass"
 	varKeyUniq = "uniq"
diff --git a/web/web.go b/web/web.go
index e47f5ece..1368c4ba 100644
--- a/web/web.go
+++ b/web/web.go
@@ -76,6 +76,10 @@ func Setup(opts *cnf.Options) (err error) {
 
 	// Setup authentication
 
+	s.Use(sess())
+
+	// Setup authentication
+
 	s.Use(auth())
 
 	// Setup live queries