From 69f8366c36ba0b707fe3f115a5d3a5528d502533 Mon Sep 17 00:00:00 2001
From: Aaron Harlap <aaron@chronosphere.io>
Date: Thu, 22 Dec 2022 12:41:28 -0500
Subject: [PATCH] Append query options to context instead of overwriting.
 (#860)

* Append query options to context instead of overwriting.

* lint.

* Address feedback.

* Use asserts.
---
 context.go      |  4 +---
 context_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 context_test.go

diff --git a/context.go b/context.go
index 588382edf9..bb56c6e14c 100644
--- a/context.go
+++ b/context.go
@@ -132,9 +132,7 @@ func WithStdAsync(wait bool) QueryOption {
 }
 
 func Context(parent context.Context, options ...QueryOption) context.Context {
-	opt := QueryOptions{
-		settings: make(Settings),
-	}
+	opt := queryOptions(parent)
 	for _, f := range options {
 		f(&opt)
 	}
diff --git a/context_test.go b/context_test.go
new file mode 100644
index 0000000000..9e087fa47d
--- /dev/null
+++ b/context_test.go
@@ -0,0 +1,35 @@
+package clickhouse
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestContext(t *testing.T) {
+	t.Run("call context multiple times making sure query options are persisted across calls",
+		func(t *testing.T) {
+			ctx := Context(context.Background(), WithQueryID("a"))
+			ctx = Context(ctx, WithQuotaKey("b"))
+			ctx = Context(ctx, WithSettings(Settings{
+				"c": "d",
+			}))
+
+			opts := queryOptions(ctx)
+			assert.Equal(t, "a", opts.queryID)
+			assert.Equal(t, "b", opts.quotaKey)
+			assert.Equal(t, "d", opts.settings["c"])
+		},
+	)
+
+	t.Run("call context multiple times making sure query options are persisted across calls",
+		func(t *testing.T) {
+			ctx := Context(context.Background(), WithQueryID("a"))
+			ctx = Context(ctx, WithQueryID("b"))
+
+			opts := queryOptions(ctx)
+			assert.Equal(t, "b", opts.queryID)
+		},
+	)
+}