Add a ODATA_SERVER_MONGO_COUNT_MAX_TIME_MS setting

Author: aarranz

odata_server/settings.py

@@ -1,3 +1,13 @@
+# Copyright (c) 2024 Future Internet Consulting and Development Solutions S.L.
+
 import os
 
-MONGO_SEARCH_MAX_TIME_MS = int(os.getenv("MONGO_SEARCH_MAX_TIME_MS", "30000"))
+MONGO_COUNT_MAX_TIME_MS = int(
+ os.getenv(
+ "ODATA_SERVER_MONGO_COUNT_MAX_TIME_MS",
+ os.getenv("ODATA_SERVER_MONGO_SEARCH_MAX_TIME_MS", "30000"),
+ )
+)
+MONGO_SEARCH_MAX_TIME_MS = int(
+ os.getenv("ODATA_SERVER_MONGO_SEARCH_MAX_TIME_MS", "30000")
+)

odata_server/utils/__init__.py

@@ -3,10 +3,11 @@
 import re
 
 import abnf
+import pymongo.database
 from bson.son import SON
 from flask import abort, request, url_for
 
-from odata_server import edm
+from odata_server import edm, settings
 
 from .common import crop_result, format_key_predicate
 from .flask import add_odata_annotations
@@ -148,9 +149,9 @@ def process_common_expr(tree, filters, entity_type, prefix, joinop="andExpr"):
  regex_literal = re.escape(parse_primitive_literal(args[1].children[0]))
  if methodExpr.name == "containsMethodCallExpr":
  filters[-1][field] = {
- "$regex": "(?!{})".format(regex_literal)
- if negation
- else regex_literal
+ "$regex": (
+  "(?!{})".format(regex_literal) if negation else regex_literal
+ )
  }
  elif methodExpr.name == "startsWithMethodCallExpr":
  filters[-1][field] = {
@@ -398,7 +399,14 @@ def prepare_anonymous_result(result, RootEntitySet, expand_details, prefix):
  return expand_result(RootEntitySet, expand_details, croped_result, prefix=prefix)
 
 
-def get_collection(mongo, RootEntitySet, subject, prefers, filters=None, count=False):
+def get_collection(
+ db: pymongo.database.Database,
+ RootEntitySet,
+ subject,
+ prefers,
+ filters=None,
+ count=False,
+):
  qs = parse_qs(request.query_string)
  anonymous = not isinstance(subject, edm.EntitySet)
 
@@ -440,7 +448,7 @@ def get_collection(mongo, RootEntitySet, subject, prefers, filters=None, count=F
  orderby = parse_orderby(qs.get("$orderby", ""))
 
  # Get the results
- mongo_collection = mongo.get_collection(RootEntitySet.mongo_collection)
+ mongo_collection = db.get_collection(RootEntitySet.mongo_collection)
  if prefix:
  seq_filter = {"Seq": filters.pop("Seq")} if "Seq" in filters else None
  pipeline = [
@@ -464,19 +472,29 @@ def get_collection(mongo, RootEntitySet, subject, prefers, filters=None, count=F
  pipeline.append({"$project": projection})
  pipeline.append({"$skip": offset})
  pipeline.append({"$limit": limit})
- results = mongo_collection.aggregate(pipeline)
+ results = mongo_collection.aggregate(
+ pipeline, maxTimeMS=settings.MONGO_SEARCH_MAX_TIME_MS
+ )
  else:
- cursor = mongo_collection.find(filters, projection)
+ cursor = mongo_collection.find(filters, projection).max_time_ms(
+ settings.MONGO_SEARCH_MAX_TIME_MS
+ )
  if len(orderby) > 0:
  cursor = cursor.sort(orderby)
  results = cursor.skip(offset).limit(limit)
 
  if count:
  if prefix == "":
- count = mongo_collection.count_documents(filters)
+ count = mongo_collection.count_documents(
+ filters, maxTimeMS=settings.MONGO_COUNT_MAX_TIME_MS
+ )
  else:
  basepipeline.append({"$count": "count"})
- result = tuple(mongo_collection.aggregate(basepipeline))
+ result = tuple(
+ mongo_collection.aggregate(
+ basepipeline, maxTimeMS=settings.MONGO_COUNT_MAX_TIME_MS
+ )
+ )
  count = 0 if len(result) == 0 else result[0]["count"]
  odata_count = count
  else:

odata_server/utils/json.py

@@ -43,7 +43,7 @@ def generate_collection_response(
  page_limit,
  prepare,
  odata_context,
- odata_count=None,
+ odata_count: int | None = None,
  prepare_kwargs={},
 ):
 

tests/test_flask.py

@@ -257,7 +257,9 @@ def test_get_entity_collection_api_orderby(self):
  )
  mongo.reset_mock()
  response = self.app.get("/Products?$orderby={}".format(orderby_expr))
- mongo.get_collection().find().sort.assert_called_once_with(expected)
+ mongo.get_collection().find().max_time_ms().sort.assert_called_once_with(
+ expected
+ )
  self.assertEqual(response.status_code, 200)
 
  def test_get_entity_collection_api_empty_filter(self):
@@ -277,7 +279,9 @@ def test_get_entity_collection_api_basic_filters(self):
  )
  for label, filter_expr in test_data:
  with self.subTest(msg=label):
- mongo.get_collection().find().skip().limit.return_value = iter(())
+ mongo.get_collection().find().max_time_ms().skip().limit.return_value = iter(
+ ()
+ )
  response = self.app.get("/Products?$filter={}".format(filter_expr))
  self.assertEqual(response.status_code, 200)
  # force processing of the body generator