diff --git a/solutions/system_design/query_cache/query_cache_snippets.py b/solutions/system_design/query_cache/query_cache_snippets.py
index 19d3f5cdd63..4be7ab5df2e 100644
--- a/solutions/system_design/query_cache/query_cache_snippets.py
+++ b/solutions/system_design/query_cache/query_cache_snippets.py
@@ -58,7 +58,7 @@ def get(self, query):
 
         Accessing a node updates its position to the front of the LRU list.
         """
-        node = self.lookup[query]
+        node = self.lookup.get(query)
         if node is None:
             return None
         self.linked_list.move_to_front(node)
@@ -71,7 +71,7 @@ def set(self, results, query):
         If the entry is new and the cache is at capacity, removes the oldest entry
         before the new entry is added.
         """
-        node = self.map[query]
+        node = self.lookup.get(query)
         if node is not None:
             # Key exists in cache, update the value
             node.results = results