fix: Fix registry cache initialization

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

Author: ntkathole

sdk/python/feast/infra/registry/base_registry.py

@@ -784,7 +784,7 @@ def list_projects(
         raise NotImplementedError
 
     @abstractmethod
-    def proto(self) -> RegistryProto:
+    def proto(self, force_refresh: bool = False) -> RegistryProto:
         """
         Retrieves a proto version of the registry.
 

sdk/python/feast/infra/registry/caching_registry.py

@@ -426,55 +426,46 @@ def list_projects(
 
     def refresh(self, project: Optional[str] = None):
         try:
-            self.cached_registry_proto = self.proto()
+            self.cached_registry_proto = self.proto(force_refresh=True)
             self.cached_registry_proto_created = _utc_now()
         except Exception as e:
             logger.debug(f"Error while refreshing registry: {e}", exc_info=True)
 
     def _refresh_cached_registry_if_necessary(self):
         if self.cache_mode == "sync":
-
-            def is_cache_expired():
+            # Try acquiring the lock without blocking
+            if not self._refresh_lock.acquire(blocking=False):
+                logger.debug(
+                    "Skipping refresh if lock is already held by another thread"
+                )
+                return
+            try:
                 if self.cached_registry_proto == RegistryProto():
-                    if self.cached_registry_proto_ttl.total_seconds() == 0:
-                        return False
-                    else:
-                        return True
-
-                # Cache is expired if it's None or creation time is None
-                if (
-                    self.cached_registry_proto is None
-                    or not hasattr(self, "cached_registry_proto_created")
-                    or self.cached_registry_proto_created is None
-                ):
-                    return True
-
-                # Cache is expired if TTL > 0 and current time exceeds creation + TTL
-                if self.cached_registry_proto_ttl.total_seconds() > 0 and _utc_now() > (
-                    self.cached_registry_proto_created + self.cached_registry_proto_ttl
-                ):
-                    return True
-
-                return False
-
-            if is_cache_expired():
-                if not self._refresh_lock.acquire(blocking=False):
-                    logger.debug(
-                        "Skipping refresh if lock is already held by another thread"
-                    )
-                    return
-                try:
-                    logger.info(
-                        f"Registry cache expired(ttl: {self.cached_registry_proto_ttl.total_seconds()} seconds), so refreshing"
+                    expired = False
+                else:
+                    expired = (
+                        self.cached_registry_proto is None
+                        or self.cached_registry_proto_created is None
+                    ) or (
+                        self.cached_registry_proto_ttl.total_seconds() > 0
+                        and (
+                            _utc_now()
+                            > (
+                                self.cached_registry_proto_created
+                                + self.cached_registry_proto_ttl
+                            )
+                        )
                     )
+                if expired:
+                    logger.debug("Registry cache expired, so refreshing")
                     self.refresh()
-                except Exception as e:
-                    logger.debug(
-                        f"Error in _refresh_cached_registry_if_necessary: {e}",
-                        exc_info=True,
-                    )
-                finally:
-                    self._refresh_lock.release()
+            except Exception as e:
+                logger.debug(
+                    f"Error in _refresh_cached_registry_if_necessary: {e}",
+                    exc_info=True,
+                )
+            finally:
+                self._refresh_lock.release()  # Always release the lock safely
 
     def _start_thread_async_refresh(self, cache_ttl_seconds):
         self.refresh()

sdk/python/feast/infra/registry/registry.py

@@ -890,7 +890,7 @@ def teardown(self):
         """Tears down (removes) the registry."""
         self._registry_store.teardown()
 
-    def proto(self) -> RegistryProto:
+    def proto(self, force_refresh: bool = False) -> RegistryProto:
         return self.cached_registry_proto or RegistryProto()
 
     def _prepare_registry_for_changes(self, project_name: str):

sdk/python/feast/infra/registry/remote.py

@@ -594,7 +594,7 @@ def get_project_metadata(self, project: str, key: str) -> Optional[str]:
                 return meta[key]
         return None
 
-    def proto(self) -> RegistryProto:
+    def proto(self, force_refresh: bool = False) -> RegistryProto:
         return self.stub.Proto(Empty())
 
     def commit(self):

sdk/python/feast/infra/registry/snowflake.py

@@ -493,9 +493,11 @@ def _delete_object(
             if cursor.rowcount < 1 and not_found_exception:  # type: ignore
                 raise not_found_exception(name, project)
             self._set_last_updated_metadata(_utc_now(), project)
+            rowcount = cursor.rowcount
 
-            self.refresh()
-            return cursor.rowcount
+        self.refresh()
+
+        return rowcount
 
     def delete_permission(self, name: str, project: str, commit: bool = True):
         return self._delete_object(
@@ -1119,7 +1121,7 @@ def get_user_metadata(
         else:
             raise FeatureViewNotFoundException(feature_view.name, project=project)
 
-    def proto(self) -> RegistryProto:
+    def proto(self, force_refresh: bool = False) -> RegistryProto:
         r = RegistryProto()
         last_updated_timestamps = []
 
@@ -1128,6 +1130,18 @@ def process_project(project: Project):
             project_name = project.name
             last_updated_timestamp = project.last_updated_timestamp
 
+            try:
+                cached_project = self.get_project(project_name, True)
+            except ProjectObjectNotFoundException:
+                cached_project = None
+
+            allow_cache = False
+
+            if cached_project is not None and not force_refresh:
+                allow_cache = (
+                    last_updated_timestamp <= cached_project.last_updated_timestamp
+                )
+
             r.projects.extend([project.to_proto()])
             last_updated_timestamps.append(last_updated_timestamp)
 
@@ -1142,7 +1156,7 @@ def process_project(project: Project):
                 (self.list_validation_references, r.validation_references),
                 (self.list_permissions, r.permissions),
             ]:
-                objs: List[Any] = lister(project_name, allow_cache=False)  # type: ignore
+                objs: List[Any] = lister(project_name, allow_cache)  # type: ignore
                 if objs:
                     obj_protos = [obj.to_proto() for obj in objs]
                     for obj_proto in obj_protos:

sdk/python/feast/infra/registry/sql.py

@@ -838,7 +838,7 @@ def get_user_metadata(
             else:
                 raise FeatureViewNotFoundException(feature_view.name, project=project)
 
-    def proto(self) -> RegistryProto:
+    def proto(self, force_refresh: bool = False) -> RegistryProto:
         r = RegistryProto()
         last_updated_timestamps = []
 
@@ -847,6 +847,18 @@ def process_project(project: Project):
             project_name = project.name
             last_updated_timestamp = project.last_updated_timestamp
 
+            try:
+                cached_project = self.get_project(project_name, True)
+            except ProjectObjectNotFoundException:
+                cached_project = None
+
+            allow_cache = False
+
+            if cached_project is not None and not force_refresh:
+                allow_cache = (
+                    last_updated_timestamp <= cached_project.last_updated_timestamp
+                )
+
             r.projects.extend([project.to_proto()])
             last_updated_timestamps.append(last_updated_timestamp)
 
@@ -861,7 +873,7 @@ def process_project(project: Project):
                 (self.list_validation_references, r.validation_references),
                 (self.list_permissions, r.permissions),
             ]:
-                objs: List[Any] = lister(project_name, allow_cache=False)  # type: ignore
+                objs: List[Any] = lister(project_name, allow_cache)  # type: ignore
                 if objs:
                     obj_protos = [obj.to_proto() for obj in objs]
                     for obj_proto in obj_protos:
@@ -1055,9 +1067,12 @@ def _delete_object(
             if not self.purge_feast_metadata:
                 self._set_last_updated_metadata(_utc_now(), project, conn)
 
-            if self.cache_mode == "sync":
-                self.refresh()
-            return rows.rowcount
+            rowcount = rows.rowcount
+
+        if self.cache_mode == "sync":
+            self.refresh()
+
+        return rowcount
 
     def _get_object(
         self,

sdk/python/tests/unit/infra/registry/test_registry.py

@@ -189,22 +189,6 @@ def test_cache_expiry_triggers_refresh(registry):
         mock_refresh.assert_called_once()
 
 
-def test_empty_cache_refresh_with_ttl(registry):
-    """Test that empty cache is refreshed when TTL > 0"""
-    # Set up empty cache with TTL > 0
-    registry.cached_registry_proto = RegistryProto()
-    registry.cached_registry_proto_created = datetime.now(timezone.utc)
-    registry.cached_registry_proto_ttl = timedelta(seconds=10)  # TTL > 0
-
-    # Mock refresh to check if it's called
-    with patch.object(
-        CachingRegistry, "refresh", wraps=registry.refresh
-    ) as mock_refresh:
-        registry._refresh_cached_registry_if_necessary()
-        # Should refresh because cache is empty and TTL > 0
-        mock_refresh.assert_called_once()
-
-
 def test_empty_cache_no_refresh_with_infinite_ttl(registry):
     """Test that empty cache is not refreshed when TTL = 0 (infinite)"""
     # Set up empty cache with TTL = 0 (infinite)
@@ -227,7 +211,7 @@ def test_concurrent_cache_refresh_race_condition(registry):
     import time
 
     # Set up expired cache
-    registry.cached_registry_proto = RegistryProto()
+    registry.cached_registry_proto = "some_cached_data"  # Not empty
     registry.cached_registry_proto_created = datetime.now(timezone.utc) - timedelta(
         seconds=5
     )