feat(cache): implement per-table tagging and raw SQL bypass

cmd/wavehouse/main.go

@@ -205,6 +205,7 @@ func run() int {
 		ctx,
 		embeddedMQ.NatsConn(),
 		chConn,
+		tiered,
 		cfg.ClickHouse.Addr,
 		cfg.ClickHouse.HTTPPort, // Uses 8123 by default
 		cfg.ClickHouse.HTTPScheme,

internal/api/pipes.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"regexp"
 	"time"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
@@ -13,6 +14,25 @@ import (
 	"golang.org/x/sync/singleflight"
 )
 
+// Naive regex to extract table names from queries
+var extractTablesRe = regexp.MustCompile(`(?i)\b(?:FROM|JOIN)\s+([a-zA-Z_][a-zA-Z0-9_]*)\b`)
+
+func extractTables(sql string) []string {
+	matches := extractTablesRe.FindAllStringSubmatch(sql, -1)
+	var tables []string
+	seen := make(map[string]bool)
+	for _, m := range matches {
+		if len(m) > 1 {
+			tbl := m[1]
+			if !seen[tbl] {
+				seen[tbl] = true
+				tables = append(tables, tbl)
+			}
+		}
+	}
+	return tables
+}
+
 // PipesHandler handles named query pipe endpoints.
 type PipesHandler struct {
 	Store      *pipes.Store
@@ -134,7 +154,6 @@ func (h *PipesHandler) Execute(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	// Execute with singleflight.
 	v, err, _ := h.sf.Do(cacheKey, func() (interface{}, error) {
 		queryCtx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 		defer cancel()
@@ -151,7 +170,9 @@ func (h *PipesHandler) Execute(w http.ResponseWriter, r *http.Request) {
 		}
 
 		if h.Cache != nil {
-			_ = h.Cache.Set(r.Context(), cacheKey, data, h.DefaultTTL)
+			// Extract tags and assign to the cache entry
+			tags := extractTables(sql)
+			_ = h.Cache.Set(r.Context(), cacheKey, data, h.DefaultTTL, tags)
 		}
 		return data, nil
 	})

internal/api/query.go

@@ -8,13 +8,19 @@ import (
 	"fmt"
 	"net/http"
 	"reflect"
+	"regexp"
 	"time"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 	"github.com/Wave-RF/WaveHouse/internal/cache"
 	"github.com/Wave-RF/WaveHouse/internal/policy"
 	"github.com/google/uuid"
-	"golang.org/x/sync/singleflight"
+)
+
+// Regex to detect mutating queries and extract table names
+var (
+	mutationRe         = regexp.MustCompile(`(?i)^\s*(INSERT|UPDATE|DELETE|ALTER|DROP|TRUNCATE)`)
+	extractTablesRawRe = regexp.MustCompile(`(?i)\b(?:FROM|JOIN|INTO|UPDATE|TABLE)\s+([a-zA-Z_][a-zA-Z0-9_]*)\b`)
 )
 
 // QueryHandler handles POST /v1/query.
@@ -23,7 +29,6 @@ type QueryHandler struct {
 	Cache       *cache.TieredCache
 	DefaultTTL  time.Duration
 	PolicyStore *policy.Store
-	sf          singleflight.Group
 }
 
 func NewQueryHandler(conn driver.Conn, c *cache.TieredCache, defaultTTL time.Duration) *QueryHandler {
@@ -40,7 +45,6 @@ func (h *QueryHandler) Handle(w http.ResponseWriter, r *http.Request) {
 	if h.PolicyStore != nil {
 		role := RoleFromContext(r.Context())
 		if role != "" && role != "admin" && role != "service" {
-			// Check if the role has raw_sql permission on any table.
 			p := h.PolicyStore.Get()
 			if p != nil {
 				claims, _ := ClaimsFromContext(r.Context())
@@ -70,45 +74,45 @@ func (h *QueryHandler) Handle(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	cacheKey := queryCacheKey(req.SQL, req.Params)
+	queryCtx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
+	defer cancel()
 
-	// Try cache.
-	if h.Cache != nil {
-		if data, _, err := h.Cache.Get(r.Context(), cacheKey); err == nil && data != nil {
-			w.Header().Set("Content-Type", "application/json")
-			w.Header().Set("X-Cache", "HIT")
-			_, _ = w.Write(data)
-			return
-		}
+	// Execute immediately (Bypassing the Cache)
+	result, err := h.executeQuery(queryCtx, req.SQL, req.Params)
+	if err != nil {
+		writeJSONError(w, http.StatusInternalServerError, err.Error())
+		return
 	}
 
-	// Execute query with singleflight to protect ClickHouse from thundering herds.
-	v, err, _ := h.sf.Do(cacheKey, func() (interface{}, error) {
-		queryCtx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
-		defer cancel()
-		result, err := h.executeQuery(queryCtx, req.SQL, req.Params)
-		if err != nil {
-			return nil, err
+	// CRUDE INVALIDATION STOP-GAP:
+	// If this raw SQL was a mutation, extract the table names and invalidate them.
+	if h.Cache != nil && mutationRe.MatchString(req.SQL) {
+		matches := extractTablesRawRe.FindAllStringSubmatch(req.SQL, -1)
+		var tags []string
+		seen := make(map[string]bool)
+		for _, m := range matches {
+			if len(m) > 1 {
+				tbl := m[1]
+				if !seen[tbl] {
+					seen[tbl] = true
+					tags = append(tags, tbl)
+				}
+			}
 		}
-
-		data, err := json.Marshal(result)
-		if err != nil {
-			return nil, err
+		if len(tags) > 0 {
+			_ = h.Cache.InvalidateByTags(r.Context(), tags)
 		}
+	}
 
-		if h.Cache != nil {
-			_ = h.Cache.Set(r.Context(), cacheKey, data, h.DefaultTTL)
-		}
-		return data, nil
-	})
+	data, err := json.Marshal(result)
 	if err != nil {
 		writeJSONError(w, http.StatusInternalServerError, err.Error())
 		return
 	}
 
 	w.Header().Set("Content-Type", "application/json")
-	w.Header().Set("X-Cache", "MISS")
-	_, _ = w.Write(v.([]byte))
+	w.Header().Set("X-Cache", "BYPASS")
+	_, _ = w.Write(data)
 }
 
 func (h *QueryHandler) executeQuery(ctx context.Context, sql string, params []any) ([]map[string]any, error) {

internal/api/structured_query.go

@@ -146,7 +146,7 @@ func (h *StructuredQueryHandler) Handle(w http.ResponseWriter, r *http.Request)
 		}
 
 		if h.Cache != nil {
-			_ = h.Cache.Set(r.Context(), cacheKey, data, ttl)
+			_ = h.Cache.Set(r.Context(), cacheKey, data, ttl, []string{table})
 		}
 		return data, nil
 	})

internal/cache/cache.go

@@ -10,8 +10,11 @@ type Cache interface {
 	// Get retrieves a value and its remaining TTL. Returns nil, 0, nil on miss.
 	Get(ctx context.Context, key string) ([]byte, time.Duration, error)
 
-	// Set stores a value with the given TTL.
-	Set(ctx context.Context, key string, value []byte, ttl time.Duration) error
+	// Set stores a value with the given TTL and tags.
+	Set(ctx context.Context, key string, value []byte, ttl time.Duration, tags []string) error
+
+	// InvalidateByTags drops all keys associated with the provided tags.
+	InvalidateByTags(ctx context.Context, tags []string) error
 
 	// Close releases resources.
 	Close() error

internal/cache/local.go

@@ -10,8 +10,9 @@ import (
 
 // LocalCache is an L1 in-process cache backed by Ristretto.
 type LocalCache struct {
-	cache *ristretto.Cache[string, []byte]
-	ttls  sync.Map // key -> expiry time.Time
+	cache   *ristretto.Cache[string, []byte]
+	ttls    sync.Map // key -> expiry time.Time
+	tagsMap sync.Map // tag (string) -> *sync.Map (key -> struct{})
 }
 
 // NewLocal creates a new Ristretto-backed local cache.
@@ -34,27 +35,53 @@ func (l *LocalCache) Get(_ context.Context, key string) ([]byte, time.Duration,
 	}
 
 	var remaining time.Duration
-	if exp, ok := l.ttls.Load(key); ok {
-		remaining = time.Until(exp.(time.Time))
-		if remaining <= 0 {
-			l.cache.Del(key)
-			l.ttls.Delete(key)
-			return nil, 0, nil
+	if expVal, ok := l.ttls.Load(key); ok {
+		exp := expVal.(time.Time)
+		if !exp.IsZero() {
+			remaining = time.Until(exp)
+			if remaining <= 0 {
+				l.cache.Del(key)
+				l.ttls.Delete(key)
+				return nil, 0, nil
+			}
 		}
 	}
 	return val, remaining, nil
 }
 
-func (l *LocalCache) Set(_ context.Context, key string, value []byte, ttl time.Duration) error {
+func (l *LocalCache) Set(_ context.Context, key string, value []byte, ttl time.Duration, tags []string) error {
 	l.cache.SetWithTTL(key, value, int64(len(value)), ttl)
+	var exp time.Time
 	if ttl > 0 {
-		l.ttls.Store(key, time.Now().Add(ttl))
+		exp = time.Now().Add(ttl)
+	}
+	l.ttls.Store(key, exp)
+
+	// Register key under each tag
+	for _, tag := range tags {
+		m, _ := l.tagsMap.LoadOrStore(tag, &sync.Map{})
+		m.(*sync.Map).Store(key, struct{}{})
+	}
+	return nil
+}
+
+func (l *LocalCache) InvalidateByTags(_ context.Context, tags []string) error {
+	for _, tag := range tags {
+		if keysMap, ok := l.tagsMap.Load(tag); ok {
+			// Iterate through all keys associated with this tag
+			keysMap.(*sync.Map).Range(func(key, _ any) bool {
+				k := key.(string)
+				l.cache.Del(k)
+				l.ttls.Delete(k)
+				return true
+			})
+			// Clear the tag index entirely
+			l.tagsMap.Delete(tag)
+		}
 	}
 	return nil
 }
 
-// Wait blocks until all buffered writes have been applied.
-// Exposed for testing; production callers rarely need this.
 func (l *LocalCache) Wait() {
 	l.cache.Wait()
 }

internal/cache/local_test.go

@@ -28,7 +28,7 @@ func TestLocalCache_SetAndGet(t *testing.T) {
 	defer func() { _ = c.Close() }()
 
 	ctx := context.Background()
-	err = c.Set(ctx, "key1", []byte("hello"), 10*time.Second)
+	err = c.Set(ctx, "key1", []byte("hello"), 10*time.Second, nil)
 	assert.NoError(t, err)
 
 	// Ristretto uses async admission — wait briefly for it to be admitted.
@@ -48,7 +48,7 @@ func TestLocalCache_ExpiredKey(t *testing.T) {
 
 	ctx := context.Background()
 	// Set with very short TTL.
-	err = c.Set(ctx, "expires", []byte("data"), 1*time.Millisecond)
+	err = c.Set(ctx, "expires", []byte("data"), 1*time.Millisecond, nil)
 	assert.NoError(t, err)
 
 	// Ensure async admission completes, then wait for expiry.
@@ -67,9 +67,9 @@ func TestLocalCache_Overwrite(t *testing.T) {
 	defer func() { _ = c.Close() }()
 
 	ctx := context.Background()
-	require.NoError(t, c.Set(ctx, "key", []byte("v1"), 10*time.Second))
+	require.NoError(t, c.Set(ctx, "key", []byte("v1"), 10*time.Second, nil))
 	c.Wait()
-	require.NoError(t, c.Set(ctx, "key", []byte("v2"), 10*time.Second))
+	require.NoError(t, c.Set(ctx, "key", []byte("v2"), 10*time.Second, nil))
 	c.Wait()
 
 	val, _, err := c.Get(ctx, "key")
@@ -84,7 +84,7 @@ func TestLocalCache_ZeroTTL(t *testing.T) {
 	defer func() { _ = c.Close() }()
 
 	ctx := context.Background()
-	err = c.Set(ctx, "notimed", []byte("data"), 0)
+	err = c.Set(ctx, "notimed", []byte("data"), 0, nil)
 	assert.NoError(t, err)
 
 	time.Sleep(10 * time.Millisecond)
@@ -96,3 +96,36 @@ func TestLocalCache_ZeroTTL(t *testing.T) {
 		assert.Zero(t, ttl, "expected zero remaining TTL for key without TTL")
 	}
 }
+
+func TestLocalCache_InvalidateByTags(t *testing.T) {
+	t.Parallel()
+	c, err := NewLocal(1 << 20)
+	require.NoError(t, err)
+	defer func() { _ = c.Close() }()
+
+	ctx := context.Background()
+
+	// Pass the table names as the tags slice
+	require.NoError(t, c.Set(ctx, "query:clicks:123", []byte("val1"), 10*time.Second, []string{"clicks"}))
+	require.NoError(t, c.Set(ctx, "query:clicks:456", []byte("val2"), 0, []string{"clicks"}))
+
+	// This key belongs to 'users', so it won't be wiped when we invalidate 'clicks'
+	require.NoError(t, c.Set(ctx, "query:users:789", []byte("val3"), 10*time.Second, []string{"users"}))
+
+	c.Wait()
+
+	// Wipe out the clicks table cache
+	err = c.InvalidateByTags(ctx, []string{"clicks"})
+	assert.NoError(t, err)
+
+	// Verify the target keys are gone
+	val1, _, _ := c.Get(ctx, "query:clicks:123")
+	assert.Nil(t, val1)
+
+	val2, _, _ := c.Get(ctx, "query:clicks:456")
+	assert.Nil(t, val2)
+
+	// Verify the innocent key is still there
+	val3, _, _ := c.Get(ctx, "query:users:789")
+	assert.Equal(t, []byte("val3"), val3)
+}

internal/cache/tiered.go

@@ -32,7 +32,8 @@ func (t *TieredCache) Get(ctx context.Context, key string) ([]byte, time.Duratio
 			val, ttl, err := t.l2.Get(ctx, key)
 			if err == nil && val != nil {
 				// Backfill L1 with remaining TTL from L2.
-				_ = t.l1.Set(ctx, key, val, ttl)
+				// Note: We don't have the original tags here, but L2 is out-of-scope for v1 Standalone.
+				_ = t.l1.Set(ctx, key, val, ttl, nil)
 				return &cacheResult{data: val, ttl: ttl}, nil
 			}
 		}
@@ -48,12 +49,24 @@ func (t *TieredCache) Get(ctx context.Context, key string) ([]byte, time.Duratio
 	return nil, 0, nil
 }
 
-func (t *TieredCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error {
-	if err := t.l1.Set(ctx, key, value, ttl); err != nil {
+// UPDATE: Accept tags and pass them down
+func (t *TieredCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration, tags []string) error {
+	if err := t.l1.Set(ctx, key, value, ttl, tags); err != nil {
 		return err
 	}
 	if t.l2 != nil {
-		return t.l2.Set(ctx, key, value, ttl)
+		return t.l2.Set(ctx, key, value, ttl, tags)
+	}
+	return nil
+}
+
+// UPDATE: Replace InvalidateByPrefix with InvalidateByTags
+func (t *TieredCache) InvalidateByTags(ctx context.Context, tags []string) error {
+	if err := t.l1.InvalidateByTags(ctx, tags); err != nil {
+		return err
+	}
+	if t.l2 != nil {
+		return t.l2.InvalidateByTags(ctx, tags)
 	}
 	return nil
 }