add support for querying entry expiration.

Author: raulk

datastore.go

@@ -73,8 +73,7 @@ func NewDatastore(path string, options *Options) (*Datastore, error) {
 	}
 
 	return &Datastore{
-		DB: kv,
-
+		DB:             kv,
 		gcDiscardRatio: gcDiscardRatio,
 	}, nil
 }
@@ -125,6 +124,13 @@ func (d *Datastore) SetTTL(key ds.Key, ttl time.Duration) error {
 	return txn.Commit()
 }
 
+func (d *Datastore) GetExpiration(key ds.Key) (time.Time, error) {
+	txn := d.newImplicitTransaction(false).(*txn)
+	defer txn.Discard()
+
+	return txn.GetExpiration(key)
+}
+
 func (d *Datastore) Get(key ds.Key) (value []byte, err error) {
 	txn := d.newImplicitTransaction(true)
 	defer txn.Discard()
@@ -192,6 +198,16 @@ func (t *txn) PutWithTTL(key ds.Key, value []byte, ttl time.Duration) error {
 	return t.txn.SetWithTTL(key.Bytes(), value, ttl)
 }
 
+func (t *txn) GetExpiration(key ds.Key) (time.Time, error) {
+	item, err := t.txn.Get(key.Bytes())
+	if err == badger.ErrKeyNotFound {
+		return time.Time{}, ds.ErrNotFound
+	} else if err != nil {
+		return time.Time{}, err
+	}
+	return time.Unix(int64(item.ExpiresAt()), 0), nil
+}
+
 func (t *txn) SetTTL(key ds.Key, ttl time.Duration) error {
 	data, err := t.Get(key)
 	if err != nil {
@@ -286,6 +302,10 @@ func (t *txn) Query(q dsq.Query) (dsq.Results, error) {
 				result = dsq.Result{Entry: e}
 			}
 
+			if q.ReturnExpirations {
+				result.Expiration = time.Unix(int64(item.ExpiresAt()), 0)
+			}
+
 			select {
 			case qrb.Output <- result:
 			case <-worker.Closing(): // client told us to close early

ds_test.go

@@ -672,3 +672,76 @@ func TestTTL(t *testing.T) {
 
 	d.Close()
 }
+
+func TestExpirations(t *testing.T) {
+	var err error
+
+	d, done := newDS(t)
+	defer done()
+
+	txn := d.NewTransaction(false)
+	ttltxn := txn.(ds.TTLDatastore)
+	defer txn.Discard()
+
+	key := ds.NewKey("/abc/def")
+	val := make([]byte, 32)
+	if n, err := rand.Read(val); n != 32 || err != nil {
+		t.Fatal("source of randomness failed")
+	}
+
+	ttl := time.Hour
+	now := time.Now()
+	tgt := now.Add(ttl)
+
+	if err = ttltxn.PutWithTTL(key, val, ttl); err != nil {
+		t.Fatalf("adding with ttl failed: %v", err)
+	}
+
+	if err = txn.Commit(); err != nil {
+		t.Fatalf("commiting transaction failed: %v", err)
+	}
+
+	// Second transaction to retrieve expirations.
+	txn = d.NewTransaction(true)
+	ttltxn = txn.(ds.TTLDatastore)
+	defer txn.Discard()
+
+	// GetExpiration returns expected value.
+	var dsExp time.Time
+	if dsExp, err = ttltxn.GetExpiration(key); err != nil {
+		t.Fatalf("getting expiration failed: %v", err)
+	} else if tgt.Sub(dsExp) >= 5*time.Second {
+		t.Fatal("expiration returned by datastore not within the expected range (tolerance: 5 seconds)")
+	} else if tgt.Sub(dsExp) < 0 {
+		t.Fatal("expiration returned by datastore was earlier than expected")
+	}
+
+	// Iterator returns expected value.
+	q := dsq.Query{
+		ReturnExpirations: true,
+		KeysOnly:          true,
+	}
+	var ress dsq.Results
+	if ress, err = txn.Query(q); err != nil {
+		t.Fatalf("querying datastore failed: %v", err)
+	}
+
+	defer ress.Close()
+	if res, ok := ress.NextSync(); !ok {
+		t.Fatal("expected 1 result in iterator")
+	} else if res.Expiration != dsExp {
+		t.Fatalf("expiration returned from iterator differs from GetExpiration, expected: %v, actual: %v", dsExp, res.Expiration)
+	}
+
+	if _, ok := ress.NextSync(); ok {
+		t.Fatal("expected no more results in iterator")
+	}
+
+	// Datastore->GetExpiration()
+	if exp, err := d.GetExpiration(key); err != nil {
+		t.Fatalf("querying datastore failed: %v", err)
+	} else if exp != dsExp {
+		t.Fatalf("expiration returned from DB differs from that returned by txn, expected: %v, actual: %v", dsExp, exp)
+	}
+
+}