From daf4fe97dc011c09899a5c19065f6106ea35ce3c Mon Sep 17 00:00:00 2001
From: saji <saji@saji.ev>
Date: Wed, 6 Mar 2024 17:09:02 -0600
Subject: [PATCH] added DocumentNotFound error

---
 db.go      | 38 ++++++++++++++++++++++++++++++++++++--
 db_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/db.go b/db.go
index 818d572..5e9b00f 100644
--- a/db.go
+++ b/db.go
@@ -321,6 +321,14 @@ func (tdb *TelemDb) AddDocument(ctx context.Context, obj json.RawMessage) error
 	return err
 }
 
+// DocumentNotFoundError is when the underlying document cannot be found.
+type DocumentNotFoundError string
+
+func (e DocumentNotFoundError) Error() string {
+	return fmt.Sprintf("document could not find key: %s", string(e))
+}
+
+
 // UpdateDocument replaces the entire contents of a document matching
 // the given key. Note that the key is derived from the document,
 // and no checks are done to ensure that the new key is the same.
@@ -328,10 +336,21 @@ func (tdb *TelemDb) UpdateDocument(ctx context.Context, key string,
 	obj json.RawMessage) error {
 
 	const upd = `UPDATE openmct_objects SET data = json(?) WHERE key IS ?`
-	_, err := tdb.db.ExecContext(ctx, upd, obj, key)
+	r, err := tdb.db.ExecContext(ctx, upd, obj, key)
+	if err != nil {
+		return err
+	}
+	n, err := r.RowsAffected()
+	if err != nil {
+		return err
+	}
+	if n != 1 {
+		return DocumentNotFoundError(key)
+	}
 	return err
 }
 
+
 // GetDocument gets the document matching the corresponding key.
 func (tdb *TelemDb) GetDocument(ctx context.Context, key string) (json.RawMessage, error) {
 	const get = `SELECT data FROM openmct_objects WHERE key IS ?`
@@ -341,6 +360,11 @@ func (tdb *TelemDb) GetDocument(ctx context.Context, key string) (json.RawMessag
 	var res []byte // VERY important, json.RawMessage won't work here
 	// since the scan function does not look at underlying types.
 	row.Scan(&res)
+
+	if len(res) == 0 {
+		return nil, DocumentNotFoundError(key)
+	}
+
 	return res, nil
 }
 
@@ -366,6 +390,16 @@ func (tdb *TelemDb) GetAllDocuments(ctx context.Context) ([]json.RawMessage, err
 // if it does not exist.
 func (tdb *TelemDb) DeleteDocument(ctx context.Context, key string) error {
 	const del = `DELETE FROM openmct_objects WHERE key IS ?`
-	_, err := tdb.db.ExecContext(ctx, del, key)
+	res, err := tdb.db.ExecContext(ctx, del, key)
+	if err != nil {
+		return err
+	}
+	n, err := res.RowsAffected()
+	if err != nil {
+		return err
+	}
+	if n != 1 {
+		return DocumentNotFoundError(key)
+	}
 	return err
 }
diff --git a/db_test.go b/db_test.go
index fee8b8f..a7120f3 100644
--- a/db_test.go
+++ b/db_test.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/rand"
 	"reflect"
@@ -232,4 +233,52 @@ func TestDbDocuments(t *testing.T) {
 		}
 
 	})
+
+	t.Run("test getting nonexistent document", func(t *testing.T) {
+		tdb := MakeMockDatabase(t.Name())
+		tdb.db.Ping()
+		ctx := context.Background()
+
+		res, err := tdb.GetDocument(ctx, "hi")
+
+		if err == nil || !errors.Is(err, DocumentNotFoundError("hi")){
+			t.Fatalf("GetDocument expected DocumentNotFoundError, got %v", err)
+		}
+		if res != nil {
+			t.Fatalf("GetDocument expected nil result, got %v", res)
+		}
+	})
+
+	t.Run("test update document", func(t *testing.T) {
+		tdb := MakeMockDatabase(t.Name())
+		tdb.db.Ping()
+		ctx := context.Background()
+		doc1 := MockDocument("hi")
+		doc2 := MockDocument("hi") // same key, we want to update.
+
+		tdb.AddDocument(ctx, doc1)
+		err := tdb.UpdateDocument(ctx, "hi", doc2)
+		if err != nil {
+			t.Fatalf("UpdateDocument expected no error, got err=%v", err)
+		}
+
+		// compare.
+		res, _ := tdb.GetDocument(ctx, "hi")
+		if !reflect.DeepEqual(res, doc2) {
+			t.Fatalf("UpdateDocument did not return new doc, got %s", res)
+		}
+
+	})
+
+	t.Run("test update nonexistent document", func(t *testing.T) {
+		tdb := MakeMockDatabase(t.Name())
+		tdb.db.Ping()
+		ctx := context.Background()
+		doc := MockDocument("hi")
+		err := tdb.UpdateDocument(ctx, "badKey", doc)
+		if err == nil {
+			t.Fatalf("UpdateDocument expected error, got nil")
+		}
+	})
+
 }