From 694688477a75ed1358e8328bcd52ebc1fdef7d77 Mon Sep 17 00:00:00 2001
From: steebchen <contact@luca-steeb.com>
Date: Sun, 17 Mar 2024 17:51:17 +0700
Subject: [PATCH] improvements

---
 internal/config/loader/loader.go         | 11 +++++++++--
 internal/repository/cache/cache.go       |  6 ++++++
 internal/repository/prisma/repository.go | 18 ++++++++++--------
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/internal/config/loader/loader.go b/internal/config/loader/loader.go
index 535b148ef..14222bee3 100644
--- a/internal/config/loader/loader.go
+++ b/internal/config/loader/loader.go
@@ -27,6 +27,7 @@ import (
 	"github.com/hatchet-dev/hatchet/internal/integrations/vcs/github"
 	"github.com/hatchet-dev/hatchet/internal/logger"
 	"github.com/hatchet-dev/hatchet/internal/msgqueue/rabbitmq"
+	"github.com/hatchet-dev/hatchet/internal/repository/cache"
 	"github.com/hatchet-dev/hatchet/internal/repository/prisma"
 	"github.com/hatchet-dev/hatchet/internal/repository/prisma/db"
 	"github.com/hatchet-dev/hatchet/internal/services/ingestor"
@@ -149,9 +150,15 @@ func GetDatabaseConfigFromConfigFile(cf *database.ConfigFile) (res *database.Con
 		return nil, fmt.Errorf("could not connect to database: %w", err)
 	}
 
+	ch := cache.New(cf.CacheDuration)
+
 	return &database.Config{
-		Disconnect: c.Prisma.Disconnect,
-		Repository: prisma.NewPrismaRepository(c, pool, prisma.WithLogger(&l), prisma.WithCacheDuration(cf.CacheDuration)),
+		Disconnect: func() error {
+			ch.Purge()
+
+			return c.Prisma.Disconnect()
+		},
+		Repository: prisma.NewPrismaRepository(c, pool, prisma.WithLogger(&l), prisma.WithCache(ch)),
 		Seed:       cf.Seed,
 	}, nil
 }
diff --git a/internal/repository/cache/cache.go b/internal/repository/cache/cache.go
index 844f4dcd2..0875aba77 100644
--- a/internal/repository/cache/cache.go
+++ b/internal/repository/cache/cache.go
@@ -12,6 +12,8 @@ type Cacheable interface {
 
 	// Get gets a value from the cache with the given key
 	Get(key string) (interface{}, bool)
+
+	Purge()
 }
 
 type Cache struct {
@@ -26,6 +28,10 @@ func (c *Cache) Get(key string) (interface{}, bool) {
 	return c.cache.Get(key)
 }
 
+func (c *Cache) Purge() {
+	c.cache.Purge()
+}
+
 func New(duration time.Duration) *Cache {
 	if duration == 0 {
 		// consider a duration of 0 a very short expiry instead of no expiry
diff --git a/internal/repository/prisma/repository.go b/internal/repository/prisma/repository.go
index 11500142e..f4d847460 100644
--- a/internal/repository/prisma/repository.go
+++ b/internal/repository/prisma/repository.go
@@ -37,9 +37,9 @@ type prismaRepository struct {
 type PrismaRepositoryOpt func(*PrismaRepositoryOpts)
 
 type PrismaRepositoryOpts struct {
-	v             validator.Validator
-	l             *zerolog.Logger
-	cacheDuration time.Duration
+	v     validator.Validator
+	l     *zerolog.Logger
+	cache cache.Cacheable
 }
 
 func defaultPrismaRepositoryOpts() *PrismaRepositoryOpts {
@@ -60,9 +60,9 @@ func WithLogger(l *zerolog.Logger) PrismaRepositoryOpt {
 	}
 }
 
-func WithCacheDuration(duration time.Duration) PrismaRepositoryOpt {
+func WithCache(cache cache.Cacheable) PrismaRepositoryOpt {
 	return func(opts *PrismaRepositoryOpts) {
-		opts.cacheDuration = duration
+		opts.cache = cache
 	}
 }
 
@@ -76,13 +76,15 @@ func NewPrismaRepository(client *db.PrismaClient, pool *pgxpool.Pool, fs ...Pris
 	newLogger := opts.l.With().Str("service", "database").Logger()
 	opts.l = &newLogger
 
-	c := cache.New(opts.cacheDuration)
+	if opts.cache == nil {
+		opts.cache = cache.New(1 * time.Millisecond)
+	}
 
 	return &prismaRepository{
-		apiToken:       NewAPITokenRepository(client, opts.v, c),
+		apiToken:       NewAPITokenRepository(client, opts.v, opts.cache),
 		event:          NewEventRepository(client, pool, opts.v, opts.l),
 		log:            NewLogRepository(client, pool, opts.v, opts.l),
-		tenant:         NewTenantRepository(client, opts.v, c),
+		tenant:         NewTenantRepository(client, opts.v, opts.cache),
 		tenantInvite:   NewTenantInviteRepository(client, opts.v),
 		workflow:       NewWorkflowRepository(client, pool, opts.v, opts.l),
 		workflowRun:    NewWorkflowRunRepository(client, pool, opts.v, opts.l),