Thanks to visit codestin.com
Credit goes to github.com

Skip to content

fix: add option to disable client-side pk generation #930

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jun 4, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 23 additions & 5 deletions django_spanner/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,11 @@
# do that.
from uuid import uuid4

RANDOM_ID_GENERATION_ENABLED_SETTING = "RANDOM_ID_GENERATION_ENABLED"

import pkg_resources
from django.conf.global_settings import DATABASES
from django.db import DEFAULT_DB_ALIAS
from google.cloud.spanner_v1 import JsonObject
from django.db.models.fields import (
NOT_PROVIDED,
Expand Down Expand Up @@ -64,11 +68,25 @@ def autofield_init(self, *args, **kwargs):
kwargs["blank"] = True
Field.__init__(self, *args, **kwargs)

if (
django.db.connection.settings_dict["ENGINE"] == "django_spanner"
and self.default == NOT_PROVIDED
):
self.default = gen_rand_int64
# The following behavior is chosen to prevent breaking changes with the original behavior.
# 1. We use a client-side randomly generated int64 value for autofields if Spanner is the
# default database, and DISABLE_RANDOM_ID_GENERATION has not been set.
# 2. If Spanner is one of the non-default databases, and no value at all has been set for
# DISABLE_RANDOM_ID_GENERATION, then we do not enable it. If there is a value for this
# configuration option, then we use that value.
databases = django.db.connections.databases
for db, config in databases.items():
default_enabled = str(db == DEFAULT_DB_ALIAS)
if (
config["ENGINE"] == "django_spanner"
and self.default == NOT_PROVIDED
and config.get(
RANDOM_ID_GENERATION_ENABLED_SETTING, default_enabled
).lower()
== "true"
):
self.default = gen_rand_int64
break


AutoField.__init__ = autofield_init
Expand Down
15 changes: 11 additions & 4 deletions tests/mockserver_tests/mock_server_test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import os
import unittest

from django.db import connection
from django.db import connection, connections
from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
import google.cloud.spanner_v1.types.type as spanner_type
import google.cloud.spanner_v1.types.result_set as result_set
Expand All @@ -37,6 +37,7 @@
start_mock_server,
)
from tests.mockserver_tests.mock_database_admin import DatabaseAdminServicer
from tests.settings import DATABASES


def add_result(sql: str, result: ResultSet):
Expand Down Expand Up @@ -178,11 +179,17 @@ def teardown_class(cls):
MockServerTestBase.server = None

def setup_method(self, test_method):
connection.settings_dict["OPTIONS"]["client"] = self.client
connection.settings_dict["OPTIONS"]["pool"] = self.pool
for db, config in DATABASES.items():
if config["ENGINE"] == "django_spanner":
connections[db].settings_dict["OPTIONS"][
"client"
] = self.client
connections[db].settings_dict["OPTIONS"]["pool"] = self.pool

def teardown_method(self, test_method):
connection.close()
for db, config in DATABASES.items():
if config["ENGINE"] == "django_spanner":
connections[db].close()
MockServerTestBase.spanner_service.clear_requests()
MockServerTestBase.database_admin_service.clear_requests()
self._client = None
Expand Down
72 changes: 70 additions & 2 deletions tests/mockserver_tests/test_basics.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,19 +11,21 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from google.cloud.spanner_v1 import (
BatchCreateSessionsRequest,
ExecuteSqlRequest,
CommitRequest,
)
from tests.mockserver_tests.mock_server_test_base import (
MockServerTestBase,
add_select1_result,
add_singer_query_result,
add_update_count,
)
from django.db import connection
from django.db import connection, models

from tests.mockserver_tests.models import Singer
from tests.settings import DATABASES


class TestBasics(MockServerTestBase):
Expand Down Expand Up @@ -60,3 +62,69 @@ def test_django_select_singer(self):
self.assertEqual(len(requests), 2)
self.assertIsInstance(requests[0], BatchCreateSessionsRequest)
self.assertIsInstance(requests[1], ExecuteSqlRequest)

def test_django_select_singer_using_other_db(self):
add_singer_query_result(
"SELECT tests_singer.id, tests_singer.first_name, tests_singer.last_name FROM tests_singer"
)
singers = Singer.objects.using("secondary").all()
self.assertEqual(len(singers), 2)
requests = self.spanner_service.requests
self.assertEqual(len(requests), 2)
self.assertIsInstance(requests[0], BatchCreateSessionsRequest)
self.assertIsInstance(requests[1], ExecuteSqlRequest)

def test_insert_singer(self):
add_update_count(
"INSERT INTO tests_singer "
"(id, first_name, last_name) "
"VALUES (@a0, @a1, @a2)",
1,
)
singer = Singer(first_name="test", last_name="test")
singer.save()
requests = self.spanner_service.requests
self.assertEqual(len(requests), 3)
self.assertIsInstance(requests[0], BatchCreateSessionsRequest)
self.assertIsInstance(requests[1], ExecuteSqlRequest)
self.assertIsInstance(requests[2], CommitRequest)
# The ExecuteSqlRequest should have 3 parameters:
# 1. first_name
# 2. last_name
# 3. client-side auto-generated primary key
self.assertEqual(len(requests[1].params), 3)

def test_insert_singer_with_disabled_random_primary_key(self):
for db, config in DATABASES.items():
if config["ENGINE"] == "django_spanner":
config["RANDOM_ID_GENERATION_ENABLED"] = "false"

# Define a class locally in this test method to ensure that
# it is initialized after disabling random ID generation.
class LocalSinger(models.Model):
first_name = models.CharField(max_length=200)
last_name = models.CharField(max_length=200)

try:
add_update_count(
"INSERT INTO tests_localsinger "
"(first_name, last_name) "
"VALUES (@a0, @a1)",
1,
)
singer = LocalSinger(first_name="test", last_name="test")
singer.save()
requests = self.spanner_service.requests
self.assertEqual(len(requests), 3)
self.assertIsInstance(requests[0], BatchCreateSessionsRequest)
self.assertIsInstance(requests[1], ExecuteSqlRequest)
self.assertIsInstance(requests[2], CommitRequest)
# The ExecuteSqlRequest should have 2 parameters:
# 1. first_name
# 2. last_name
# There should be no client-side auto-generated primary key.
self.assertEqual(len(requests[1].params), 2)
finally:
for db, config in DATABASES.items():
if config["ENGINE"] == "django_spanner":
config.pop("DISABLE_RANDOM_ID_GENERATION", None)
12 changes: 11 additions & 1 deletion tests/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,17 @@
"INSTANCE": INSTANCE_ID,
"NAME": DATABASE_NAME,
"TEST": {"NAME": DATABASE_NAME},
}
},
"secondary": {
"ENGINE": "django_spanner",
"PROJECT": PROJECT_ID,
"INSTANCE": INSTANCE_ID,
"NAME": DATABASE_NAME,
"TEST": {"NAME": DATABASE_NAME},
},
"other": {
"ENGINE": "django.db.backends.sqlite3",
},
}

SECRET_KEY = "spanner env secret key"
Expand Down
26 changes: 25 additions & 1 deletion tests/unit/django_spanner/test_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@


from .models import Author
from django.db import NotSupportedError, connection
from django.db import NotSupportedError, connection, connections
from django.db.models import Index
from django.db.models.fields import AutoField, IntegerField
from django_spanner import gen_rand_int64
Expand Down Expand Up @@ -433,3 +433,27 @@ def test_autofield_not_spanner_w_default(self):
assert gen_rand_int64 != field.default
assert mock_func == field.default
connection.settings_dict["ENGINE"] = "django_spanner"

def test_autofield_spanner_as_non_default_db_random_generation_enabled(
self,
):
"""Not Spanner as the default db, default for field not provided."""
connections.settings["default"]["ENGINE"] = "another_db"
connections.settings["secondary"]["ENGINE"] = "django_spanner"
connections.settings["secondary"][
"RANDOM_ID_GENERATION_ENABLED"
] = "true"
field = AutoField(name="field_name")
assert gen_rand_int64 == field.default
connections.settings["default"]["ENGINE"] = "django_spanner"
connections.settings["secondary"]["ENGINE"] = "django_spanner"
del connections.settings["secondary"]["RANDOM_ID_GENERATION_ENABLED"]

def test_autofield_random_generation_disabled(self):
"""Spanner, default is not provided."""
connections.settings["default"][
"RANDOM_ID_GENERATION_ENABLED"
] = "false"
field = AutoField(name="field_name")
assert gen_rand_int64 != field.default
del connections.settings["default"]["RANDOM_ID_GENERATION_ENABLED"]