Fix bug with custom column names

It would fail :-)

Author: Photonios

psqlextra/compiler.py

@@ -124,13 +124,8 @@ def _rewrite_insert_nothing(self, sql, params, returning):
         # for conflicts
         conflict_target = self._build_conflict_target()
 
-        def format_field_name(field_name):
-            if isinstance(field_name, tuple):
-                return self.qn(field_name[0])
-            return self.qn(field_name)
-
         where_clause = ', '.join([
-            '{0} = %s'.format(format_field_name(field_name))
+            '{0} = %s'.format(self._format_field_name(field_name))
             for field_name in self.query.conflict_target
         ])
 
@@ -188,10 +183,20 @@ def _assert_valid_field(field_name):
                 'names and hstore key.'
             ) % str(field_name))
 
-        for field in self.query.conflict_target:
-            _assert_valid_field(field)
-            conflict_target.append(
-                self._format_field_name(field))
+        for field_name in self.query.conflict_target:
+            _assert_valid_field(field_name)
+
+            # special handling for hstore keys
+            if isinstance(field_name, tuple):
+                conflict_target.append(
+                    '(%s->\'%s\')' % (
+                        self._format_field_name(field_name),
+                        field_name[1]
+                    )
+                )
+            else:
+                conflict_target.append(
+                    self._format_field_name(field_name))
 
         return '(%s)' % ','.join(conflict_target)
 
@@ -207,28 +212,30 @@ def _get_model_field(self, name: str):
             no such field exists.
         """
 
+        field_name = name
+        if isinstance(field_name, tuple):
+            field_name = field_name[0]
+
         for field in self.query.model._meta.local_concrete_fields:
-            if field.name == name or field.column == name:
+            if field.name == field_name or field.column == field_name:
                 return field
 
         return None
 
-    def _format_field_name(self, name) -> str:
+    def _format_field_name(self, field_name) -> str:
         """Formats a field's name for usage in SQL.
 
         Arguments:
-            name:
+            field_name:
                 The field name to format.
 
         Returns:
             The specified field name formatted for
             usage in SQL.
         """
 
-        if isinstance(name, tuple):
-            return '(%s->\'%s\')' % name
-
-        return self.qn(name)
+        field = self._get_model_field(field_name)
+        return self.qn(field.column)
 
     def _format_field_value(self, field_name) -> str:
         """Formats a field's value for usage in SQL.
@@ -243,11 +250,9 @@ def _format_field_value(self, field_name) -> str:
             in SQL.
         """
 
-        if isinstance(field_name, tuple):
-            field_name, _ = field_name
-
+        field = self._get_model_field(field_name)
         return SQLInsertCompiler.prepare_value(
             self,
-            self._get_model_field(field_name),
-            getattr(self.query.objs[0], field_name)
+            field,
+            getattr(self.query.objs[0], field.name)
         )

tests/test_on_conflict.py

@@ -214,3 +214,19 @@ def test_on_conflict_outdated_model(conflict_action):
         .on_conflict(['title'], conflict_action)
         .insert_and_get(title='beer')
     )
+
+@pytest.mark.parametrize("conflict_action", CONFLICT_ACTIONS)
+def test_on_conflict_custom_column_names(conflict_action):
+    """Asserts that models with custom column names (models
+    where the column and field name are different) work properly."""
+
+    model = get_fake_model({
+        'title': models.CharField(max_length=140, unique=True, db_column='beer'),
+        'description': models.CharField(max_length=255, db_column='desc')
+    })
+
+    id = (
+        model.objects
+        .on_conflict(['title'], conflict_action)
+        .insert(title='yeey', description='great thing')
+    )