diff --git a/backend/python/app/connectors/services/base_arango_service.py b/backend/python/app/connectors/services/base_arango_service.py index 9009e2a51..e657015f7 100644 --- a/backend/python/app/connectors/services/base_arango_service.py +++ b/backend/python/app/connectors/services/base_arango_service.py @@ -1,5 +1,3 @@ -"""ArangoDB service for interacting with the database""" - # pylint: disable=E1101, W0718 import asyncio import datetime @@ -9,50 +7,41 @@ from typing import Any, Dict, List, Optional, Set, Tuple import aiohttp # type: ignore -from arango import ArangoClient # type: ignore -from arango.database import TransactionDatabase # type: ignore -from fastapi import Request # type: ignore - from app.config.configuration_service import ConfigurationService -from app.config.constants.arangodb import ( - CollectionNames, - Connectors, - DepartmentNames, - GraphNames, - LegacyGraphNames, - OriginTypes, - RecordTypes, -) +from app.config.constants.arangodb import (CollectionNames, Connectors, + DepartmentNames, GraphNames, + LegacyGraphNames, OriginTypes, + RecordTypes) from app.config.constants.http_status_code import HttpStatusCode -from app.config.constants.service import DefaultEndpoints, config_node_constants +from app.config.constants.service import (DefaultEndpoints, + config_node_constants) from app.connectors.services.kafka_service import KafkaService -from app.models.entities import AppUserGroup, FileRecord, Record, RecordGroup, User -from app.schema.arango.documents import ( - agent_schema, - agent_template_schema, - app_schema, - department_schema, - file_record_schema, - mail_record_schema, - orgs_schema, - record_group_schema, - record_schema, - team_schema, - ticket_record_schema, - user_schema, - webpage_record_schema, -) -from app.schema.arango.edges import ( - basic_edge_schema, - belongs_to_schema, - is_of_type_schema, - permissions_schema, - record_relations_schema, - user_app_relation_schema, - user_drive_relation_schema, -) +from app.models.entities import (AppUserGroup, FileRecord, Record, RecordGroup, + User) +from app.schema.arango.documents import (agent_schema, agent_template_schema, + app_schema, department_schema, + file_record_schema, + mail_record_schema, orgs_schema, + record_group_schema, record_schema, + team_schema, ticket_record_schema, + user_schema, webpage_record_schema) +from app.schema.arango.edges import (basic_edge_schema, belongs_to_schema, + is_of_type_schema, permissions_schema, + record_relations_schema, + user_app_relation_schema, + user_drive_relation_schema) from app.schema.arango.graph import EDGE_DEFINITIONS from app.utils.time_conversion import get_epoch_timestamp_in_ms +from arango import ArangoClient # type: ignore +from arango.database import TransactionDatabase # type: ignore +from codeflash.code_utils.codeflash_wrap_decorator import \ + codeflash_performance_async +from fastapi import Request # type: ignore + +"""ArangoDB service for interacting with the database""" + + + # Collection definitions with their schemas NODE_COLLECTIONS = [ @@ -3576,6 +3565,7 @@ async def get_record_by_conversation_index( ) return None + @codeflash_performance_async async def get_record_owner_source_user_email( self, record_id: str, @@ -3604,7 +3594,11 @@ async def get_record_owner_source_user_email( """ db = transaction if transaction else self.db - cursor = db.aql.execute(query, bind_vars={"record_id": record_id}) + + # Offload the blocking db.aql.execute to a thread and make it async + cursor = await asyncio.to_thread( + db.aql.execute, query, bind_vars={"record_id": record_id} + ) result = next(cursor, None) return result @@ -3640,7 +3634,7 @@ async def get_record_by_path( db = transaction if transaction else self.db cursor = db.aql.execute( - query, bind_vars={"path": path} + query, bind_vars={"path": path}, batch_size=1 ) result = next(cursor, None)