715-get-or-create-with-null (piccolo-orm#716)

make sure that `get_or_create` works correctly with null values
powellnorma · Dec 13, 2022 · c1a95ae · c1a95ae
1 parent 7911252
commit c1a95ae
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 7 deletions.
diff --git a/piccolo/columns/combination.py b/piccolo/columns/combination.py
@@ -2,7 +2,11 @@
 
 import typing as t
 
-from piccolo.columns.operators.comparison import ComparisonOperator, Equal
+from piccolo.columns.operators.comparison import (
+    ComparisonOperator,
+    Equal,
+    IsNull,
+)
 from piccolo.custom_types import Combinable, Iterable
 from piccolo.querystring import QueryString
 from piccolo.utils.sql_values import convert_to_sql_value
@@ -52,25 +56,31 @@ def get_column_values(self) -> t.Dict[Column, t.Any]:
         This is used by `get_or_create` to know which values to assign if
         the row doesn't exist in the database.
 
-        For example, if we have:
+        For example, if we have::
 
-        (Band.name == 'Pythonistas') & (Band.popularity == 1000)
+            (Band.name == 'Pythonistas') & (Band.popularity == 1000)
 
-        We will return {Band.name: 'Pythonistas', Band.popularity: 1000}.
+        We will return::
+
+            {Band.name: 'Pythonistas', Band.popularity: 1000}.
 
         If the operator is anything besides equals, we don't return it, for
-        example:
+        example::
+
+            (Band.name == 'Pythonistas') & (Band.popularity > 1000)
 
-        (Band.name == 'Pythonistas') & (Band.popularity > 1000)
+        Returns::
 
-        Returns {Band.name: 'Pythonistas'}
+            {Band.name: 'Pythonistas'}
 
         """
         output = {}
         for combinable in (self.first, self.second):
             if isinstance(combinable, Where):
                 if combinable.operator == Equal:
                     output[combinable.column] = combinable.value
+                elif combinable.operator == IsNull:
+                    output[combinable.column] = None
             elif isinstance(combinable, And):
                 output.update(combinable.get_column_values())
 

diff --git a/tests/columns/test_combination.py b/tests/columns/test_combination.py
@@ -36,3 +36,21 @@ def test_not_in(self):
 
         with self.assertRaises(ValueError):
             Band.name.not_in([])
+
+
+class TestAnd(unittest.TestCase):
+    def test_get_column_values(self):
+        """
+        Make sure that we can extract the column values from an ``And``.
+
+        There was a bug with ``None`` values not working:
+
+        https://github.com/piccolo-orm/piccolo/issues/715
+
+        """
+        And_ = (Band.manager.is_null()) & (Band.name == "Pythonistas")
+        column_values = And_.get_column_values()
+
+        self.assertDictEqual(
+            column_values, {Band.name: "Pythonistas", Band.manager: None}
+        )