Merge pull request #16932 from gavindsouza/db.query++

refactor: frappe.db.query

frappe/database/database.py

@@ -1026,14 +1026,10 @@ class Database(object):
 			if cache_count is not None:
 				return cache_count
 		query = self.query.get_sql(table=dt, filters=filters, fields=Count("*"), distinct=distinct)
-		if filters:
-			count = self.sql(query, debug=debug)[0][0]
-			return count
-		else:
-			count = self.sql(query, debug=debug)[0][0]
-			if cache:
-				frappe.cache().set_value("doctype:count:{}".format(dt), count, expires_in_sec=86400)
-			return count
+		count = self.sql(query, debug=debug)[0][0]
+		if not filters and cache:
+			frappe.cache().set_value("doctype:count:{}".format(dt), count, expires_in_sec=86400)
+		return count
 
 	@staticmethod
 	def format_date(date):

frappe/database/query.py

@@ -1,9 +1,12 @@
 import operator
 import re
+from functools import cached_property
 from typing import Any, Dict, List, Tuple, Union
 
 import frappe
 from frappe import _
+from frappe.boot import get_additional_filters_from_hooks
+from frappe.model.db_query import get_timespan_date_range
 from frappe.query_builder import Criterion, Field, Order, Table
 
 
@@ -90,6 +93,20 @@ def func_is(key, value):
 	return Field(key).isnotnull() if value.lower() == "set" else Field(key).isnull()
 
 
+def func_timespan(key: Field, value: str) -> frappe.qb:
+	"""Wrapper method for `TIMESPAN`
+
+	Args:
+	        key (str): field
+	        value (str): criterion
+
+	Returns:
+	        frappe.qb: `frappe.qb object with `TIMESPAN`
+	"""
+
+	return func_between(key, get_timespan_date_range(value))
+
+
 def make_function(key: Any, value: Union[int, str]):
 	"""returns fucntion query
 
@@ -100,7 +117,7 @@ def make_function(key: Any, value: Union[int, str]):
 	Returns:
 	        frappe.qb: frappe.qb object
 	"""
-	return OPERATOR_MAP[value[0]](key, value[1])
+	return OPERATOR_MAP[value[0].casefold()](key, value[1])
 
 
 def change_orderby(order: str):
@@ -123,7 +140,8 @@ def change_orderby(order: str):
 	return order[0], Order.desc
 
 
-OPERATOR_MAP = {
+# default operators
+OPERATOR_MAP: Dict[str, "function"] = {
 	"+": operator.add,
 	"=": operator.eq,
 	"-": operator.sub,
@@ -141,12 +159,36 @@ OPERATOR_MAP = {
 	"regex": func_regex,
 	"between": func_between,
 	"is": func_is,
+	"timespan": func_timespan,
+	# TODO: Add support for nested set
+	# TODO: Add support for custom operators (WIP) - via filters_config hooks
 }
 
 
 class Query:
 	tables: dict = {}
 
+	@cached_property
+	def OPERATOR_MAP(self):
+		# default operators
+		all_operators = OPERATOR_MAP.copy()
+
+		# update with site-specific custom operators
+		additional_filters_config = get_additional_filters_from_hooks()
+
+		if additional_filters_config:
+			from frappe.utils.commands import warn
+
+			warn("'filters_config' hook is not completely implemented yet in frappe.db.query engine")
+
+		for _operator, function in additional_filters_config.items():
+			if callable(function):
+				all_operators.update({_operator.casefold(): function})
+			elif isinstance(function, dict):
+				all_operators[_operator.casefold()] = frappe.get_attr(function.get("get_field"))()["operator"]
+
+		return all_operators
+
 	def get_condition(self, table: Union[str, Table], **kwargs) -> frappe.qb:
 		"""Get initial table object
 
@@ -227,14 +269,14 @@ class Query:
 		if isinstance(filters, list):
 			for f in filters:
 				if not isinstance(f, (list, tuple)):
-					_operator = OPERATOR_MAP[filters[1]]
+					_operator = self.OPERATOR_MAP[filters[1].casefold()]
 					if not isinstance(filters[0], str):
 						conditions = make_function(filters[0], filters[2])
 						break
 					conditions = conditions.where(_operator(Field(filters[0]), filters[2]))
 					break
 				else:
-					_operator = OPERATOR_MAP[f[-2]]
+					_operator = self.OPERATOR_MAP[f[-2].casefold()]
 					if len(f) == 4:
 						table_object = self.get_table(f[0])
 						_field = table_object[f[1]]
@@ -263,18 +305,14 @@ class Query:
 
 		for key in filters:
 			value = filters.get(key)
-			_operator = OPERATOR_MAP["="]
+			_operator = self.OPERATOR_MAP["="]
 
 			if not isinstance(key, str):
 				conditions = conditions.where(make_function(key, value))
 				continue
 			if isinstance(value, (list, tuple)):
-				if isinstance(value[1], (list, tuple)) or value[0] in list(OPERATOR_MAP.keys())[-4:]:
-					_operator = OPERATOR_MAP[value[0]]
-					conditions = conditions.where(_operator(Field(key), value[1]))
-				else:
-					_operator = OPERATOR_MAP[value[0]]
-					conditions = conditions.where(_operator(Field(key), value[1]))
+				_operator = self.OPERATOR_MAP[value[0].casefold()]
+				conditions = conditions.where(_operator(Field(key), value[1]))
 			else:
 				if value is not None:
 					conditions = conditions.where(_operator(Field(key), value))

frappe/model/db_query.py

@@ -1041,7 +1041,7 @@ def get_additional_filter_field(additional_filters_config, f, value):
 	return f
 
 
-def get_date_range(operator, value):
+def get_date_range(operator: str, value: str):
 	timespan_map = {
 		"1 week": "week",
 		"1 month": "month",
@@ -1054,7 +1054,10 @@ def get_date_range(operator, value):
 		"next": "next",
 	}
 
-	timespan = period_map[operator] + " " + timespan_map[value] if operator != "timespan" else value
+	if operator != "timespan":
+		timespan = f"{period_map[operator]} {timespan_map[value]}"
+	else:
+		timespan = value
 
 	return get_timespan_date_range(timespan)
 

frappe/utils/commands.py

@@ -59,7 +59,7 @@ def log(message, colour=""):
 	print(colour + message + end_line)
 
 
-def warn(message, category=None):
+def warn(message, category=None, stacklevel=2):
 	from warnings import warn
 
-	warn(message=message, category=category, stacklevel=2)
+	warn(message=message, category=category, stacklevel=stacklevel)

frappe/utils/data.py

@@ -714,7 +714,7 @@ def get_weekday(datetime: Optional[datetime.datetime] = None) -> str:
 	return weekdays[datetime.weekday()]
 
 
-def get_timespan_date_range(timespan):
+def get_timespan_date_range(timespan: str) -> Tuple[datetime.datetime, datetime.datetime]:
 	today = nowdate()
 	date_range_map = {
 		"last week": lambda: (