Issue #7 -- Use temp table for large number of parameters inside in clause

Author: absci

mssql/functions.py

@@ -6,7 +6,7 @@
 from django.db.models.functions import Cast
 from django.db.models.functions.math import ATan2, Log, Ln, Mod, Round
 from django.db.models.expressions import Case, Exists, OrderBy, When
-from django.db.models.lookups import Lookup
+from django.db.models.lookups import Lookup, In
 
 DJANGO3 = VERSION[0] >= 3
 
@@ -91,11 +91,29 @@ def sqlserver_orderby(self, compiler, connection):
     return self.as_sql(compiler, connection, template=template)
 
 
+def split_parameter_list_as_sql(self, compiler, connection):
+    # Insert In clause parameters 1000 at a time into a temp table.
+    lhs, _ = self.process_lhs(compiler, connection)
+    _, rhs_params = self.batch_process_rhs(compiler, connection)
+
+    with connection.cursor() as cursor:
+        cursor.execute("IF OBJECT_ID('tempdb.dbo.#Temp_params', 'U') IS NOT NULL DROP TABLE #Temp_params; ")
+        cursor.execute("CREATE TABLE #Temp_params (params int)")
+        for offset in range(0, len(rhs_params), 1000):
+            sqls_params = rhs_params[offset: offset + 1000]
+            sqls_params = ", ".join("({})".format(item) for item in sqls_params)
+            cursor.execute("INSERT INTO #Temp_params VALUES %s" % sqls_params)
+
+    in_clause = lhs + ' IN ' + '(SELECT params from #Temp_params)'
+
+    return in_clause, ()
+
 ATan2.as_microsoft = sqlserver_atan2
 Log.as_microsoft = sqlserver_log
 Ln.as_microsoft = sqlserver_ln
 Mod.as_microsoft = sqlserver_mod
 Round.as_microsoft = sqlserver_round
+In.split_parameter_list_as_sql = split_parameter_list_as_sql
 
 if DJANGO3:
     Lookup.as_microsoft = sqlserver_lookup

mssql/operations.py

@@ -20,6 +20,11 @@ class DatabaseOperations(BaseDatabaseOperations):
 
     cast_char_field_without_max_length = 'nvarchar(max)'
 
+    def max_in_list_size(self):
+        # The driver might add a few parameters
+        # chose a reasonable number less than 2100 limit
+        return 2048
+
     def _convert_field_to_tz(self, field_name, tzname):
         if settings.USE_TZ and not tzname == 'UTC':
             offset = self._get_utcoffset(tzname)

testapp/migrations/0010_pizza_topping.py

@@ -0,0 +1,28 @@
+# Generated by Django 3.1.7 on 2021-03-08 18:28
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('testapp', '0009_test_drop_table_with_foreign_key_reference_part2'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='Topping',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('name', models.CharField(max_length=30)),
+            ],
+        ),
+        migrations.CreateModel(
+            name='Pizza',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('name', models.CharField(max_length=50)),
+                ('toppings', models.ManyToManyField(to='testapp.Topping')),
+            ],
+        ),
+    ]

testapp/models.py

@@ -74,3 +74,17 @@ class TestRemoveOneToOneFieldModel(models.Model):
     # thats already is removed.
     # b = models.OneToOneField('self', on_delete=models.SET_NULL, null=True)
     a = models.CharField(max_length=50)
+
+
+class Topping(models.Model):
+    name = models.CharField(max_length=30)
+
+class Pizza(models.Model):
+    name = models.CharField(max_length=50)
+    toppings = models.ManyToManyField(Topping)
+
+    def __str__(self):
+        return "%s (%s)" % (
+            self.name,
+            ", ".join(topping.name for topping in self.toppings.all()),
+        )

testapp/tests/test_lookups.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from django.test import TestCase
+from ..models import Pizza, Topping
+
+class TestLookups(TestCase):
+    def test_large_number_of_params(self):
+        iterations = 3000
+        for i in range(iterations):
+            Pizza.objects.create(name="Pizza" + str(i))
+            Topping.objects.create(name="Topping" + str(i))
+        prefetch_result = Pizza.objects.prefetch_related('toppings')
+
+        self.assertEqual(len(prefetch_result), iterations)