Repository: G-Research/fasttrackml
Pull request #488 (2023-10-18 15:00:17)

Implement namespace caching layer.
What's inside the box:

- caching layer for namespace logic.
- attach namespace middleware to only mlflow and aim routes instead of global router.


Patch:
diff --git a/go.mod b/go.mod
index ddf8187..0eb2f27 100644
--- a/go.mod
+++ b/go.mod
@@ -23,6 +23,7 @@ require (
 	github.com/google/uuid v1.4.0
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/hetiansu5/urlquery v1.2.7
+	github.com/jackc/pgx/v5 v5.4.3
 	github.com/mattn/go-sqlite3 v1.14.16
 	github.com/pkg/errors v0.9.1
 	github.com/rotisserie/eris v0.5.4
@@ -76,7 +77,6 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
-	github.com/jackc/pgx/v5 v5.4.3 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/klauspost/asmfmt v1.3.2 // indirect
diff --git a/pkg/api/admin/service/namespace/service.go b/pkg/api/admin/service/namespace/service.go
index df5071a..f6e4008 100644
--- a/pkg/api/admin/service/namespace/service.go
+++ b/pkg/api/admin/service/namespace/service.go
@@ -100,7 +100,7 @@ func (s Service) DeleteNamespace(ctx context.Context, id uint) error {
 	if namespace == nil {
 		return eris.Errorf("namespace not found by id: %d", id)
 	}
-	if err := s.namespaceRepository.Delete(ctx, id); err != nil {
+	if err := s.namespaceRepository.Delete(ctx, namespace); err != nil {
 		return eris.Wrap(err, "error deleting namespace")
 	}
 	return nil
diff --git a/pkg/api/mlflow/dao/listener.go b/pkg/api/mlflow/dao/listener.go
new file mode 100644
index 0000000..72e7dae
--- /dev/null
+++ b/pkg/api/mlflow/dao/listener.go
@@ -0,0 +1,99 @@
+package dao
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/stdlib"
+	"github.com/rotisserie/eris"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+)
+
+// EventListenerProvider provides an interface to work with database event listener.
+type EventListenerProvider interface {
+	// Listen listens for incoming database events.
+	Listen() <-chan string
+	// GetChannelName returns channel name.
+	GetChannelName() string
+}
+
+// EventListener represents database event listener.
+type EventListener struct {
+	ctx        context.Context
+	channel    string
+	connection *stdlib.Conn
+}
+
+// NewEventListener creates new database event listener.
+func NewEventListener(ctx context.Context, db *gorm.DB, channel string) (*EventListener, error) {
+	eventListener := EventListener{
+		ctx:     ctx,
+		channel: channel,
+	}
+
+	switch db.Dialector.Name() {
+	case "postgres":
+		sqlDB, err := db.DB()
+		if err != nil {
+			return nil, eris.Wrap(err, "error getting db instance")
+		}
+		driverConnection, err := sqlDB.Conn(ctx)
+		if err != nil {
+			return nil, eris.Wrap(err, "error getting database connection")
+		}
+
+		if err := driverConnection.Raw(func(driverConn any) error {
+			var ok bool
+			if eventListener.connection, ok = driverConn.(*stdlib.Conn); !ok {
+				return eris.New(
+					"error getting underlying driver connection. driver connection has no type *stdlib.Conn",
+				)
+			}
+			return nil
+		}); err != nil {
+			return nil, eris.Wrap(err, "error getting underlying driver connection")
+		}
+
+		if _, err := eventListener.connection.Conn().Exec(
+			ctx, fmt.Sprintf("listen %s", channel),
+		); err != nil {
+			return nil, eris.Wrapf(err, "error creating listener for %s channel", channel)
+		}
+	}
+
+	return &eventListener, nil
+}
+
+// NewNamespaceListener creates new database event listener for Namespace entity.
+func NewNamespaceListener(ctx context.Context, db *gorm.DB) (*EventListener, error) {
+	return NewEventListener(ctx, db, "namespace_update_events")
+}
+
+// Listen listens for incoming database events.
+func (el EventListener) Listen() <-chan string {
+	ch := make(chan string)
+	// if listener not nil, then listen for incoming events from database.
+	// if listener is nil, then just return closed channel to do not do anything further.
+	if el.connection != nil {
+		go func() {
+			defer close(ch)
+			for {
+				notification, err := el.connection.Conn().WaitForNotification(context.Background())
+				if err != nil {
+					log.Errorf("error occurred while listening for the event: %+v", err)
+					return
+				}
+				ch <- notification.Payload
+			}
+		}()
+	} else {
+		close(ch)
+	}
+	return ch
+}
+
+// GetChannelName returns current channel name.
+func (el EventListener) GetChannelName() string {
+	return el.channel
+}
diff --git a/pkg/api/mlflow/dao/models/namespace.go b/pkg/api/mlflow/dao/models/namespace.go
index 252ad0c..ee0ff81 100644
--- a/pkg/api/mlflow/dao/models/namespace.go
+++ b/pkg/api/mlflow/dao/models/namespace.go
@@ -8,12 +8,12 @@ import (
 
 // Namespace represents model to work with `namespaces` table.
 type Namespace struct {
-	ID                  uint   `gorm:"primaryKey;autoIncrement"`
-	Code                string `gorm:"unique;index;not null"`
-	Description         string
-	CreatedAt           time.Time
-	UpdatedAt           time.Time
-	DeletedAt           gorm.DeletedAt `gorm:"index"`
-	DefaultExperimentID *int32         `gorm:"not null"`
-	Experiments         []Experiment   `gorm:"constraint:OnDelete:CASCADE"`
+	ID                  uint           `gorm:"primaryKey;autoIncrement" json:"id"`
+	Code                string         `gorm:"unique;index;not null" json:"code"`
+	Description         string         `json:"description"`
+	CreatedAt           time.Time      `json:"created_at"`
+	UpdatedAt           time.Time      `json:"updated_at"`
+	DeletedAt           gorm.DeletedAt `gorm:"index" json:"deleted_at"`
+	DefaultExperimentID *int32         `gorm:"not null" json:"default_experiment_id"`
+	Experiments         []Experiment   `gorm:"constraint:OnDelete:CASCADE" json:"experiments"`
 }
diff --git a/pkg/api/mlflow/dao/repositories/mock_namespace_repository_provider.go b/pkg/api/mlflow/dao/repositories/mock_namespace_repository_provider.go
index 5a49853..232284e 100644
--- a/pkg/api/mlflow/dao/repositories/mock_namespace_repository_provider.go
+++ b/pkg/api/mlflow/dao/repositories/mock_namespace_repository_provider.go
@@ -28,13 +28,13 @@ func (_m *MockNamespaceRepositoryProvider) Create(ctx context.Context, namespace
 	return r0
 }
 
-// Delete provides a mock function with given fields: ctx, id
-func (_m *MockNamespaceRepositoryProvider) Delete(ctx context.Context, id uint) error {
-	ret := _m.Called(ctx, id)
+// Delete provides a mock function with given fields: ctx, namespace
+func (_m *MockNamespaceRepositoryProvider) Delete(ctx context.Context, namespace *models.Namespace) error {
+	ret := _m.Called(ctx, namespace)
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, uint) error); ok {
-		r0 = rf(ctx, id)
+	if rf, ok := ret.Get(0).(func(context.Context, *models.Namespace) error); ok {
+		r0 = rf(ctx, namespace)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -42,25 +42,25 @@ func (_m *MockNamespaceRepositoryProvider) Delete(ctx context.Context, id uint)
 	return r0
 }
 
-// GetByCode provides a mock function with given fields: ctx, code
-func (_m *MockNamespaceRepositoryProvider) GetByCode(ctx context.Context, code string) (*models.Namespace, error) {
-	ret := _m.Called(ctx, code)
+// GetByCode provides a mock function with given fields: ctx, noCache, code
+func (_m *MockNamespaceRepositoryProvider) GetByCode(ctx context.Context, noCache bool, code string) (*models.Namespace, error) {
+	ret := _m.Called(ctx, noCache, code)
 
 	var r0 *models.Namespace
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, string) (*models.Namespace, error)); ok {
-		return rf(ctx, code)
+	if rf, ok := ret.Get(0).(func(context.Context, bool, string) (*models.Namespace, error)); ok {
+		return rf(ctx, noCache, code)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, string) *models.Namespace); ok {
-		r0 = rf(ctx, code)
+	if rf, ok := ret.Get(0).(func(context.Context, bool, string) *models.Namespace); ok {
+		r0 = rf(ctx, noCache, code)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*models.Namespace)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
-		r1 = rf(ctx, code)
+	if rf, ok := ret.Get(1).(func(context.Context, bool, string) error); ok {
+		r1 = rf(ctx, noCache, code)
 	} else {
 		r1 = ret.Error(1)
 	}
diff --git a/pkg/api/mlflow/dao/repositories/namespace.go b/pkg/api/mlflow/dao/repositories/namespace.go
index a388c9c..c1d96ff 100644
--- a/pkg/api/mlflow/dao/repositories/namespace.go
+++ b/pkg/api/mlflow/dao/repositories/namespace.go
@@ -17,9 +17,9 @@ type NamespaceRepositoryProvider interface {
 	// Update modifies the existing models.Namespace entity.
 	Update(ctx context.Context, namespace *models.Namespace) error
 	// Delete removes a namespace and it's associated experiments by its ID.
-	Delete(ctx context.Context, id uint) error
+	Delete(ctx context.Context, namespace *models.Namespace) error
 	// GetByCode returns namespace by its Code.
-	GetByCode(ctx context.Context, code string) (*models.Namespace, error)
+	GetByCode(ctx context.Context, noCache bool, code string) (*models.Namespace, error)
 	// GetByID returns namespace by its ID.
 	GetByID(ctx context.Context, id uint) (*models.Namespace, error)
 	// List returns all namespaces.
@@ -46,7 +46,7 @@ func (r NamespaceRepository) Create(ctx context.Context, namespace *models.Names
 	return nil
 }
 
-// Update persists modifications to models.Namespace entity.
+// Update modifies the existing models.Namespace entity.
 func (r NamespaceRepository) Update(ctx context.Context, namespace *models.Namespace) error {
 	if err := r.db.WithContext(ctx).Updates(namespace).Error; err != nil {
 		return eris.Wrap(err, "error updating namespace entity")
@@ -54,16 +54,16 @@ func (r NamespaceRepository) Update(ctx context.Context, namespace *models.Names
 	return nil
 }
 
-// Delete removes the  models.Namespace entity.
-func (r NamespaceRepository) Delete(ctx context.Context, id uint) error {
-	if err := r.db.WithContext(ctx).Delete(&models.Namespace{}, id).Error; err != nil {
+// Delete removes a namespace and it's associated experiments by its ID.
+func (r NamespaceRepository) Delete(ctx context.Context, namespace *models.Namespace) error {
+	if err := r.db.WithContext(ctx).Delete(namespace).Error; err != nil {
 		return eris.Wrap(err, "error deleting namespace entity")
 	}
 	return nil
 }
 
 // GetByCode returns namespace by its Code.
-func (r NamespaceRepository) GetByCode(ctx context.Context, code string) (*models.Namespace, error) {
+func (r NamespaceRepository) GetByCode(ctx context.Context, _ bool, code string) (*models.Namespace, error) {
 	var namespace models.Namespace
 	if err := r.db.WithContext(ctx).Where(
 		"code = ?", code,
diff --git a/pkg/api/mlflow/dao/repositories/namespace_cached.go b/pkg/api/mlflow/dao/repositories/namespace_cached.go
new file mode 100644
index 0000000..46a5590
--- /dev/null
+++ b/pkg/api/mlflow/dao/repositories/namespace_cached.go
@@ -0,0 +1,201 @@
+package repositories
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/gofiber/fiber/v2/log"
+	lru "github.com/hashicorp/golang-lru/v2"
+	"github.com/rotisserie/eris"
+	"gorm.io/gorm"
+
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao"
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/models"
+	"github.com/G-Research/fasttrackml/pkg/database"
+)
+
+// NamespaceEventAction represents Event action.
+type NamespaceEventAction string
+
+// Supported event actions.
+const (
+	NamespaceEventActionFetched = "fetched"
+	NamespaceEventActionCreated = "created"
+	NamespaceEventActionDeleted = "deleted"
+	NamespaceEventActionUpdated = "updated"
+)
+
+// NamespaceEvent represents database event.
+type NamespaceEvent struct {
+	Action    NamespaceEventAction `json:"action"`
+	Namespace models.Namespace     `json:"namespace"`
+}
+
+// NamespaceCachedRepository cached repository to work with `namespace` entity.
+type NamespaceCachedRepository struct {
+	db                  *gorm.DB
+	cache               *lru.Cache[string, models.Namespace]
+	listener            dao.EventListenerProvider
+	namespaceRepository NamespaceRepositoryProvider
+}
+
+// NewNamespaceCachedRepository creates new instance of cached repository to work with `namespace` entity.
+func NewNamespaceCachedRepository(
+	db *gorm.DB, listener dao.EventListenerProvider, namespaceRepository NamespaceRepositoryProvider,
+) (*NamespaceCachedRepository, error) {
+	cache, err := lru.New[string, models.Namespace](1000)
+	if err != nil {
+		return nil, eris.Wrap(err, "error creating lru cache for namespace entities")
+	}
+
+	// pre load all namespaces into cache.
+	var namespaces []models.Namespace
+	if err := db.Find(&namespaces).Error; err != nil {
+		return nil, eris.Wrapf(err, "error getting namespaces")
+	}
+	for _, namespace := range namespaces {
+		cache.Add(namespace.Code, namespace)
+	}
+
+	repository := NamespaceCachedRepository{
+		db:                  db,
+		cache:               cache,
+		listener:            listener,
+		namespaceRepository: namespaceRepository,
+	}
+
+	go func() {
+		for data := range listener.Listen() {
+			if err := repository.processEvent(data); err != nil {
+				log.Errorf(`error processing incoming event: %s, error: %+v`, data, err)
+			}
+		}
+	}()
+	return &repository, nil
+}
+
+// Create creates new models.Namespace entity.
+func (r NamespaceCachedRepository) Create(ctx context.Context, namespace *models.Namespace) error {
+	if err := r.namespaceRepository.Create(ctx, namespace); err != nil {
+		return eris.Wrap(err, "error creating cached namespace entity")
+	}
+
+	// trigger database event to notify current instance and
+	// other instances to create record in theirs local cache.
+	if err := r.sendEvent(NamespaceEventActionCreated, namespace); err != nil {
+		return eris.Wrap(err, "error sending database event")
+	}
+	return nil
+}
+
+// Update updates existing models.Namespace entity.
+func (r NamespaceCachedRepository) Update(ctx context.Context, namespace *models.Namespace) error {
+	if err := r.namespaceRepository.Update(ctx, namespace); err != nil {
+		return eris.Wrap(err, "error updating cached namespace entity")
+	}
+
+	// trigger database event to notify current instance and
+	// other instances to update record in theirs local cache.
+	if err := r.sendEvent(NamespaceEventActionUpdated, namespace); err != nil {
+		return eris.Wrap(err, "error sending database event")
+	}
+	return nil
+}
+
+// GetByCode returns namespace by its Code.
+func (r NamespaceCachedRepository) GetByCode(
+	ctx context.Context, noCache bool, code string,
+) (*models.Namespace, error) {
+	// TODO:dsuhinin - it is temporary solution for integration tests.
+	// Right now we have to disable cache, because in integration tests we directly work with database.
+	// It leads that cache state and database state is out of sync and different tests could just fail
+	if !noCache {
+		result, ok := r.cache.Get(code)
+		if ok {
+			return &result, nil
+		}
+	}
+
+	namespace, err := r.namespaceRepository.GetByCode(ctx, noCache, code)
+	if err != nil {
+		return nil, eris.Wrapf(err, "error getting cached namespace by code: %s", code)
+	}
+	if namespace == nil {
+		return nil, nil
+	}
+
+	// trigger database event to notify current instance and
+	// other instances to add record to theirs local cache.
+	if err := r.sendEvent(NamespaceEventActionFetched, namespace); err != nil {
+		return nil, eris.Wrap(err, "error sending database event")
+	}
+	return namespace, nil
+}
+
+// GetByID returns namespace by its ID.
+func (r NamespaceCachedRepository) GetByID(ctx context.Context, id uint) (*models.Namespace, error) {
+	return r.namespaceRepository.GetByID(ctx, id)
+}
+
+// Delete deletes existing models.Namespace entity.
+func (r NamespaceCachedRepository) Delete(ctx context.Context, namespace *models.Namespace) error {
+	if err := r.namespaceRepository.Delete(ctx, namespace); err != nil {
+		return eris.Wrap(err, "error deleting cached namespace entity")
+	}
+
+	// trigger database event to notify current instance and
+	// other instances to remove record from theirs local cache.
+	if err := r.sendEvent(NamespaceEventActionDeleted, namespace); err != nil {
+		return eris.Wrap(err, "error sending database event")
+	}
+	return nil
+}
+
+// List returns all namespaces.
+func (r NamespaceCachedRepository) List(ctx context.Context) ([]models.Namespace, error) {
+	return r.namespaceRepository.List(ctx)
+}
+
+// processEvent process incoming event from database.
+func (r NamespaceCachedRepository) processEvent(data string) error {
+	log.Debugf("got incoming namespace event: %s", data)
+	event := NamespaceEvent{}
+	if err := json.Unmarshal([]byte(data), &event); err != nil {
+		return eris.Wrap(err, "error unmarshaling incoming database event")
+	}
+	switch event.Action {
+	case NamespaceEventActionFetched:
+		r.cache.Add(event.Namespace.Code, event.Namespace)
+	case NamespaceEventActionCreated:
+		r.cache.Add(event.Namespace.Code, event.Namespace)
+	case NamespaceEventActionUpdated:
+		r.cache.Add(event.Namespace.Code, event.Namespace)
+	case NamespaceEventActionDeleted:
+		r.cache.Remove(event.Namespace.Code)
+	}
+	log.Debugf("namespace keys in local cache: %+v", r.cache.Keys())
+	return nil
+}
+
+// sendEvent sends database event.
+func (r NamespaceCachedRepository) sendEvent(action NamespaceEventAction, namespace *models.Namespace) error {
+	// skip event processing if current database is not a `postgres`.
+	if r.db.Dialector.Name() != database.PostgresDialectorName {
+		return nil
+	}
+
+	data, err := json.Marshal(NamespaceEvent{
+		Action:    action,
+		Namespace: *namespace,
+	})
+	if err != nil {
+		return eris.Wrap(err, "error serializing NamespaceEvent event")
+	}
+	if err := r.db.Exec(
+		fmt.Sprintf(`SELECT pg_notify('%s', '%s')`, r.listener.GetChannelName(), data),
+	).Error; err != nil {
+		return eris.Wrap(err, "error triggering 'pg_notify'")
+	}
+	return nil
+}
diff --git a/pkg/cmd/server.go b/pkg/cmd/server.go
index 8056f02..a9ca7b6 100644
--- a/pkg/cmd/server.go
+++ b/pkg/cmd/server.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"os"
@@ -25,6 +26,8 @@ import (
 	mlflowAPI "github.com/G-Research/fasttrackml/pkg/api/mlflow"
 	mlflowConfig "github.com/G-Research/fasttrackml/pkg/api/mlflow/config"
 	mlflowController "github.com/G-Research/fasttrackml/pkg/api/mlflow/controller"
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao"
+	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/repositories"
 	mlflowRepositories "github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/repositories"
 	mlflowService "github.com/G-Research/fasttrackml/pkg/api/mlflow/service"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/service/artifact"
@@ -65,13 +68,28 @@ func serverCmd(cmd *cobra.Command, args []string) error {
 	//nolint:errcheck
 	defer db.Close()
 
-	namespaceRepository := mlflowRepositories.NewNamespaceRepository(db.GormDB())
+	// 3. create database listeners.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	namespaceListener, err := dao.NewNamespaceListener(ctx, db.GormDB())
+	if err != nil {
+		return eris.Wrap(err, "error creating namespace notification listener")
+	}
 
-	// 3. init main HTTP server.
+	// 4. init main HTTP server.
+	namespaceRepository, err := repositories.NewNamespaceCachedRepository(
+		db.GormDB(), namespaceListener, repositories.NewNamespaceRepository(db.GormDB()),
+	)
+	if err != nil {
+		return eris.Wrap(err, "error creating cached namespace repository")
+	}
+
+	// 5. init main HTTP server.
 	server := initServer(mlflowConfig, namespaceRepository)
 
-	// 4. init `aim` api and ui routes.
-	aimAPI.AddRoutes(server.Group("/aim/api/"))
+	// 6. init `aim` api and ui routes.
+	router := server.Group("/aim/api/")
+	aimAPI.AddRoutes(router)
 	aimUI.AddRoutes(server)
 
 	artifactStorageFactory, err := storage.NewArtifactStorageFactory(mlflowConfig)
@@ -79,7 +97,7 @@ func serverCmd(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	// 5. init `mlflow` api and ui routes.
+	// 7. init `mlflow` api and ui routes.
 	// TODO:DSuhinin right now it might look scary. we prettify it a bit later.
 	mlflowAPI.NewRouter(
 		mlflowController.NewController(
@@ -108,14 +126,14 @@ func serverCmd(cmd *cobra.Command, args []string) error {
 	).Init(server)
 	mlflowUI.AddRoutes(server)
 
-	// 6. init `admin` UI routes.
+	// 8. init `admin` UI routes.
 	adminUI.NewRouter(
 		adminUIController.NewController(
 			namespace.NewService(namespaceRepository),
 		),
 	).Init(server)
 
-	// 7. init `chooser` ui routes.
+	// 9. init `chooser` ui routes.
 	chooser.NewRouter(
 		chooserController.NewController(
 			namespace.NewService(namespaceRepository),
@@ -182,7 +200,7 @@ func initDB(config *mlflowConfig.ServiceConfig) (database.DBProvider, error) {
 // initServer init HTTP server with base configuration.
 func initServer(
 	config *mlflowConfig.ServiceConfig,
-	namespaceRepository mlflowRepositories.NamespaceRepositoryProvider,
+	namespaceRepository repositories.NamespaceRepositoryProvider,
 ) *fiber.App {
 	server := fiber.New(fiber.Config{
 		BodyLimit:             16 * 1024 * 1024,
@@ -234,6 +252,7 @@ func initServer(
 		Format: "${status} - ${latency} ${method} ${path}\n",
 		Output: log.StandardLogger().Writer(),
 	}))
+
 	server.Use(namespaceMiddleware.New(namespaceRepository))
 
 	server.Get("/health", func(c *fiber.Ctx) error {
diff --git a/pkg/common/middleware/namespace/namespace.go b/pkg/common/middleware/namespace/namespace.go
index 9d17181..6dcbcca 100644
--- a/pkg/common/middleware/namespace/namespace.go
+++ b/pkg/common/middleware/namespace/namespace.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
+	"strconv"
 	"strings"
 
 	"github.com/gofiber/fiber/v2"
@@ -26,11 +27,15 @@ var namespaceRegexp = regexp.MustCompile(`^/ns/([^/]+)/`)
 // New creates new Middleware instance
 func New(namespaceRepository repositories.NamespaceRepositoryProvider) fiber.Handler {
 	return func(c *fiber.Ctx) (err error) {
-		log.Debugf("checking namespace for path: %s", c.Path())
+		noCache, err := strconv.ParseBool(string(c.Request().Header.Peek("no-cache")))
+		if err != nil {
+			noCache = false
+		}
+		log.Debugf("checking namespace for path: %s. no-cache: %t", c.Path(), noCache)
 		// if namespace exists in the request then try to process it, otherwise fallback to default namespace.
 		if matches := namespaceRegexp.FindStringSubmatch(c.Path()); matches != nil {
 			namespaceCode := strings.Clone(matches[1])
-			namespace, err := namespaceRepository.GetByCode(c.Context(), namespaceCode)
+			namespace, err := namespaceRepository.GetByCode(c.Context(), noCache, namespaceCode)
 			if err != nil {
 				return c.JSON(api.NewInternalError("error getting namespace with code: %s", namespaceCode))
 			}
@@ -45,7 +50,7 @@ func New(namespaceRepository repositories.NamespaceRepositoryProvider) fiber.Han
 			c.Locals(namespaceContextKey, namespace)
 			c.Path(strings.TrimPrefix(c.Path(), fmt.Sprintf("/ns/%s", namespaceCode)))
 		} else {
-			namespace, err := namespaceRepository.GetByCode(c.Context(), defaultNamespaceCode)
+			namespace, err := namespaceRepository.GetByCode(c.Context(), noCache, defaultNamespaceCode)
 			if err != nil {
 				return c.JSON(api.NewInternalError("error getting namespace with code: %s", defaultNamespaceCode))
 			}


Test patch:
diff --git a/pkg/api/admin/service/namespace/service_test.go b/pkg/api/admin/service/namespace/service_test.go
index 528c630..955d668 100644
--- a/pkg/api/admin/service/namespace/service_test.go
+++ b/pkg/api/admin/service/namespace/service_test.go
@@ -159,7 +159,7 @@ func TestService_DeleteNamespace_Ok(t *testing.T) {
 		Code: "code",
 	}
 	namespaceRepository.On(
-		"Delete", context.TODO(), uint(0),
+		"Delete", context.TODO(), &ns,
 	).Return(nil).On(
 		"GetByID", context.TODO(), uint(0),
 	).Return(&ns, nil)
diff --git a/tests/integration/golang/helpers/client.go b/tests/integration/golang/helpers/client.go
index 0d2b7ce..9175ac5 100644
--- a/tests/integration/golang/helpers/client.go
+++ b/tests/integration/golang/helpers/client.go
@@ -171,6 +171,9 @@ func (c *HttpClient) DoRequest(uri string, values ...any) error {
 			req.Header.Set(key, value)
 		}
 	}
+	// TODO:dsuhinin - right now set `no-cache` for all the requests to avoid any
+	// problem related to namespace caching logic. Remove it a bit later.
+	req.Header.Set("no-cache", "true")
 
 	// 7. send request data.
 	resp, err := c.client.Do(req)