CovenantSQL · leventeliu · Apr 24, 2019 · Apr 22, 2019 · Apr 23, 2019 · Apr 23, 2019
diff --git a/client/conn.go b/client/conn.go
@@ -286,7 +286,6 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 		affectedRows: affectedRows,
 		lastInsertID: lastInsertID,
 	}
-
 	return
 }
 
@@ -429,6 +428,13 @@ func (c *conn) sendQuery(ctx context.Context, queryType types.QueryType, queries
 		return
 	}
 
+	// set receipt if key exists in context
+	if val := ctx.Value(&ctxReceiptKey); val != nil {
+		val.(*atomic.Value).Store(&Receipt{
+			RequestHash: req.Header.Hash(),
+		})
+	}
+
 	var response types.Response
 	if err = uc.pCaller.Call(route.DBSQuery.String(), req, &response); err != nil {
 		return

diff --git a/client/conn_test.go b/client/conn_test.go
@@ -17,6 +17,7 @@
 package client
 
 import (
+	"context"
 	"database/sql"
 	"sync"
 	"testing"
@@ -29,8 +30,11 @@ import (
 func TestConn(t *testing.T) {
 	log.SetLevel(log.DebugLevel)
 	Convey("test connection", t, func() {
-		var stopTestService func()
-		var err error
+		var (
+			stopTestService func()
+			ok              bool
+			err             error
+		)
 		stopTestService, _, err = startTestService()
 		So(err, ShouldBeNil)
 		defer stopTestService()
@@ -40,18 +44,35 @@ func TestConn(t *testing.T) {
 		So(db, ShouldNotBeNil)
 		So(err, ShouldBeNil)
 
-		_, err = db.Exec("create table test (test int)")
+		ctx := WithReceipt(context.Background())
+		rec, ok := GetReceipt(ctx)
+		So(ok, ShouldBeFalse)
+		So(rec, ShouldBeNil)
+
+		_, err = db.ExecContext(ctx, "create table test (test int)")
 		So(err, ShouldBeNil)
-		_, err = db.Exec("insert into test values (1)")
+		rec, ok = GetReceipt(ctx)
+		So(ok, ShouldBeTrue)
+		So(rec, ShouldNotBeNil)
+
+		_, err = db.ExecContext(ctx, "insert into test values (1)")
 		So(err, ShouldBeNil)
+		rec2, ok := GetReceipt(ctx)
+		So(ok, ShouldBeTrue)
+		So(rec2, ShouldNotBeNil)
+		So(rec, ShouldNotEqual, rec2) // receipt should be reset
 
 		// test with query
 		var rows *sql.Rows
 		var result int
-		rows, err = db.Query("select * from test")
+		rows, err = db.QueryContext(ctx, "select * from test")
 		So(err, ShouldBeNil)
 		So(rows, ShouldNotBeNil)
 		So(rows.Next(), ShouldBeTrue)
+		rec, ok = GetReceipt(ctx)
+		So(ok, ShouldBeTrue)
+		So(rec, ShouldNotBeNil)
+
 		err = rows.Scan(&result)
 		So(err, ShouldBeNil)
 		So(result, ShouldEqual, 1)

diff --git a/client/receipt.go b/client/receipt.go
@@ -0,0 +1,63 @@
+/*
+ * Copyright 2019 The CovenantSQL Authors.
+ *
+ * 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 client
+
+import (
+	"context"
+	"sync/atomic"
+
+	"github.com/CovenantSQL/CovenantSQL/crypto/hash"
+)
+
+var (
+	ctxReceiptKey = "_cql_receipt"
+)
+
+// Receipt defines a receipt of CovenantSQL query request.
+type Receipt struct {
+	RequestHash hash.Hash
+}
+
+// WithReceipt returns a context who holds a *atomic.Value. A *Receipt will be set to this value
+// after the query succeeds.
+//
+// Note that this context is safe for concurrent queries, but the value may be reset in another
+// goroutines. So if you want to make use of Receipt in several goroutines, you should call this
+// method to get separated child context in each goroutine.
+func WithReceipt(ctx context.Context) context.Context {
+	var value atomic.Value
+	value.Store((*Receipt)(nil))
+	return context.WithValue(ctx, &ctxReceiptKey, &value)
+}
+
+// GetReceipt tries to get *Receipt from context.
+func GetReceipt(ctx context.Context) (rec *Receipt, ok bool) {
+	vali := ctx.Value(&ctxReceiptKey)
+	if vali == nil {
+		return
+	}
+	value, ok := vali.(*atomic.Value)
+	if !ok {
+		return
+	}
+	reci := value.Load()
+	rec, ok = reci.(*Receipt)
+	if rec == nil {
+		ok = false
+	}
+	return
+}