mik-laj commented on a change in pull request #7516: [AIRFLOW-6894] Prevent DB
query in example_dags
URL: https://github.com/apache/airflow/pull/7516#discussion_r383040768
##########
File path: tests/test_utils/asserts.py
##########
@@ -16,9 +16,53 @@
# under the License.
import re
+from contextlib import contextmanager
+
+from sqlalchemy import event
+
+from airflow.settings import engine
def assert_equal_ignore_multiple_spaces(case, first, second, msg=None):
def _trim(s):
return re.sub(r"\s+", " ", s.strip())
return case.assertEqual(_trim(first), _trim(second), msg)
+
+
+class CountQueriesResult:
+ def __init__(self):
+ self.count = 0
+
+
+class CountQueries(object):
+ """
+ Counts the number of queries sent to Airflow Database in a given context.
+
+ Does not support multiple processes. When a new process is started in
context, its queries will
+ not be included.
+ """
+ def __init__(self):
+ self.result = CountQueriesResult()
+
+ def __enter__(self):
+ event.listen(engine, "after_cursor_execute", self.after_cursor_execute)
+ return self.result
+
+ def __exit__(self, type, value, traceback):
+ event.remove(engine, "after_cursor_execute", self.after_cursor_execute)
+
+ def after_cursor_execute(self, *args, **kwargs):
+ self.result.count += 1
+
+
+count_queries = CountQueries
+
+
+@contextmanager
+def assert_queries_count(expected_count, message_fmt=None):
Review comment:
This decorator also wants to use it to prevent regression in scheduler
performance. Some methods are critical and I have optimized it to use very few
queries, but it can be easily broken. This context manager will allow us to
detect a regression regarding it.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
