From 91c8f9fa9ac4698e0c79917b01b230bb623b143b Mon Sep 17 00:00:00 2001
From: Markus Unterwaditzer <markus-honeypot@unterwaditzer.net>
Date: Tue, 9 Aug 2022 18:32:01 +0200
Subject: [PATCH] Fix support for autofields in non-default databases

Fix https://github.com/googleapis/python-spanner-django/issues/783

Also refactored code a bit such that version detection is no longer used
to determine whether BigAutoField exists. BigAutoField exists in Django
versions older than 3.0, and the way this is currently monkeypatched
prevents us from using any BigAutoField in Spanner with Django 2.2.
---
 django_spanner/__init__.py | 46 ++++++++++++++++----------------------
 1 file changed, 19 insertions(+), 27 deletions(-)

diff --git a/django_spanner/__init__.py b/django_spanner/__init__.py
index ad4bc4e38d..ea42dbfb99 100644
--- a/django_spanner/__init__.py
+++ b/django_spanner/__init__.py
@@ -14,11 +14,7 @@
 
 import pkg_resources
 from google.cloud.spanner_v1 import JsonObject
-from django.db.models.fields import (
-    NOT_PROVIDED,
-    AutoField,
-    Field,
-)
+from django.db.models import fields
 
 from .expressions import register_expressions
 from .functions import register_functions
@@ -35,10 +31,6 @@
     USING_DJANGO_3 = True
 
 if USING_DJANGO_3:
-    from django.db.models.fields import (
-        SmallAutoField,
-        BigAutoField,
-    )
     from django.db.models import JSONField
 
 __version__ = pkg_resources.get_distribution("django-google-spanner").version
@@ -58,29 +50,29 @@ def gen_rand_int64():
     # Credit to https://stackoverflow.com/a/3530326.
     return uuid4().int & 0x7FFFFFFFFFFFFFFF
 
+def _fix_id_generator(cls):
+    old_get_db_prep_value = cls.get_db_prep_value
 
-def autofield_init(self, *args, **kwargs):
-    kwargs["blank"] = True
-    Field.__init__(self, *args, **kwargs)
+    def spanner_autofield_get_db_prep_value(self, value, connection, prepared=False):
+        value = old_get_db_prep_value(self, value, connection, prepared)
 
-    if (
-        django.db.connection.settings_dict["ENGINE"] == "django_spanner"
-        and self.default == NOT_PROVIDED
-    ):
-        self.default = gen_rand_int64
+        if (
+            connection.settings_dict["ENGINE"] == "django_spanner"
+            and value is None
+        ):
+            value = gen_rand_int64()
 
+        return value
 
-AutoField.__init__ = autofield_init
-AutoField.db_returning = False
-AutoField.validators = []
-if USING_DJANGO_3:
-    SmallAutoField.__init__ = autofield_init
-    BigAutoField.__init__ = autofield_init
-    SmallAutoField.db_returning = False
-    BigAutoField.db_returning = False
-    SmallAutoField.validators = []
-    BigAutoField.validators = []
+    cls.get_db_prep_value = spanner_autofield_get_db_prep_value
+    cls.db_returning = False
+    cls.validators = []
 
+for field_cls_name in ("AutoField", "BigAutoField", "SmallAutoField"):
+    if hasattr(fields, field_cls_name):
+        _fix_id_generator(getattr(fields, field_cls_name))
+
+if USING_DJANGO_3:
     def get_prep_value(self, value):
         # Json encoding and decoding for spanner is done in python-spanner.
         if not isinstance(value, JsonObject) and isinstance(value, dict):