From 2eac8ed64fad571f6b288d337716af61e44c9e7a Mon Sep 17 00:00:00 2001
From: Dominic Wong <domwongemail@googlemail.com>
Date: Tue, 18 Aug 2020 18:30:05 +0100
Subject: [PATCH] Fix cockroach store not respecting WriteTTL option (#1943)

* cockroach fixes for expiry

* cockroach should run in the background
---
 .github/workflows/tests.yml       |  5 ++-
 store/cockroach/cockroach.go      |  3 ++
 store/cockroach/cockroach_test.go | 51 +++++++++++++++++++++++++------
 3 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 8b108ca8..d24a5ecc 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -25,7 +25,10 @@ jobs:
       id: tests
       env:
         IN_TRAVIS_CI: yes
-      run: go test -v ./...
+      run: |
+        wget -qO- https://binaries.cockroachdb.com/cockroach-v20.1.4.linux-amd64.tgz | tar  xvz
+        cockroach-v20.1.4.linux-amd64/cockroach start-single-node --insecure &
+        go test -v ./...
   
     - name: Notify of test failure
       if: failure()
diff --git a/store/cockroach/cockroach.go b/store/cockroach/cockroach.go
index 9ab13f58..2b2f3987 100644
--- a/store/cockroach/cockroach.go
+++ b/store/cockroach/cockroach.go
@@ -429,6 +429,9 @@ func (s *sqlStore) Write(r *store.Record, opts ...store.WriteOption) error {
 	} else if r.Expiry != 0 {
 		expiry = time.Now().Add(r.Expiry)
 	}
+	if options.TTL != 0 {
+		expiry = time.Now().Add(options.TTL)
+	}
 
 	if expiry.IsZero() {
 		_, err = st.Exec(r.Key, r.Value, metadata, nil)
diff --git a/store/cockroach/cockroach_test.go b/store/cockroach/cockroach_test.go
index 1c7eabf5..2b794fb1 100644
--- a/store/cockroach/cockroach_test.go
+++ b/store/cockroach/cockroach_test.go
@@ -3,7 +3,6 @@ package cockroach
 import (
 	"database/sql"
 	"fmt"
-	"os"
 	"testing"
 	"time"
 
@@ -12,10 +11,6 @@ import (
 )
 
 func TestSQL(t *testing.T) {
-	if len(os.Getenv("IN_TRAVIS_CI")) != 0 {
-		t.Skip()
-	}
-
 	connection := fmt.Sprintf(
 		"host=%s port=%d user=%s sslmode=disable dbname=%s",
 		"localhost",
@@ -109,9 +104,9 @@ func TestSQL(t *testing.T) {
 		break
 	}
 	sqlStore.Delete("bar")
-	sqlStore.Write(&store.Record{Key: "aaa", Value: []byte("bbb"), Expiry: 5 * time.Second})
-	sqlStore.Write(&store.Record{Key: "aaaa", Value: []byte("bbb"), Expiry: 5 * time.Second})
-	sqlStore.Write(&store.Record{Key: "aaaaa", Value: []byte("bbb"), Expiry: 5 * time.Second})
+	sqlStore.Write(&store.Record{Key: "aaa", Value: []byte("bbb"), Expiry: 10 * time.Second})
+	sqlStore.Write(&store.Record{Key: "aaaa", Value: []byte("bbb"), Expiry: 10 * time.Second})
+	sqlStore.Write(&store.Record{Key: "aaaaa", Value: []byte("bbb"), Expiry: 10 * time.Second})
 	results, err := sqlStore.Read("a", store.ReadPrefix())
 	if err != nil {
 		t.Error(err)
@@ -119,7 +114,7 @@ func TestSQL(t *testing.T) {
 	if len(results) != 3 {
 		t.Fatal("Results should have returned 3 records")
 	}
-	time.Sleep(6 * time.Second)
+	time.Sleep(10 * time.Second)
 	results, err = sqlStore.Read("a", store.ReadPrefix())
 	if err != nil {
 		t.Error(err)
@@ -127,4 +122,42 @@ func TestSQL(t *testing.T) {
 	if len(results) != 0 {
 		t.Fatal("Results should have returned 0 records")
 	}
+
+	sqlStore.Write(&store.Record{Key: "bbb", Value: []byte("bbb")}, store.WriteExpiry(time.Now().Add(10*time.Second)))
+	sqlStore.Write(&store.Record{Key: "bbbb", Value: []byte("bbb")}, store.WriteExpiry(time.Now().Add(10*time.Second)))
+	sqlStore.Write(&store.Record{Key: "bbbbb", Value: []byte("bbb")}, store.WriteExpiry(time.Now().Add(10*time.Second)))
+	results, err = sqlStore.Read("b", store.ReadPrefix())
+	if err != nil {
+		t.Error(err)
+	}
+	if len(results) != 3 {
+		t.Fatalf("Results should have returned 3 records. Received %d", len(results))
+	}
+	time.Sleep(10 * time.Second)
+	results, err = sqlStore.Read("b", store.ReadPrefix())
+	if err != nil {
+		t.Error(err)
+	}
+	if len(results) != 0 {
+		t.Fatalf("Results should have returned 0 records. Received %d", len(results))
+	}
+
+	sqlStore.Write(&store.Record{Key: "ccc", Value: []byte("bbb")}, store.WriteTTL(10*time.Second))
+	sqlStore.Write(&store.Record{Key: "cccc", Value: []byte("bbb")}, store.WriteTTL(10*time.Second))
+	sqlStore.Write(&store.Record{Key: "ccccc", Value: []byte("bbb")}, store.WriteTTL(10*time.Second))
+	results, err = sqlStore.Read("c", store.ReadPrefix())
+	if err != nil {
+		t.Error(err)
+	}
+	if len(results) != 3 {
+		t.Fatalf("Results should have returned 3 records. Received %d", len(results))
+	}
+	time.Sleep(10 * time.Second)
+	results, err = sqlStore.Read("c", store.ReadPrefix())
+	if err != nil {
+		t.Error(err)
+	}
+	if len(results) != 0 {
+		t.Fatalf("Results should have returned 0 records. Received %d", len(results))
+	}
 }