feat: support AUTO_ID_CACHE (#56)

README.md

@@ -110,6 +110,25 @@ Migrate from `AUTO_INCREMENT` to `AUTO_RANDOM`:
 
 3. Finnaly, migrate it to `BigAutoRandomField(bigint)`.
 
+### Using `AUTO_ID_CACHE`
+
+[`AUTO_ID_CACHE`](https://docs.pingcap.com/tidb/stable/auto-increment#auto_id_cache) allow users to set the cache size for allocating the auto-increment ID, as you may know, TiDB guarantees that AUTO_INCREMENT values are monotonic (always increasing) on a per-server basis, but its value may appear to jump dramatically if an INSERT operation is performed against another TiDB Server, This is caused by the fact that each server has its own cache which is controlled by `AUTO_ID_CACHE`. But from TiDB v6.4.0, it introduces a centralized auto-increment ID allocating service, you can enable [*MySQL compatibility mode*](https://docs.pingcap.com/tidb/stable/auto-increment#mysql-compatibility-mode) by set `AUTO_ID_CACHE` to `1` when creating a table without losing performance.
+
+To use `AUTO_ID_CACHE` in Django, you can specify `tidb_auto_id_cache` in the model's `Meta` class as shown below when creating a new table:
+
+```python
+class MyModel(models.Model):
+    title = models.CharField(max_length=200)
+
+    class Meta:
+        tidb_auto_id_cache = 1
+```
+
+But there are some limitations:
+
+- `tidb_auto_id_cache` can only affect the table creation, after that it will be ignored even if you change it.
+- `tidb_auto_id_cache` only affects the `AUTO_INCREMENT` column.
+
 ## Supported versions
 
 - TiDB 4.0 and newer

django_tidb/__init__.py

@@ -13,8 +13,10 @@
 
 # Check Django compatibility before other imports which may fail if the
 # wrong version of Django is installed.
-from .functions import register_functions
+
+from .patch import monkey_patch
 
 __version__ = "3.2.1"
 
-register_functions()
+
+monkey_patch()

django_tidb/patch.py

@@ -0,0 +1,47 @@
+# Copyright 2021 PingCAP, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from django.db.models.functions import Chr
+from django.db.models import options
+from django.db.migrations import state
+
+
+def char(self, compiler, connection, **extra_context):
+    # TiDB doesn't support utf16
+    return self.as_sql(
+        compiler,
+        connection,
+        function="CHAR",
+        template="%(function)s(%(expressions)s USING utf8mb4)",
+        **extra_context,
+    )
+
+
+def patch_model_functions():
+    Chr.as_mysql = char
+
+
+def patch_model_options():
+    # Patch `tidb_auto_id_cache` to options.DEFAULT_NAMES,
+    # so that user can define it in model's Meta class.
+    options.DEFAULT_NAMES += ("tidb_auto_id_cache",)
+    # Because Django named import DEFAULT_NAMES in migrations,
+    # so we need to patch it again here.
+    # Django will record `tidb_auto_id_cache` in migration files,
+    # and then restore it when applying migrations.
+    state.DEFAULT_NAMES += ("tidb_auto_id_cache",)
+
+
+def monkey_patch():
+    patch_model_functions()
+    patch_model_options()

django_tidb/schema.py

@@ -60,3 +60,10 @@ def add_field(self, model, field):
             self.execute(self._create_unique_sql(model, [field.column]))
         else:
             super().add_field(model, field)
+
+    def table_sql(self, model):
+        sql, params = super().table_sql(model)
+        tidb_auto_id_cache = getattr(model._meta, "tidb_auto_id_cache", None)
+        if tidb_auto_id_cache is not None:
+            sql += " AUTO_ID_CACHE %s" % tidb_auto_id_cache
+        return sql, params

tests/test_tidb_auto_id_cache.py

@@ -0,0 +1,83 @@
+import re
+
+from django.db import models, connection
+from django.db.utils import ProgrammingError
+from django.test import TransactionTestCase
+from django.test.utils import isolate_apps
+
+
+AUTO_ID_CACHE_PATTERN = re.compile(r"\/\*T!\[auto_id_cache\] AUTO_ID_CACHE=(\d+) \*\/")
+
+
+class TiDBAutoIDCacheTests(TransactionTestCase):
+    available_apps = ["tidb"]
+
+    def get_auto_id_cache_info(self, table):
+        with connection.cursor() as cursor:
+            cursor.execute(
+                # It seems that SHOW CREATE TABLE is the only way to get the auto_random info.
+                # Use parameterized query will add quotes to the table name, which will cause syntax error.
+                f"SHOW CREATE TABLE {table}",
+            )
+            row = cursor.fetchone()
+            if row is None:
+                return None
+            match = AUTO_ID_CACHE_PATTERN.search(row[1])
+            if match:
+                return match.groups()[0]
+            return None
+
+    @isolate_apps("tidb")
+    def test_create_table_with_tidb_auto_id_cache_1(self):
+        class AutoIDCacheNode1(models.Model):
+            title = models.CharField(max_length=255)
+
+            class Meta:
+                app_label = "tidb"
+                tidb_auto_id_cache = 1
+
+        with connection.schema_editor() as editor:
+            editor.create_model(AutoIDCacheNode1)
+        self.assertEqual(
+            self.get_auto_id_cache_info(AutoIDCacheNode1._meta.db_table), "1"
+        )
+
+    @isolate_apps("tidb")
+    def test_create_table_with_tidb_auto_id_cache_non_1(self):
+        class AutoIDCacheNode2(models.Model):
+            title = models.CharField(max_length=255)
+
+            class Meta:
+                app_label = "tidb"
+                tidb_auto_id_cache = 10
+
+        with connection.schema_editor() as editor:
+            editor.create_model(AutoIDCacheNode2)
+        self.assertEqual(
+            self.get_auto_id_cache_info(AutoIDCacheNode2._meta.db_table), "10"
+        )
+
+    @isolate_apps("tidb")
+    def test_create_table_with_invalid_tidb_auto_id_cache(self):
+        class AutoIDCacheNode3(models.Model):
+            title = models.CharField(max_length=255)
+
+            class Meta:
+                app_label = "tidb"
+                tidb_auto_id_cache = "invalid"
+
+        with self.assertRaises(ProgrammingError):
+            with connection.schema_editor() as editor:
+                editor.create_model(AutoIDCacheNode3)
+
+    @isolate_apps("tidb")
+    def test_create_table_without_tidb_auto_id_cache(self):
+        class AutoIDCacheNode4(models.Model):
+            title = models.CharField(max_length=255)
+
+            class Meta:
+                app_label = "tidb"
+
+        with connection.schema_editor() as editor:
+            editor.create_model(AutoIDCacheNode4)
+        self.assertIsNone(self.get_auto_id_cache_info(AutoIDCacheNode4._meta.db_table))