Merge pull request #10 from iNishant/feature/get-or-none-wrapper

QuerysetGetOrNoneWrapper, type checking

README.md

@@ -61,8 +61,9 @@ assert results == [queryset1.count(), list(queryset2), ...]
 - Add a diagram in README depicting the time saved during network trips
 - Anything else which makes this better, open to ideas
 - Better readable way of accessing results (instead of `results[0]`, `results[1]`)
-- `FirstWrapper`, `LastWrapper` etc for lazy evaluating `.first()` and `.last()`
+- `QuerysetFirstWrapper`, `QuerysetLastWrapper` etc for lazy evaluating `.first()` and `.last()`
 - MySQL support as an experiment
+- "How it works" section/diagram?
 
 
 ## Notes

check-lint.sh

@@ -5,5 +5,5 @@ set -o pipefail
 
 ruff check .
 ruff format . --check
-
+mypy --ignore-missing-imports django_querysets_single_query_fetch/service.py
 echo 'Check lint passed!'

dev-requirements.txt

@@ -1,4 +1,6 @@
 ruff==0.3.7
 model-bakery==1.17.0
 build==1.2.1
-twine==5.0.0
+twine==5.0.0
+mypy==1.10.0
+mypy-extensions==1.0.0

django_querysets_single_query_fetch/service.py

@@ -3,7 +3,7 @@
 import logging
 import operator
 from decimal import Decimal
-from typing import Any, Union
+from typing import Any, List, Tuple, Union
 from uuid import UUID
 
 from django.core.exceptions import EmptyResultSet
@@ -38,6 +38,20 @@ def __init__(self, queryset: QuerySet) -> None:
         self.queryset = queryset
 
 
+class QuerysetGetOrNoneWrapper:
+    """
+    Wrapper around queryset to indicate that we want to fetch .get() or None
+    NOTE: this uses LIMIT 1 query so does not raise MultipleObjectsReturned
+    only returns the actual row or None (in case of no match)
+    """
+
+    def __init__(self, queryset: QuerySet) -> None:
+        self.queryset = queryset[:1]  # force limit 1
+
+
+QuerysetWrapperType = Union[QuerySet, QuerysetCountWrapper, QuerysetGetOrNoneWrapper]
+
+
 class QuerysetsSingleQueryFetch:
     """
     Executes multiple querysets in a single db query using json_build_object and returns results which
@@ -57,7 +71,7 @@ class QuerysetsSingleQueryFetch:
     assert results == [list(queryset) for queryset in querysets]
     """
 
-    def __init__(self, querysets: list[QuerySet]) -> None:
+    def __init__(self, querysets: list[QuerysetWrapperType]) -> None:
         self.querysets = querysets
 
     def _get_fetch_count_compiler_from_queryset(
@@ -154,22 +168,26 @@ def _get_fetch_count_compiler_from_queryset(
         outer_query.select_related = False
         return outer_query.get_compiler(using, elide_empty=elide_empty)
 
-    def _get_compiler_from_queryset(
-        self, queryset: Union[QuerySet, QuerysetCountWrapper]
-    ) -> Any:
+    def _get_compiler_from_queryset(self, queryset: QuerysetWrapperType) -> Any:
         """
         if queryset is wrapped in QuerysetCountWrapper, then we need to call _get_fetch_count_compiler_from_queryset
         else we can call get_compiler directly from queryset's query
         """
-        return (
-            self._get_fetch_count_compiler_from_queryset(
+
+        if isinstance(queryset, QuerysetCountWrapper):
+            compiler = self._get_fetch_count_compiler_from_queryset(
                 queryset.queryset, using=queryset.queryset.db
             )
-            if isinstance(queryset, QuerysetCountWrapper)
-            else queryset.query.get_compiler(using=queryset.db)
-        )
+        elif isinstance(queryset, QuerysetGetOrNoneWrapper):
+            _queryset = queryset.queryset
+            compiler = _queryset.query.get_compiler(using=_queryset.db)
+        else:
+            # queryset is the normal django queryset not wrapped by anything
+            compiler = queryset.query.get_compiler(using=queryset.db)
+
+        return compiler
 
-    def _get_django_sql_for_queryset(self, queryset: QuerySet) -> str:
+    def _get_django_sql_for_queryset(self, queryset: QuerysetWrapperType) -> str:
         """
         gets the sql that django would normally execute for the queryset, return empty string
         if queryset will always return empty irrespective of params ()
@@ -178,7 +196,7 @@ def _get_django_sql_for_queryset(self, queryset: QuerySet) -> str:
         # handle param quoting for IN queries (TODO: find some psycopg2 way to do this)
         # this is a bit hacky, but it works for now
 
-        quoted_params = ()
+        quoted_params: Tuple[Any, ...] = ()
         compiler = self._get_compiler_from_queryset(queryset=queryset)
         try:
             sql, params = compiler.as_sql(
@@ -326,28 +344,39 @@ def _get_instances_from_results_for_model_iterable(
         return instances
 
     def _convert_raw_results_to_final_queryset_results(
-        self, queryset: QuerySet, queryset_raw_results: list
+        self, queryset: QuerysetWrapperType, queryset_raw_results: list
     ):
         if isinstance(queryset, QuerysetCountWrapper):
             queryset_results = queryset_raw_results[0]["__count"]
-        elif issubclass(queryset._iterable_class, ModelIterable):
-            queryset_results = self._get_instances_from_results_for_model_iterable(
-                queryset=queryset, results=queryset_raw_results
-            )
-        elif issubclass(queryset._iterable_class, ValuesIterable):
-            queryset_results = queryset_raw_results
-        elif issubclass(queryset._iterable_class, FlatValuesListIterable):
-            queryset_results = [
-                list(row_dict.values())[0] for row_dict in queryset_raw_results
-            ]
-        elif issubclass(queryset._iterable_class, ValuesListIterable):
-            queryset_results = [
-                list(row_dict.values()) for row_dict in queryset_raw_results
-            ]
         else:
-            raise ValueError(
-                f"Unsupported queryset iterable class: {queryset._iterable_class}"
-            )
+            if isinstance(queryset, QuerysetGetOrNoneWrapper):
+                django_queryset = queryset.queryset
+            else:
+                django_queryset = queryset
+
+            if issubclass(django_queryset._iterable_class, ModelIterable):
+                queryset_results = self._get_instances_from_results_for_model_iterable(
+                    queryset=django_queryset, results=queryset_raw_results
+                )
+            elif issubclass(django_queryset._iterable_class, ValuesIterable):
+                queryset_results = queryset_raw_results
+            elif issubclass(django_queryset._iterable_class, FlatValuesListIterable):
+                queryset_results = [
+                    list(row_dict.values())[0] for row_dict in queryset_raw_results
+                ]
+            elif issubclass(django_queryset._iterable_class, ValuesListIterable):
+                queryset_results = [
+                    list(row_dict.values()) for row_dict in queryset_raw_results
+                ]
+            else:
+                raise ValueError(
+                    f"Unsupported queryset iterable class: {django_queryset._iterable_class}"
+                )
+
+        if isinstance(queryset, QuerysetGetOrNoneWrapper):
+            # convert queryset_results to either row or none
+            queryset_results = queryset_results[0] if queryset_results else None
+
         return queryset_results
 
     def execute(self) -> list[list[Any]]:
@@ -356,7 +385,7 @@ def execute(self) -> list[list[Any]]:
             for queryset in self.querysets
         ]
 
-        final_result_list = []
+        final_result_list: List[Any] = []
 
         for queryset_sql in django_sqls_for_querysets:
             if not queryset_sql:

setup.py

@@ -2,9 +2,9 @@
 
 setuptools.setup(
     name="django_querysets_single_query_fetch",
-    version="0.0.8",
+    version="0.0.9",
     description="Execute multiple Django querysets in a single SQL query",
-    long_description="",
+    long_description="Utility which executes multiple Django querysets over a single network/query call and returns results which would have been returned in normal evaluation of querysets",
     author="Nishant Singh",
     author_email="nishant.singh@mydukaan.io",
     license="Apache Software License",

testapp/tests/test_behaviour_for_postgres.py

@@ -6,7 +6,6 @@
 
 from django_querysets_single_query_fetch.service import (
     QuerysetsSingleQueryFetch,
-    QuerysetCountWrapper,
 )
 from testapp.models import OnlineStore, StoreProduct, StoreProductCategory
 
@@ -123,51 +122,6 @@ def test_executing_multiple_querysets_which_are_always_empty_is_handled(self):
             self.assertEqual(results[0], [])
             self.assertEqual(results[1], [])
 
-    def test_fetch_count(self):
-        """
-        - test fetch count works in single query
-        _ test fetch count works with filter querysets
-        _ test fetch count works with other querysets
-        """
-        # test fetch count works in single query
-        count_queryset = StoreProduct.objects.filter()
-        with self.assertNumQueries(1):
-            results = QuerysetsSingleQueryFetch(
-                querysets=[
-                    QuerysetCountWrapper(queryset=count_queryset),
-                ]
-            ).execute()
-        self.assertEqual(len(results), 1)
-        self.assertEqual(results[0], count_queryset.count())
-
-        # test fetch count works with filter querysets
-        count_filter_queryset = StoreProduct.objects.filter(id=self.product_1.id)
-        with self.assertNumQueries(1):
-            results = QuerysetsSingleQueryFetch(
-                querysets=[
-                    QuerysetCountWrapper(queryset=count_filter_queryset),
-                ]
-            ).execute()
-        self.assertEqual(len(results), 1)
-        self.assertEqual(results[0], count_filter_queryset.count())
-
-        # test fetch count works with other querysets
-        count_queryset = StoreProduct.objects.filter()
-        count_filter_queryset = StoreProduct.objects.filter(id=self.product_1.id)
-        queryset = StoreProduct.objects.filter()
-        with self.assertNumQueries(1):
-            results = QuerysetsSingleQueryFetch(
-                querysets=[
-                    QuerysetCountWrapper(queryset=count_queryset),
-                    QuerysetCountWrapper(queryset=count_filter_queryset),
-                    queryset,
-                ]
-            ).execute()
-        self.assertEqual(len(results), 3)
-        self.assertEqual(results[0], count_queryset.count())
-        self.assertEqual(results[1], count_filter_queryset.count())
-        self.assertEqual(results[2], list(queryset))
-
     def test_quotes_inside_the_string_in_select_query_will_work(self):
         name = "Ap's"
         product = baker.make(

testapp/tests/test_count_wrapper_for_postgres.py

@@ -0,0 +1,70 @@
+from datetime import datetime, timezone
+
+from django.test import TransactionTestCase
+from model_bakery import baker
+
+from django_querysets_single_query_fetch.service import (
+    QuerysetsSingleQueryFetch,
+    QuerysetCountWrapper,
+)
+from testapp.models import OnlineStore, StoreProduct, StoreProductCategory
+
+
+class QuerysetCountWrapperPostgresTestCase(TransactionTestCase):
+    def setUp(self) -> None:
+        self.today = datetime.now(tz=timezone.utc)
+        self.store = baker.make(OnlineStore, expired_on=self.today)
+        self.store = OnlineStore.objects.get(
+            id=self.store.id
+        )  # force refresh from db so that types are the default
+        # types
+        self.category = baker.make(StoreProductCategory, store=self.store)
+        self.product_1 = baker.make(StoreProduct, store=self.store, selling_price=50.22)
+        self.product_2 = baker.make(
+            StoreProduct, store=self.store, category=self.category, selling_price=100.33
+        )
+
+    def test_fetch_count(self):
+        """
+        - test fetch count works in single query
+        _ test fetch count works with filter querysets
+        _ test fetch count works with other querysets
+        """
+        # test fetch count works in single query
+        count_queryset = StoreProduct.objects.filter()
+        with self.assertNumQueries(1):
+            results = QuerysetsSingleQueryFetch(
+                querysets=[
+                    QuerysetCountWrapper(queryset=count_queryset),
+                ]
+            ).execute()
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0], count_queryset.count())
+
+        # test fetch count works with filter querysets
+        count_filter_queryset = StoreProduct.objects.filter(id=self.product_1.id)
+        with self.assertNumQueries(1):
+            results = QuerysetsSingleQueryFetch(
+                querysets=[
+                    QuerysetCountWrapper(queryset=count_filter_queryset),
+                ]
+            ).execute()
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0], count_filter_queryset.count())
+
+        # test fetch count works with other querysets
+        count_queryset = StoreProduct.objects.filter()
+        count_filter_queryset = StoreProduct.objects.filter(id=self.product_1.id)
+        queryset = StoreProduct.objects.filter()
+        with self.assertNumQueries(1):
+            results = QuerysetsSingleQueryFetch(
+                querysets=[
+                    QuerysetCountWrapper(queryset=count_queryset),
+                    QuerysetCountWrapper(queryset=count_filter_queryset),
+                    queryset,
+                ]
+            ).execute()
+        self.assertEqual(len(results), 3)
+        self.assertEqual(results[0], count_queryset.count())
+        self.assertEqual(results[1], count_filter_queryset.count())
+        self.assertEqual(results[2], list(queryset))

testapp/tests/test_get_or_none_wrapper_for_postgres.py

@@ -0,0 +1,62 @@
+from datetime import datetime, timezone
+from django_querysets_single_query_fetch.service import (
+    QuerysetsSingleQueryFetch,
+    QuerysetGetOrNoneWrapper,
+)
+from django.test import TransactionTestCase
+from model_bakery import baker
+
+from testapp.models import OnlineStore, StoreProduct, StoreProductCategory
+
+
+class QuerysetGetOrNoneWrapperPostgresTestCase(TransactionTestCase):
+    def setUp(self) -> None:
+        self.today = datetime.now(tz=timezone.utc)
+        self.store = baker.make(OnlineStore, expired_on=self.today)
+        self.store = OnlineStore.objects.get(
+            id=self.store.id
+        )  # force refresh from db so that types are the default
+        # types
+        self.category = baker.make(StoreProductCategory, store=self.store)
+        self.product_1 = baker.make(StoreProduct, store=self.store, selling_price=50.22)
+        self.product_2 = baker.make(
+            StoreProduct, store=self.store, category=self.category, selling_price=100.33
+        )
+
+    def test_get_or_none_wrapper_with_single_row_matching(self):
+        with self.assertNumQueries(1):
+            results = QuerysetsSingleQueryFetch(
+                querysets=[
+                    QuerysetGetOrNoneWrapper(
+                        StoreProduct.objects.filter(id=self.product_1.id)
+                    ),
+                ]
+            ).execute()
+        self.assertEqual(len(results), 1)
+        product = results[0]
+        self.assertEqual(product.id, self.product_1.id)
+
+    def test_get_or_none_wrapper_with_no_row_matching(self):
+        with self.assertNumQueries(1):
+            results = QuerysetsSingleQueryFetch(
+                querysets=[
+                    QuerysetGetOrNoneWrapper(StoreProduct.objects.filter(id=-1)),
+                ]
+            ).execute()
+        self.assertEqual(len(results), 1)
+        product = results[0]
+        self.assertIsNone(product)
+
+    def test_get_or_none_wrapper_with_multiple_rows_matching(self):
+        with self.assertNumQueries(1):
+            # get in this case can return either product 1 or product 2
+            results = QuerysetsSingleQueryFetch(
+                querysets=[
+                    QuerysetGetOrNoneWrapper(StoreProduct.objects.all()),
+                ]
+            ).execute()
+        self.assertEqual(len(results), 1)
+        product = results[0]
+        self.assertTrue(
+            (product.id == self.product_1.id) or (product.id == self.product_2.id)
+        )