22284276152

Committed 22 Feb 2026 08:01PM UTC coverage: 49.631% (+0.09%) from 49.538%

Build # 22284276152

Build Type

push

github

Committed by

mostafa

Commit Message

Fix lint: add checked type assertion and allow protobuf/proto in depguard

Coverage Stats

3 of 3 new or added lines in 1 file covered. (100.0%)

23 existing lines in 1 file now uncovered.

269 of 542 relevant lines covered (49.63%)

2.13 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

79.21

/plugin/plugin.go

package plugin

import (
        "context"
        "encoding/base64"
        "strings"
        "sync"
        "time"

        sdkAct "github.com/gatewayd-io/gatewayd-plugin-sdk/act"
        "github.com/gatewayd-io/gatewayd-plugin-sdk/databases/postgres"
        sdkPlugin "github.com/gatewayd-io/gatewayd-plugin-sdk/plugin"
        v1 "github.com/gatewayd-io/gatewayd-plugin-sdk/plugin/v1"
        apiV1 "github.com/gatewayd-io/gatewayd/api/v1"
        "github.com/hashicorp/go-hclog"
        goplugin "github.com/hashicorp/go-plugin"
        goRedis "github.com/redis/go-redis/v9"
        "github.com/spf13/cast"
        "google.golang.org/grpc"
        "google.golang.org/protobuf/proto"
)

type Plugin struct {
        goplugin.GRPCPlugin
        v1.GatewayDPluginServiceServer

        Logger hclog.Logger

        APIClient apiV1.GatewayDAdminAPIServiceClient

        // Cache configuration.
        RedisClient        *goRedis.Client
        RedisURL           string
        Expiry             time.Duration
        DefaultDBName      string
        ScanCount          int64
        ExitOnStartupError bool

        UpdateCacheChannel chan *v1.Struct
        WaitGroup          *sync.WaitGroup

        // Periodic invalidator configuration.
        PeriodicInvalidatorEnabled    bool
        PeriodicInvalidatorStartDelay time.Duration
        PeriodicInvalidatorInterval   time.Duration
}

type CachePlugin struct {
        goplugin.NetRPCUnsupportedPlugin
        Impl Plugin
}

// Define a set for PostgreSQL date/time functions
// https://www.postgresql.org/docs/8.2/functions-datetime.html
var pgDateTimeFunctions = map[string]struct{}{
        "AGE":                   {},
        "CLOCK_TIMESTAMP":       {},
        "CURRENT_DATE":          {},
        "CURRENT_TIME":          {},
        "CURRENT_TIMESTAMP":     {},
        "LOCALTIME":             {},
        "LOCALTIMESTAMP":        {},
        "NOW":                   {},
        "STATEMENT_TIMESTAMP":   {},
        "TIMEOFDAY":             {},
        "TRANSACTION_TIMESTAMP": {},
}

// NewCachePlugin returns a new instance of the CachePlugin.
func NewCachePlugin(impl Plugin) *CachePlugin {
        return &CachePlugin{
                NetRPCUnsupportedPlugin: goplugin.NetRPCUnsupportedPlugin{},
                Impl:                    impl,
        }
}

// GRPCServer registers the plugin with the gRPC server.
func (p *CachePlugin) GRPCServer(_ *goplugin.GRPCBroker, s *grpc.Server) error {
        v1.RegisterGatewayDPluginServiceServer(s, &p.Impl)
        return nil
}

// GRPCClient returns the plugin client.
func (p *CachePlugin) GRPCClient(_ context.Context, _ *goplugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
        return v1.NewGatewayDPluginServiceClient(c), nil
}

// GetPluginConfig returns the plugin config.
func (p *Plugin) GetPluginConfig(
        _ context.Context, _ *v1.Struct,
) (*v1.Struct, error) {
        GetPluginConfigCounter.Inc()
        return v1.NewStruct(PluginConfig)
}

// OnTrafficFromClient is called when a request is received by GatewayD from the client.
func (p *Plugin) OnTrafficFromClient(
        ctx context.Context, req *v1.Struct,
) (*v1.Struct, error) {
        OnTrafficFromClientCounter.Inc()
        req, err := postgres.HandleClientMessage(req, p.Logger)
        if err != nil {
                p.Logger.Info("Failed to handle client message", "error", err)
        }

        // This is used as a fallback if the database is not found in the startup message.
        database := p.DefaultDBName
        if database == "" {
                client := cast.ToStringMapString(sdkPlugin.GetAttr(req, "client", nil))
                database = p.getDBFromStartupMessage(ctx, req, database, client)

                // Get the database from the cache if it's not found in the startup message or
                // if the current request is not a startup message.
                if database == "" {
                        database, err = p.RedisClient.Get(ctx, client["remote"]).Result()
                        if err != nil {
                                CacheErrorsCounter.Inc()
                                p.Logger.Debug("Failed to get cache", "error", err)
                        }
                        CacheGetsCounter.Inc()
                        p.Logger.Debug("Get the database in the cache for the current session",
                                "database", database, "client", client["remote"])
                }
        }

        // If the database is still not found, return the response as is without caching.
        // This might also happen if the cache is cleared while the client is still connected.
        // In this case, the client should reconnect and the error will go away.
        preconditions := sdkPlugin.GetAttr(req, "sslRequest", "") != "" ||
                sdkPlugin.GetAttr(req, "saslInitialResponse", "") != "" ||
                sdkPlugin.GetAttr(req, "cancelRequest", "") != ""
        if database == "" && !preconditions {
                p.Logger.Error(
                        "Database name not found or set in cache, startup message or plugin config. Skipping cache")
                p.Logger.Error("Consider setting the database name in the plugin config or disabling the plugin if you don't need it")
                return req, nil
        }

        query := cast.ToString(sdkPlugin.GetAttr(req, "query", ""))
        request := cast.ToString(sdkPlugin.GetAttr(req, "request", ""))
        server := cast.ToStringMapString(sdkPlugin.GetAttr(req, "server", ""))
        cacheKey := strings.Join([]string{server["remote"], database, request}, ":")

        if query == "" {
                return req, nil
        }

        p.Logger.Trace("Query", "query", query)

        // Clear the cache if the query is an insert, update or delete query.
        p.invalidateDML(ctx, query)

        // Check if the query is cached.
        response, err := p.RedisClient.Get(ctx, cacheKey).Bytes()
        if err != nil {
                p.Logger.Debug("Failed to get cached response", "error", err)
        }
        CacheGetsCounter.Inc()

        if response == nil {
                // If the query is not cached, return the request as is.
                CacheMissesCounter.Inc()
                return req, nil
        }

        // If the query is cached, return the cached response.
        signals, err := v1.NewList([]any{
                sdkAct.Terminate().ToMap(),
                sdkAct.Log("debug", "Returning cached response", map[string]any{
                        "cacheKey": []byte(cacheKey),
                        "plugin":   PluginID.GetName(),
                }).ToMap(),
        })
        if err != nil {
                CacheErrorsCounter.Inc()
                p.Logger.Error("Failed to create signals", "error", err)
        } else {
                CacheHitsCounter.Inc()
                // Return the cached response.
                req.Fields[sdkAct.Signals] = v1.NewListValue(signals)
                req.Fields["response"] = v1.NewBytesValue(response)
        }
        return req, nil
}

// IsCacheNeeded determines if caching is needed.
func IsCacheNeeded(upperQuery string) bool {
        // Iterate over each function name in the set of PostgreSQL date/time functions.
        for function := range pgDateTimeFunctions {
                if strings.Contains(upperQuery, function) {
                        // If the query contains a date/time function, caching is not needed.
                        return false
                }
        }
        return true
}

func (p *Plugin) UpdateCache(ctx context.Context) {
        defer p.WaitGroup.Done()
        for {
                serverResponse, ok := <-p.UpdateCacheChannel
                if !ok {
                        p.Logger.Info("Channel closed, returning from function")
                        return
                }

                OnTrafficFromServerCounter.Inc()
                resp, err := postgres.HandleServerMessage(serverResponse, p.Logger)
                if err != nil {
                        p.Logger.Info("Failed to handle server message", "error", err)
                }

                rowDescription := cast.ToString(sdkPlugin.GetAttr(resp, "rowDescription", ""))
                dataRow := cast.ToStringSlice(sdkPlugin.GetAttr(resp, "dataRow", []interface{}{}))
                errorResponse := cast.ToString(sdkPlugin.GetAttr(resp, "errorResponse", ""))
                request, isOk := sdkPlugin.GetAttr(resp, "request", nil).([]byte)
                if !isOk {
                        request = []byte{}
                }

                response, isOk := sdkPlugin.GetAttr(resp, "response", nil).([]byte)
                if !isOk {
                        response = []byte{}
                }
                server := cast.ToStringMapString(sdkPlugin.GetAttr(resp, "server", ""))

                // This is used as a fallback if the database is not found in the startup message.

                database := p.DefaultDBName
                if database == "" {
                        client := cast.ToStringMapString(sdkPlugin.GetAttr(resp, "client", ""))
                        if client != nil && client["remote"] != "" {
                                database, err = p.RedisClient.Get(ctx, client["remote"]).Result()
                                if err != nil {
                                        CacheErrorsCounter.Inc()
                                        p.Logger.Debug("Failed to get cached response", "error", err)
                                }
                                CacheGetsCounter.Inc()
                        }
                }

                // If the database is still not found, return the response as is without caching.
                // This might also happen if the cache is cleared while the client is still connected.
                // In this case, the client should reconnect and the error will go away.
                if database == "" {
                        p.Logger.Debug("Database name not found or set in cache, startup message or plugin config. " +
                                "Skipping cache")
                        p.Logger.Debug("Consider setting the database name in the " +
                                "plugin config or disabling the plugin if you don't need it")
                        continue
                }

                cacheKey := strings.Join([]string{server["remote"], database, string(request)}, ":")
                if errorResponse != "" || rowDescription == "" || dataRow == nil || len(dataRow) == 0 {
                        continue
                }

                query, err := postgres.GetQueryFromRequest(request)
                if err != nil {
                        p.Logger.Debug("Failed to get query from request", "error", err)
                        continue
                }

                if !IsCacheNeeded(strings.ToUpper(query)) {
                        continue
                }

                // The request was successful and the response contains data. Cache the response.
                if err := p.RedisClient.Set(ctx, cacheKey, response, p.Expiry).Err(); err != nil {
                        CacheErrorsCounter.Inc()
                        p.Logger.Debug("Failed to set cache", "error", err)
                }
                CacheSetsCounter.Inc()

                tables, err := postgres.GetTablesFromQuery(query)
                if err != nil {
                        p.Logger.Debug("Failed to get tables from query", "error", err)
                        continue
                }

                // Cache the table(s) used in each cached request. This is used to invalidate
                // the cache when a rows is inserted, updated or deleted into that table.
                for _, table := range tables {
                        requestQueryCacheKey := strings.Join([]string{table, cacheKey}, ":")
                        if err := p.RedisClient.Set(
                                ctx, requestQueryCacheKey, "", p.Expiry).Err(); err != nil {
                                CacheErrorsCounter.Inc()
                                p.Logger.Debug("Failed to set cache", "error", err)
                        }
                        CacheSetsCounter.Inc()
                }
        }
}

// OnTrafficFromServer is called when a response is received by GatewayD from the server.
func (p *Plugin) OnTrafficFromServer(
        _ context.Context, resp *v1.Struct,
) (*v1.Struct, error) {
        p.Logger.Debug("Traffic is coming from the server side")
        if cloned, ok := proto.Clone(resp).(*v1.Struct); ok {
                p.UpdateCacheChannel <- cloned
        }
        return resp, nil
}

func (p *Plugin) OnClosed(ctx context.Context, req *v1.Struct) (*v1.Struct, error) {
        OnClosedCounter.Inc()
        client := cast.ToStringMapString(sdkPlugin.GetAttr(req, "client", nil))
        if client != nil {
                if err := p.RedisClient.Del(ctx, client["remote"]).Err(); err != nil {
                        p.Logger.Debug("Failed to delete cache", "error", err)
                        CacheErrorsCounter.Inc()
                }
                p.Logger.Debug("Client closed", "client", client["remote"])
                CacheDeletesCounter.Inc()
        }
        return req, nil
}

// invalidateDML invalidates the cache for the tables that are affected by the DML.
// This is done by getting the cached queries for each table and deleting them.
func (p *Plugin) invalidateDML(ctx context.Context, query string) {
        // Check if the query is a UPDATE, INSERT or DELETE.
        queryDecoded, err := base64.StdEncoding.DecodeString(query)
        if err != nil {
                p.Logger.Debug("Failed to decode query", "error", err)
                return
        }

        queryMessage := cast.ToStringMapString(string(queryDecoded))
        p.Logger.Trace("Query message", "query", queryMessage)

        queryString := strings.ToUpper(queryMessage["String"])
        // Ignore SELECT and WITH/SELECT queries.
        // TODO: This is a naive approach, but query parsing has a cost.
        if strings.HasPrefix(queryString, "SELECT") ||
                (strings.HasPrefix(queryString, "WITH") &&
                        strings.Contains(queryString, "SELECT")) {
                return
        }

        tables, err := postgres.GetTablesFromQuery(queryMessage["String"])
        if err != nil {
                p.Logger.Debug("Failed to get tables from query", "error", err)
                return
        }

        p.Logger.Trace("Tables", "tables", tables)
        for _, table := range tables {
                // Invalidate the cache for the table.
                // TODO: This is not efficient. We should be able to invalidate the cache
                // for a specific key instead of invalidating the entire table.
                pipeline := p.RedisClient.Pipeline()
                var cursor uint64
                for {
                        scanResult := p.RedisClient.Scan(ctx, cursor, table+":*", p.ScanCount)
                        if scanResult.Err() != nil {
                                CacheErrorsCounter.Inc()
                                p.Logger.Debug("Failed to scan keys", "error", scanResult.Err())
                                break
                        }
                        CacheScanCounter.Inc()

                        // Per each key, delete the cache entry and the table cache key itself.
                        var keys []string
                        keys, cursor = scanResult.Val()
                        CacheScanKeysCounter.Add(float64(len(keys)))
                        for _, tableKey := range keys {
                                // Invalidate the cache for the table.
                                cachedResponseKey := strings.TrimPrefix(tableKey, table+":")
                                pipeline.Del(ctx, cachedResponseKey)
                                // Invalidate the table cache key itself.
                                pipeline.Del(ctx, tableKey)
                        }

                        if cursor == 0 {
                                break
                        }
                }

                result, err := pipeline.Exec(ctx)
                if err != nil {
                        p.Logger.Debug("Failed to execute pipeline", "error", err)
                }

                for _, res := range result {
                        if res.Err() != nil {
                                CacheErrorsCounter.Inc()
                        } else {
                                CacheDeletesCounter.Inc()
                        }
                }
        }
}

// getDBFromStartupMessage gets the database name from the startup message.
func (p *Plugin) getDBFromStartupMessage(
        ctx context.Context,
        req *v1.Struct,
        database string,
        client map[string]string,
) string {
        // Try to get the database from the startup message, which is only sent once by the client.
        // Store the database in the cache so that we can use it for subsequent requests.
        startupMessageEncoded := cast.ToString(sdkPlugin.GetAttr(req, "startupMessage", ""))
        if startupMessageEncoded == "" {
                return database
        }

        startupMessageBytes, err := base64.StdEncoding.DecodeString(startupMessageEncoded)
        if err != nil {
                p.Logger.Debug("Failed to decode startup message", "error", err)
                return database
        }

        startupMessage := cast.ToStringMap(string(startupMessageBytes))
        p.Logger.Trace("Startup message", "startupMessage", startupMessage, "client", client)
        if startupMessage != nil && client != nil {
                startupMsgParams := cast.ToStringMapString(startupMessage["Parameters"])
                if startupMsgParams != nil &&
                        startupMsgParams["database"] != "" &&
                        client["remote"] != "" {
                        if err := p.RedisClient.Set(
                                ctx, client["remote"],
                                startupMsgParams["database"],
                                time.Duration(0),
                        ).Err(); err != nil {
                                CacheErrorsCounter.Inc()
                                p.Logger.Debug("Failed to set cache", "error", err)
                        }
                        CacheSetsCounter.Inc()
                        p.Logger.Debug("Set the database in the cache for the current session",
                                "database", database, "client", client["remote"])
                        return startupMsgParams["database"]
                }
        }

        return database
}

1	package plugin
2
3	import (
4	"context"
5	"encoding/base64"
6	"strings"
7	"sync"
8	"time"
9
10	sdkAct "github.com/gatewayd-io/gatewayd-plugin-sdk/act"
11	"github.com/gatewayd-io/gatewayd-plugin-sdk/databases/postgres"
12	sdkPlugin "github.com/gatewayd-io/gatewayd-plugin-sdk/plugin"
13	v1 "github.com/gatewayd-io/gatewayd-plugin-sdk/plugin/v1"
14	apiV1 "github.com/gatewayd-io/gatewayd/api/v1"
15	"github.com/hashicorp/go-hclog"
16	goplugin "github.com/hashicorp/go-plugin"
17	goRedis "github.com/redis/go-redis/v9"
18	"github.com/spf13/cast"
19	"google.golang.org/grpc"
20	"google.golang.org/protobuf/proto"
21	)
22
23	type Plugin struct {
24	goplugin.GRPCPlugin
25	v1.GatewayDPluginServiceServer
26
27	Logger hclog.Logger
28
29	APIClient apiV1.GatewayDAdminAPIServiceClient
30
31	// Cache configuration.
32	RedisClient *goRedis.Client
33	RedisURL string
34	Expiry time.Duration
35	DefaultDBName string
36	ScanCount int64
37	ExitOnStartupError bool
38
39	UpdateCacheChannel chan *v1.Struct
40	WaitGroup *sync.WaitGroup
41
42	// Periodic invalidator configuration.
43	PeriodicInvalidatorEnabled bool
44	PeriodicInvalidatorStartDelay time.Duration
45	PeriodicInvalidatorInterval time.Duration
46	}
47
48	type CachePlugin struct {
49	goplugin.NetRPCUnsupportedPlugin
50	Impl Plugin
51	}
52
53	// Define a set for PostgreSQL date/time functions
54	// https://www.postgresql.org/docs/8.2/functions-datetime.html
55	var pgDateTimeFunctions = map[string]struct{}{
56	"AGE": {},
57	"CLOCK_TIMESTAMP": {},
58	"CURRENT_DATE": {},
59	"CURRENT_TIME": {},
60	"CURRENT_TIMESTAMP": {},
61	"LOCALTIME": {},
62	"LOCALTIMESTAMP": {},
63	"NOW": {},
64	"STATEMENT_TIMESTAMP": {},
65	"TIMEOFDAY": {},
66	"TRANSACTION_TIMESTAMP": {},
67	}
68
69	// NewCachePlugin returns a new instance of the CachePlugin.
70	func NewCachePlugin(impl Plugin) *CachePlugin {	7✔
71	return &CachePlugin{	7✔
72	NetRPCUnsupportedPlugin: goplugin.NetRPCUnsupportedPlugin{},	7✔
73	Impl: impl,	7✔
74	}	7✔
75	}	7✔
76
77	// GRPCServer registers the plugin with the gRPC server.
78	func (p CachePlugin) GRPCServer(_ goplugin.GRPCBroker, s *grpc.Server) error {	×
79	v1.RegisterGatewayDPluginServiceServer(s, &p.Impl)	×
80	return nil	×
UNCOV 81	}	×
82
83	// GRPCClient returns the plugin client.
84	func (p CachePlugin) GRPCClient(_ context.Context, _ goplugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {	×
85	return v1.NewGatewayDPluginServiceClient(c), nil	×
UNCOV 86	}	×
87
88	// GetPluginConfig returns the plugin config.
89	func (p *Plugin) GetPluginConfig(
90	_ context.Context, _ *v1.Struct,
91	) (*v1.Struct, error) {	1✔
92	GetPluginConfigCounter.Inc()	1✔
93	return v1.NewStruct(PluginConfig)	1✔
94	}	1✔
95
96	// OnTrafficFromClient is called when a request is received by GatewayD from the client.
97	func (p *Plugin) OnTrafficFromClient(
98	ctx context.Context, req *v1.Struct,
99	) (*v1.Struct, error) {	4✔
100	OnTrafficFromClientCounter.Inc()	4✔
101	req, err := postgres.HandleClientMessage(req, p.Logger)	4✔
102	if err != nil {	4✔
103	p.Logger.Info("Failed to handle client message", "error", err)	×
UNCOV 104	}	×
105
106	// This is used as a fallback if the database is not found in the startup message.
107	database := p.DefaultDBName	4✔
108	if database == "" {	8✔
109	client := cast.ToStringMapString(sdkPlugin.GetAttr(req, "client", nil))	4✔
110	database = p.getDBFromStartupMessage(ctx, req, database, client)	4✔
111		4✔
112	// Get the database from the cache if it's not found in the startup message or	4✔
113	// if the current request is not a startup message.	4✔
114	if database == "" {	6✔
115	database, err = p.RedisClient.Get(ctx, client["remote"]).Result()	2✔
116	if err != nil {	2✔
117	CacheErrorsCounter.Inc()	×
118	p.Logger.Debug("Failed to get cache", "error", err)	×
UNCOV 119	}	×
120	CacheGetsCounter.Inc()	2✔
121	p.Logger.Debug("Get the database in the cache for the current session",	2✔
122	"database", database, "client", client["remote"])	2✔
123	}
124	}
125
126	// If the database is still not found, return the response as is without caching.
127	// This might also happen if the cache is cleared while the client is still connected.
128	// In this case, the client should reconnect and the error will go away.
129	preconditions := sdkPlugin.GetAttr(req, "sslRequest", "") != "" \|\|	4✔
130	sdkPlugin.GetAttr(req, "saslInitialResponse", "") != "" \|\|	4✔
131	sdkPlugin.GetAttr(req, "cancelRequest", "") != ""	4✔
132	if database == "" && !preconditions {	4✔
133	p.Logger.Error(	×
134	"Database name not found or set in cache, startup message or plugin config. Skipping cache")	×
135	p.Logger.Error("Consider setting the database name in the plugin config or disabling the plugin if you don't need it")	×
136	return req, nil	×
UNCOV 137	}	×
138
139	query := cast.ToString(sdkPlugin.GetAttr(req, "query", ""))	4✔
140	request := cast.ToString(sdkPlugin.GetAttr(req, "request", ""))	4✔
141	server := cast.ToStringMapString(sdkPlugin.GetAttr(req, "server", ""))	4✔
142	cacheKey := strings.Join([]string{server["remote"], database, request}, ":")	4✔
143		4✔
144	if query == "" {	6✔
145	return req, nil	2✔
146	}	2✔
147
148	p.Logger.Trace("Query", "query", query)	2✔
149		2✔
150	// Clear the cache if the query is an insert, update or delete query.	2✔
151	p.invalidateDML(ctx, query)	2✔
152		2✔
153	// Check if the query is cached.	2✔
154	response, err := p.RedisClient.Get(ctx, cacheKey).Bytes()	2✔
155	if err != nil {	3✔
156	p.Logger.Debug("Failed to get cached response", "error", err)	1✔
157	}	1✔
158	CacheGetsCounter.Inc()	2✔
159		2✔
160	if response == nil {	3✔
161	// If the query is not cached, return the request as is.	1✔
162	CacheMissesCounter.Inc()	1✔
163	return req, nil	1✔
164	}	1✔
165
166	// If the query is cached, return the cached response.
167	signals, err := v1.NewList([]any{	1✔
168	sdkAct.Terminate().ToMap(),	1✔
169	sdkAct.Log("debug", "Returning cached response", map[string]any{	1✔
170	"cacheKey": []byte(cacheKey),	1✔
171	"plugin": PluginID.GetName(),	1✔
172	}).ToMap(),	1✔
173	})	1✔
174	if err != nil {	1✔
175	CacheErrorsCounter.Inc()	×
UNCOV 176	p.Logger.Error("Failed to create signals", "error", err)	×
177	} else {	1✔
178	CacheHitsCounter.Inc()	1✔
179	// Return the cached response.	1✔
180	req.Fields[sdkAct.Signals] = v1.NewListValue(signals)	1✔
181	req.Fields["response"] = v1.NewBytesValue(response)	1✔
182	}	1✔
183	return req, nil	1✔
184	}
185
186	// IsCacheNeeded determines if caching is needed.
187	func IsCacheNeeded(upperQuery string) bool {	16✔
188	// Iterate over each function name in the set of PostgreSQL date/time functions.	16✔
189	for function := range pgDateTimeFunctions {	130✔
190	if strings.Contains(upperQuery, function) {	125✔
191	// If the query contains a date/time function, caching is not needed.	11✔
192	return false	11✔
193	}	11✔
194	}
195	return true	5✔
196	}
197
198	func (p *Plugin) UpdateCache(ctx context.Context) {	3✔
199	defer p.WaitGroup.Done()	3✔
200	for {	10✔
201	serverResponse, ok := <-p.UpdateCacheChannel	7✔
202	if !ok {	10✔
203	p.Logger.Info("Channel closed, returning from function")	3✔
204	return	3✔
205	}	3✔
206
207	OnTrafficFromServerCounter.Inc()	4✔
208	resp, err := postgres.HandleServerMessage(serverResponse, p.Logger)	4✔
209	if err != nil {	4✔
210	p.Logger.Info("Failed to handle server message", "error", err)	×
UNCOV 211	}	×
212
213	rowDescription := cast.ToString(sdkPlugin.GetAttr(resp, "rowDescription", ""))	4✔
214	dataRow := cast.ToStringSlice(sdkPlugin.GetAttr(resp, "dataRow", []interface{}{}))	4✔
215	errorResponse := cast.ToString(sdkPlugin.GetAttr(resp, "errorResponse", ""))	4✔
216	request, isOk := sdkPlugin.GetAttr(resp, "request", nil).([]byte)	4✔
217	if !isOk {	4✔
218	request = []byte{}	×
UNCOV 219	}	×
220
221	response, isOk := sdkPlugin.GetAttr(resp, "response", nil).([]byte)	4✔
222	if !isOk {	4✔
223	response = []byte{}	×
UNCOV 224	}	×
225	server := cast.ToStringMapString(sdkPlugin.GetAttr(resp, "server", ""))	4✔
226		4✔
227	// This is used as a fallback if the database is not found in the startup message.	4✔
228		4✔
229	database := p.DefaultDBName	4✔
230	if database == "" {	8✔
231	client := cast.ToStringMapString(sdkPlugin.GetAttr(resp, "client", ""))	4✔
232	if client != nil && client["remote"] != "" {	8✔
233	database, err = p.RedisClient.Get(ctx, client["remote"]).Result()	4✔
234	if err != nil {	5✔
235	CacheErrorsCounter.Inc()	1✔
236	p.Logger.Debug("Failed to get cached response", "error", err)	1✔
237	}	1✔
238	CacheGetsCounter.Inc()	4✔
239	}
240	}
241
242	// If the database is still not found, return the response as is without caching.
243	// This might also happen if the cache is cleared while the client is still connected.
244	// In this case, the client should reconnect and the error will go away.
245	if database == "" {	5✔
246	p.Logger.Debug("Database name not found or set in cache, startup message or plugin config. " +	1✔
247	"Skipping cache")	1✔
248	p.Logger.Debug("Consider setting the database name in the " +	1✔
249	"plugin config or disabling the plugin if you don't need it")	1✔
250	continue	1✔
251	}
252
253	cacheKey := strings.Join([]string{server["remote"], database, string(request)}, ":")	3✔
254	if errorResponse != "" \|\| rowDescription == "" \|\| dataRow == nil \|\| len(dataRow) == 0 {	3✔
UNCOV 255	continue	×
256	}
257
258	query, err := postgres.GetQueryFromRequest(request)	3✔
259	if err != nil {	3✔
260	p.Logger.Debug("Failed to get query from request", "error", err)	×
UNCOV 261	continue	×
262	}
263
264	if !IsCacheNeeded(strings.ToUpper(query)) {	4✔
265	continue	1✔
266	}
267
268	// The request was successful and the response contains data. Cache the response.
269	if err := p.RedisClient.Set(ctx, cacheKey, response, p.Expiry).Err(); err != nil {	2✔
270	CacheErrorsCounter.Inc()	×
271	p.Logger.Debug("Failed to set cache", "error", err)	×
UNCOV 272	}	×
273	CacheSetsCounter.Inc()	2✔
274		2✔
275	tables, err := postgres.GetTablesFromQuery(query)	2✔
276	if err != nil {	2✔
277	p.Logger.Debug("Failed to get tables from query", "error", err)	×
UNCOV 278	continue	×
279	}
280
281	// Cache the table(s) used in each cached request. This is used to invalidate
282	// the cache when a rows is inserted, updated or deleted into that table.
283	for _, table := range tables {	4✔
284	requestQueryCacheKey := strings.Join([]string{table, cacheKey}, ":")	2✔
285	if err := p.RedisClient.Set(	2✔
286	ctx, requestQueryCacheKey, "", p.Expiry).Err(); err != nil {	2✔
287	CacheErrorsCounter.Inc()	×
288	p.Logger.Debug("Failed to set cache", "error", err)	×
UNCOV 289	}	×
290	CacheSetsCounter.Inc()	2✔
291	}
292	}
293	}
294
295	// OnTrafficFromServer is called when a response is received by GatewayD from the server.
296	func (p *Plugin) OnTrafficFromServer(
297	_ context.Context, resp *v1.Struct,
298	) (*v1.Struct, error) {	2✔
299	p.Logger.Debug("Traffic is coming from the server side")	2✔
300	if cloned, ok := proto.Clone(resp).(*v1.Struct); ok {	4✔
301	p.UpdateCacheChannel <- cloned	2✔
302	}	2✔
303	return resp, nil	2✔
304	}
305
306	func (p Plugin) OnClosed(ctx context.Context, req v1.Struct) (*v1.Struct, error) {	2✔
307	OnClosedCounter.Inc()	2✔
308	client := cast.ToStringMapString(sdkPlugin.GetAttr(req, "client", nil))	2✔
309	if client != nil {	4✔
310	if err := p.RedisClient.Del(ctx, client["remote"]).Err(); err != nil {	2✔
311	p.Logger.Debug("Failed to delete cache", "error", err)	×
312	CacheErrorsCounter.Inc()	×
UNCOV 313	}	×
314	p.Logger.Debug("Client closed", "client", client["remote"])	2✔
315	CacheDeletesCounter.Inc()	2✔
316	}
317	return req, nil	2✔
318	}
319
320	// invalidateDML invalidates the cache for the tables that are affected by the DML.
321	// This is done by getting the cached queries for each table and deleting them.
322	func (p *Plugin) invalidateDML(ctx context.Context, query string) {	4✔
323	// Check if the query is a UPDATE, INSERT or DELETE.	4✔
324	queryDecoded, err := base64.StdEncoding.DecodeString(query)	4✔
325	if err != nil {	4✔
326	p.Logger.Debug("Failed to decode query", "error", err)	×
327	return	×
UNCOV 328	}	×
329
330	queryMessage := cast.ToStringMapString(string(queryDecoded))	4✔
331	p.Logger.Trace("Query message", "query", queryMessage)	4✔
332		4✔
333	queryString := strings.ToUpper(queryMessage["String"])	4✔
334	// Ignore SELECT and WITH/SELECT queries.	4✔
335	// TODO: This is a naive approach, but query parsing has a cost.	4✔
336	if strings.HasPrefix(queryString, "SELECT") \|\|	4✔
337	(strings.HasPrefix(queryString, "WITH") &&	4✔
338	strings.Contains(queryString, "SELECT")) {	7✔
339	return	3✔
340	}	3✔
341
342	tables, err := postgres.GetTablesFromQuery(queryMessage["String"])	1✔
343	if err != nil {	1✔
344	p.Logger.Debug("Failed to get tables from query", "error", err)	×
345	return	×
UNCOV 346	}	×
347
348	p.Logger.Trace("Tables", "tables", tables)	1✔
349	for _, table := range tables {	2✔
350	// Invalidate the cache for the table.	1✔
351	// TODO: This is not efficient. We should be able to invalidate the cache	1✔
352	// for a specific key instead of invalidating the entire table.	1✔
353	pipeline := p.RedisClient.Pipeline()	1✔
354	var cursor uint64	1✔
355	for {	2✔
356	scanResult := p.RedisClient.Scan(ctx, cursor, table+":*", p.ScanCount)	1✔
357	if scanResult.Err() != nil {	1✔
358	CacheErrorsCounter.Inc()	×
359	p.Logger.Debug("Failed to scan keys", "error", scanResult.Err())	×
UNCOV 360	break	×
361	}
362	CacheScanCounter.Inc()	1✔
363		1✔
364	// Per each key, delete the cache entry and the table cache key itself.	1✔
365	var keys []string	1✔
366	keys, cursor = scanResult.Val()	1✔
367	CacheScanKeysCounter.Add(float64(len(keys)))	1✔
368	for _, tableKey := range keys {	2✔
369	// Invalidate the cache for the table.	1✔
370	cachedResponseKey := strings.TrimPrefix(tableKey, table+":")	1✔
371	pipeline.Del(ctx, cachedResponseKey)	1✔
372	// Invalidate the table cache key itself.	1✔
373	pipeline.Del(ctx, tableKey)	1✔
374	}	1✔
375
376	if cursor == 0 {	2✔
377	break	1✔
378	}
379	}
380
381	result, err := pipeline.Exec(ctx)	1✔
382	if err != nil {	1✔
383	p.Logger.Debug("Failed to execute pipeline", "error", err)	×
UNCOV 384	}	×
385
386	for _, res := range result {	3✔
387	if res.Err() != nil {	2✔
UNCOV 388	CacheErrorsCounter.Inc()	×
389	} else {	2✔
390	CacheDeletesCounter.Inc()	2✔
391	}	2✔
392	}
393	}
394	}
395
396	// getDBFromStartupMessage gets the database name from the startup message.
397	func (p *Plugin) getDBFromStartupMessage(
398	ctx context.Context,
399	req *v1.Struct,
400	database string,
401	client map[string]string,
402	) string {	4✔
403	// Try to get the database from the startup message, which is only sent once by the client.	4✔
404	// Store the database in the cache so that we can use it for subsequent requests.	4✔
405	startupMessageEncoded := cast.ToString(sdkPlugin.GetAttr(req, "startupMessage", ""))	4✔
406	if startupMessageEncoded == "" {	6✔
407	return database	2✔
408	}	2✔
409
410	startupMessageBytes, err := base64.StdEncoding.DecodeString(startupMessageEncoded)	2✔
411	if err != nil {	2✔
412	p.Logger.Debug("Failed to decode startup message", "error", err)	×
413	return database	×
UNCOV 414	}	×
415
416	startupMessage := cast.ToStringMap(string(startupMessageBytes))	2✔
417	p.Logger.Trace("Startup message", "startupMessage", startupMessage, "client", client)	2✔
418	if startupMessage != nil && client != nil {	4✔
419	startupMsgParams := cast.ToStringMapString(startupMessage["Parameters"])	2✔
420	if startupMsgParams != nil &&	2✔
421	startupMsgParams["database"] != "" &&	2✔
422	client["remote"] != "" {	4✔
423	if err := p.RedisClient.Set(	2✔
424	ctx, client["remote"],	2✔
425	startupMsgParams["database"],	2✔
426	time.Duration(0),	2✔
427	).Err(); err != nil {	2✔
428	CacheErrorsCounter.Inc()	×
429	p.Logger.Debug("Failed to set cache", "error", err)	×
UNCOV 430	}	×
431	CacheSetsCounter.Inc()	2✔
432	p.Logger.Debug("Set the database in the cache for the current session",	2✔
433	"database", database, "client", client["remote"])	2✔
434	return startupMsgParams["database"]	2✔
435	}
436	}
437
UNCOV 438	return database	×
439	}

gatewayd-io / gatewayd-plugin-cache / 22284276152

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous