Merge pull request #740 from Unix4ever/set-default-timestamp-batch

Implemented WithTimestamp function in gocql.Batch

frame.go

@@ -1414,8 +1414,9 @@ type writeBatchFrame struct {
 	consistency Consistency
 
 	// v3+
-	serialConsistency SerialConsistency
-	defaultTimestamp  bool
+	serialConsistency     SerialConsistency
+	defaultTimestamp      bool
+	defaultTimestampValue int64
 }
 
 func (w *writeBatchFrame) writeFrame(framer *framer, streamID int) error {
@@ -1469,9 +1470,15 @@ func (f *framer) writeBatchFrame(streamID int, w *writeBatchFrame) error {
 		if w.serialConsistency > 0 {
 			f.writeConsistency(Consistency(w.serialConsistency))
 		}
+
 		if w.defaultTimestamp {
-			now := time.Now().UnixNano() / 1000
-			f.writeLong(now)
+			var ts int64
+			if w.defaultTimestampValue != 0 {
+				ts = w.defaultTimestampValue
+			} else {
+				ts = time.Now().UnixNano() / 1000
+			}
+			f.writeLong(ts)
 		}
 	}
 

session.go

@@ -1162,15 +1162,16 @@ func (n *nextIter) fetch() *Iter {
 }
 
 type Batch struct {
-	Type             BatchType
-	Entries          []BatchEntry
-	Cons             Consistency
-	rt               RetryPolicy
-	attempts         int
-	totalLatency     int64
-	serialCons       SerialConsistency
-	defaultTimestamp bool
-	context          context.Context
+	Type                  BatchType
+	Entries               []BatchEntry
+	Cons                  Consistency
+	rt                    RetryPolicy
+	attempts              int
+	totalLatency          int64
+	serialCons            SerialConsistency
+	defaultTimestamp      bool
+	defaultTimestampValue int64
+	context               context.Context
 }
 
 // NewBatch creates a new batch operation without defaults from the cluster
@@ -1263,6 +1264,18 @@ func (b *Batch) DefaultTimestamp(enable bool) *Batch {
 	return b
 }
 
+// WithTimestamp will enable the with default timestamp flag on the query
+// like DefaultTimestamp does. But also allows to define value for timestamp.
+// It works the same way as USING TIMESTAMP in the query itself, but
+// should not break prepared query optimization
+//
+// Only available on protocol >= 3
+func (b *Batch) WithTimestamp(timestamp int64) *Batch {
+	b.DefaultTimestamp(true)
+	b.defaultTimestampValue = timestamp
+	return b
+}
+
 func (b *Batch) attempt(d time.Duration) {
 	b.attempts++
 	b.totalLatency += d.Nanoseconds()