From 9ea8cadf8469cc92e031ce986f5840e3ae0be0a8 Mon Sep 17 00:00:00 2001
From: Yash Mehrotra <yashmehrotra95@gmail.com>
Date: Tue, 28 Nov 2023 13:09:54 +0530
Subject: [PATCH] feat: add postgrest response modifier

---
 cmd/serve.go     | 34 ++++++++++++++++++++++++++++++++--
 pkg/db/canary.go | 11 +++++++++--
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/cmd/serve.go b/cmd/serve.go
index 229112854..27b50efc6 100644
--- a/cmd/serve.go
+++ b/cmd/serve.go
@@ -2,12 +2,15 @@ package cmd
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	_ "net/http/pprof" // required by serve
 	"net/url"
 	"os"
 	"os/signal"
+	"strings"
 	"time"
 
 	"github.com/labstack/echo/v4"
@@ -81,6 +84,32 @@ func setup() {
 	go push.Start()
 }
 
+func postgrestResponseModifier(r *http.Response) error {
+	shouldPersistCanary := r.Request.Method == http.MethodPost &&
+		strings.TrimSuffix(r.Request.URL.Path, "/") == "/canaries" &&
+		r.StatusCode == http.StatusCreated
+
+	// If a new canary is inserted via postgrest, we need to persist the canary
+	// again so that all the checks of that canary are created in the database
+	if shouldPersistCanary {
+		var canaries []pkg.Canary
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			return fmt.Errorf("error reading response body: %w", err)
+		}
+		if err := json.Unmarshal(body, &canaries); err != nil {
+			return fmt.Errorf("error unmarshaling response body to json: %w", err)
+		}
+		for _, c := range canaries {
+			if _, err := db.PersistCanaryModel(c); err != nil {
+				logger.Errorf("Error persisting canary[%s]: %v", c.ID, err)
+			}
+		}
+	}
+
+	return nil
+}
+
 func serve() {
 	var allowedCors string
 	e := echo.New()
@@ -106,7 +135,7 @@ func serve() {
 
 	if !disablePostgrest {
 		go db.StartPostgrest()
-		forward(e, "/db", db.PostgRESTEndpoint())
+		forward(e, "/db", db.PostgRESTEndpoint(), postgrestResponseModifier)
 	}
 
 	e.Use(middleware.Logger())
@@ -158,7 +187,7 @@ func serve() {
 	}
 }
 
-func forward(e *echo.Echo, prefix string, target string) {
+func forward(e *echo.Echo, prefix string, target string, respModifierFunc func(*http.Response) error) {
 	targetURL, err := url.Parse(target)
 	if err != nil {
 		e.Logger.Fatal(err)
@@ -172,6 +201,7 @@ func forward(e *echo.Echo, prefix string, target string) {
 				URL: targetURL,
 			},
 		}),
+		ModifyResponse: respModifierFunc,
 	}))
 }
 
diff --git a/pkg/db/canary.go b/pkg/db/canary.go
index bb4988064..de8710979 100644
--- a/pkg/db/canary.go
+++ b/pkg/db/canary.go
@@ -20,10 +20,15 @@ import (
 	dutyTypes "github.com/flanksource/duty/types"
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 )
 
+var (
+	PostgresDuplicateKeyError = &pgconn.PgError{Code: "23505"}
+)
+
 func GetAllCanariesForSync(ctx context.Context, namespace string) ([]pkg.Canary, error) {
 	query := `
         SELECT json_agg(
@@ -335,9 +340,11 @@ func PersistCanaryModel(model pkg.Canary) (*pkg.Canary, error) {
 
 	// Duplicate key happens when an already created canary is persisted
 	// We will ignore this error but act on other errors
+	// In this scenario PostgresDuplicateKeyError is checked primarily and
+	// gorm.ErrDuplicatedKey is just for fallback but does not work
 	if err != nil {
-		if !errors.Is(err, gorm.ErrDuplicatedKey) {
-			return nil, err
+		if !errors.As(err, &PostgresDuplicateKeyError) && !errors.Is(err, gorm.ErrDuplicatedKey) {
+			return nil, fmt.Errorf("error persisting canary to db: %w", err)
 		}
 	}