fixup: cache server source code

Signed-off-by: kaikaila <[email protected]>

Author: kaikaila

backend/src/cache/client_manager.go

@@ -114,10 +114,14 @@ func initDBClient(params WhSvrDBParameters, initConnectionTimeout time.Duration)
 		glog.Fatalf("Failed to update the execution template type. Error: %s", err)
 	}
 
-	var tableNames []string
-	gormDB.Raw(`show tables`).Pluck("Tables_in_caches", &tableNames)
-	for _, tableName := range tableNames {
-		log.Printf("%s", tableName)
+	// List all tables using GORM Migrator API (works for both MySQL and PostgreSQL)
+	tableNames, err := gormDB.Migrator().GetTables()
+	if err != nil {
+		glog.Warningf("Failed to get table list: %v", err)
+	} else {
+		for _, tableName := range tableNames {
+			log.Printf("Table: %s", tableName)
+		}
 	}
 
 	return gormDB, dbDialect
@@ -214,6 +218,7 @@ func initDBDriver(params WhSvrDBParameters, initConnectionTimeout time.Duration)
 			db.Close()
 			glog.Fatalf("Failed to create database: %v", err)
 		}
+		log.Printf("Database created")
 		db.Close()
 
 		// Return DSN with target DB

backend/src/cache/server/admission.go

@@ -114,6 +114,11 @@ func doServeAdmitFunc(w http.ResponseWriter, r *http.Request, admit admitFunc, c
 		return errorResponse(admissionReviewReq.Request.UID, err), nil
 	}
 
+	// If no patch operations, return response without patch field
+	if patchOps == nil {
+		return allowedResponse(admissionReviewReq.Request.UID, nil), nil
+	}
+
 	patchBytes, err := json.Marshal(patchOps)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)

backend/src/cache/server/mutation.go

@@ -104,6 +104,13 @@ func MutatePodIfCached(req *v1beta1.AdmissionRequest, clientMgr ClientManagerInt
 	annotations := pod.ObjectMeta.Annotations
 	labels := pod.ObjectMeta.Labels
 
+	if annotations == nil {
+		annotations = make(map[string]string)
+	}
+	if labels == nil {
+		labels = make(map[string]string)
+	}
+
 	template, exists := getArgoTemplate(&pod)
 	if !exists {
 		return patches, nil

backend/src/cache/server/mutation_test.go

@@ -270,3 +270,19 @@ func TestMutatePodIfCachedWithTeamplateCleanup(t *testing.T) {
 	require.Equal(t, patchOperation[1].Op, OperationTypeAdd)
 	require.Equal(t, patchOperation[2].Op, OperationTypeAdd)
 }
+
+func TestMutatePodIfCachedWithNilAnnotations(t *testing.T) {
+	podWithoutAnnotations := *fakePod.DeepCopy()
+	podWithoutAnnotations.Annotations = nil
+	podWithoutAnnotations.Labels = map[string]string{
+		KFPCacheEnabledLabelKey: KFPCacheEnabledLabelValue,
+	}
+
+	patchOperations, err := MutatePodIfCached(
+		GetFakeRequestFromPod(&podWithoutAnnotations),
+		fakeClientManager,
+	)
+
+	assert.NotNil(t, patchOperations)
+	assert.Nil(t, err)
+}

backend/src/cache/storage/execution_cache_store.go

@@ -46,19 +46,31 @@ func (s *ExecutionCacheStore) GetExecutionCache(executionCacheKey string, cacheS
 	if cacheStaleness == 0 {
 		return nil, fmt.Errorf("CacheStaleness=0, Cache is disabled.")
 	}
-	r, err := s.db.Table("execution_caches").Where("ExecutionCacheKey = ?", executionCacheKey).Rows()
-	if err != nil {
-		return nil, fmt.Errorf("Failed to get execution cache: %q, err: %v", executionCacheKey, err)
+	// Use GORM's standard query method with struct field name to avoid case-sensitivity issues
+	var executionCaches []model.ExecutionCache
+	result := s.db.Where(&model.ExecutionCache{ExecutionCacheKey: executionCacheKey}).Find(&executionCaches)
+	if result.Error != nil {
+		return nil, fmt.Errorf("failed to get execution cache: %q, err: %v", executionCacheKey, result.Error)
 	}
-	defer r.Close()
-	executionCaches, err := s.scanRows(r, cacheStaleness)
-	if err != nil {
-		return nil, fmt.Errorf("Failed to scan rows on execution cache: %q, err: %v", executionCacheKey, err)
+
+	// Filter by cache staleness
+	var validCaches []*model.ExecutionCache
+	for i := range executionCaches {
+		cache := &executionCaches[i]
+		log.Println("Get id: " + strconv.FormatInt(cache.ID, 10))
+		log.Println("Get template: " + cache.ExecutionTemplate)
+		// maxCacheStaleness comes from the database entry.
+		// cacheStaleness is computed from the pods annotation and environment variables.
+		if (cache.MaxCacheStaleness < 0 || s.time.Now().UTC().Unix()-cache.StartedAtInSec <= cache.MaxCacheStaleness) &&
+			(cacheStaleness < 0 || s.time.Now().UTC().Unix()-cache.StartedAtInSec <= cacheStaleness) {
+			validCaches = append(validCaches, cache)
+		}
 	}
-	if len(executionCaches) == 0 {
+
+	if len(validCaches) == 0 {
 		return nil, fmt.Errorf("Execution cache not found with cache key: %q", executionCacheKey)
 	}
-	latestCache, err := getLatestCacheEntry(executionCaches)
+	latestCache, err := getLatestCacheEntry(validCaches)
 	if err != nil {
 		return nil, err
 	}
@@ -140,17 +152,15 @@ func getLatestCacheEntry(executionCaches []*model.ExecutionCache) (*model.Execut
 func (s *ExecutionCacheStore) CreateExecutionCache(executionCache *model.ExecutionCache) (*model.ExecutionCache, error) {
 	log.Printf("checking for existing row with cache key: %s before insertion", executionCache.ExecutionCacheKey)
 
-	r, err := s.db.Table("execution_caches").Where("ExecutionCacheKey = ?", executionCache.ExecutionCacheKey).Rows()
-	if err != nil {
-		log.Printf("Failed to get execution cache with key: %s, err: %v", executionCache.ExecutionCacheKey, err)
-		return nil, err
+	// Use GORM's standard query method with struct field name to avoid case-sensitivity issues
+	var existingCaches []model.ExecutionCache
+	result := s.db.Where(&model.ExecutionCache{ExecutionCacheKey: executionCache.ExecutionCacheKey}).Find(&existingCaches)
+	if result.Error != nil {
+		log.Printf("Failed to get execution cache with key: %s, err: %v", executionCache.ExecutionCacheKey, result.Error)
+		return nil, result.Error
 	}
 
-	rowCount := 0
-
-	for r.Next() {
-		rowCount++
-	}
+	rowCount := len(existingCaches)
 	log.Printf("number of rows returned for existing rows check: %d", rowCount)
 
 	if rowCount == 0 {