Refactor exception handling in MemgraphStorage label methods

This commit is contained in:
yangdx 2025-11-14 11:01:26 +08:00
parent 423e4e927a
commit 8283c86bce

View file

@ -1050,12 +1050,12 @@ class MemgraphStorage(BaseGraphStorage):
"Memgraph driver is not initialized. Call 'await initialize()' first." "Memgraph driver is not initialized. Call 'await initialize()' first."
) )
workspace_label = self._get_workspace_label() result = None
async with self._driver.session( try:
database=self._DATABASE, default_access_mode="READ" workspace_label = self._get_workspace_label()
) as session: async with self._driver.session(
result = None database=self._DATABASE, default_access_mode="READ"
try: ) as session:
query = f""" query = f"""
MATCH (n:`{workspace_label}`) MATCH (n:`{workspace_label}`)
WHERE n.entity_id IS NOT NULL WHERE n.entity_id IS NOT NULL
@ -1075,13 +1075,11 @@ class MemgraphStorage(BaseGraphStorage):
f"[{self.workspace}] Retrieved {len(labels)} popular labels (limit: {limit})" f"[{self.workspace}] Retrieved {len(labels)} popular labels (limit: {limit})"
) )
return labels return labels
except Exception as e: except Exception as e:
logger.error( logger.error(f"[{self.workspace}] Error getting popular labels: {str(e)}")
f"[{self.workspace}] Error getting popular labels: {str(e)}" if result is not None:
) await result.consume()
if result is not None: return []
await result.consume()
return []
async def search_labels(self, query: str, limit: int = 50) -> list[str]: async def search_labels(self, query: str, limit: int = 50) -> list[str]:
"""Search labels with fuzzy matching """Search labels with fuzzy matching
@ -1103,12 +1101,12 @@ class MemgraphStorage(BaseGraphStorage):
if not query_lower: if not query_lower:
return [] return []
workspace_label = self._get_workspace_label() result = None
async with self._driver.session( try:
database=self._DATABASE, default_access_mode="READ" workspace_label = self._get_workspace_label()
) as session: async with self._driver.session(
result = None database=self._DATABASE, default_access_mode="READ"
try: ) as session:
cypher_query = f""" cypher_query = f"""
MATCH (n:`{workspace_label}`) MATCH (n:`{workspace_label}`)
WHERE n.entity_id IS NOT NULL WHERE n.entity_id IS NOT NULL
@ -1135,8 +1133,8 @@ class MemgraphStorage(BaseGraphStorage):
f"[{self.workspace}] Search query '{query}' returned {len(labels)} results (limit: {limit})" f"[{self.workspace}] Search query '{query}' returned {len(labels)} results (limit: {limit})"
) )
return labels return labels
except Exception as e: except Exception as e:
logger.error(f"[{self.workspace}] Error searching labels: {str(e)}") logger.error(f"[{self.workspace}] Error searching labels: {str(e)}")
if result is not None: if result is not None:
await result.consume() await result.consume()
return [] return []