Repository: adamchainz/django-mysql

Issue:
LAST_INSERT_ID function
Follow up to #12 - `LAST_INSERT_ID` would be a useful one to expose, can be used for a poor-man's `UPDATE... RETURNING`.



Discussion:

Resolving patch:
diff --git a/django_mysql/models/functions.py b/django_mysql/models/functions.py
--- a/django_mysql/models/functions.py
+++ b/django_mysql/models/functions.py
@@ -1,3 +1,4 @@
+from django.db import DEFAULT_DB_ALIAS, connections
 from django.db.models import CharField, IntegerField, TextField
 
 from django_mysql.shims import Func, Value
@@ -141,3 +142,26 @@ def __init__(self, expression, hash_len=512):
                 .format(",".join(str(x) for x in self.hash_lens))
             )
         super(SHA2, self).__init__(expression, Value(hash_len))
+
+
+# Information Functions
+
+class LastInsertId(Func):
+    function = 'LAST_INSERT_ID'
+
+    def __init__(self, expression=None):
+        if expression is not None:
+            super(LastInsertId, self).__init__(expression)
+        else:
+            super(LastInsertId, self).__init__()
+
+        self.output_field = IntegerField()
+
+    @classmethod
+    def get(cls, using=DEFAULT_DB_ALIAS):
+        # N.B. did try getting it from connection.connection.insert_id() (The
+        # MySQLdb query-free method) but it did not work with non-default
+        # database connections in Django, and the reason was not clear
+        with connections[using].cursor() as cursor:
+            cursor.execute("SELECT LAST_INSERT_ID()")
+            return cursor.fetchone()[0]


Test patch:
diff --git a/tests/django_mysql_tests/test_functions.py b/tests/django_mysql_tests/test_functions.py
--- a/tests/django_mysql_tests/test_functions.py
+++ b/tests/django_mysql_tests/test_functions.py
@@ -8,7 +8,7 @@
 
 from django_mysql.models.functions import (
     CRC32, ELT, MD5, SHA1, SHA2, Abs, Ceiling, ConcatWS, Field, Floor,
-    Greatest, Least, Round, Sign
+    Greatest, LastInsertId, Least, Round, Sign
 )
 from django_mysql_tests.models import Alphabet
 
@@ -222,6 +222,33 @@ def test_sha2_bad_hash_len(self):
             SHA2('a', 123)
 
 
+@requiresDatabaseFunctions
+class InformationFunctionTests(TestCase):
+
+    def test_last_insert_id(self):
+        Alphabet.objects.create(a=7891)
+        Alphabet.objects.update(a=LastInsertId('a') + 1)
+        lid = LastInsertId.get()
+        self.assertEqual(lid, 7891)
+
+    def test_last_insert_id_secondary_connection(self):
+        Alphabet.objects.using('secondary').create(a=9191)
+        Alphabet.objects.using('secondary').update(a=LastInsertId('a') + 9)
+        lid = LastInsertId.get(using='secondary')
+        self.assertEqual(lid, 9191)
+
+    def test_last_insert_id_in_query(self):
+        ab1 = Alphabet.objects.create(a=3719, b=717612)
+        ab2 = Alphabet.objects.create(a=1838, b=12636)
+
+        # Delete but store value of b and re-assign it to first second Alphabet
+        Alphabet.objects.filter(id=ab1.id, b=LastInsertId('b')).delete()
+        Alphabet.objects.filter(id=ab2.id).update(b=LastInsertId())
+
+        ab = Alphabet.objects.get()
+        self.assertEqual(ab.b, 717612)
+
+
 @skipIf(django.VERSION >= (1, 8),
         "Requires old Django version without Database Functions")
 class OldDjangoFunctionTests(TestCase):