From e02e3906e75f4b831db62f475c29dba6665e238e 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 1/2] 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):

From 8cf9dc41c457acf81ec52eb73301ee0f3bf7935a Mon Sep 17 00:00:00 2001
From: Markus Unterwaditzer <markus-honeypot@unterwaditzer.net>
Date: Tue, 9 Aug 2022 19:32:13 +0200
Subject: [PATCH 2/2] fix: Use contribute_to_class instead

---
 django_spanner/__init__.py | 21 +++++++++------------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/django_spanner/__init__.py b/django_spanner/__init__.py
index ea42dbfb99..1e401db850 100644
--- a/django_spanner/__init__.py
+++ b/django_spanner/__init__.py
@@ -45,26 +45,23 @@
 register_functions()
 register_lookups()
 
-
 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 spanner_autofield_get_db_prep_value(self, value, connection, prepared=False):
-        value = old_get_db_prep_value(self, value, connection, prepared)
+    old_contribute_to_class = cls.contribute_to_class
+    
+    def spanner_autofield_contribute_to_class(self, cls, name, **kwargs):
+        if self.default == fields.NOT_PROVIDED:
+            using = django.db.router.db_for_write(cls)
+            if django.db.connections[using].settings_dict["ENGINE"] == "django_spanner":
+                self.default = gen_rand_int64
 
-        if (
-            connection.settings_dict["ENGINE"] == "django_spanner"
-            and value is None
-        ):
-            value = gen_rand_int64()
+        return old_contribute_to_class(self, cls, name, **kwargs)
 
-        return value
+    cls.contribute_to_class = spanner_autofield_contribute_to_class
 
-    cls.get_db_prep_value = spanner_autofield_get_db_prep_value
     cls.db_returning = False
     cls.validators = []