mirror of
https://github.com/django/django.git
synced 2024-12-28 20:16:19 +00:00
102 lines
3.6 KiB
Python
102 lines
3.6 KiB
Python
import operator
|
|
import unittest
|
|
from collections import namedtuple
|
|
from contextlib import contextmanager
|
|
|
|
from django.db import connection, models
|
|
from django.test import TestCase
|
|
|
|
from ..models import Person
|
|
|
|
|
|
@unittest.skipUnless(connection.vendor == "postgresql", "PostgreSQL tests")
|
|
class ServerSideCursorsPostgres(TestCase):
|
|
cursor_fields = (
|
|
"name, statement, is_holdable, is_binary, is_scrollable, creation_time"
|
|
)
|
|
PostgresCursor = namedtuple("PostgresCursor", cursor_fields)
|
|
|
|
@classmethod
|
|
def setUpTestData(cls):
|
|
Person.objects.create(first_name="a", last_name="a")
|
|
Person.objects.create(first_name="b", last_name="b")
|
|
|
|
def inspect_cursors(self):
|
|
with connection.cursor() as cursor:
|
|
cursor.execute(
|
|
"SELECT {fields} FROM pg_cursors;".format(fields=self.cursor_fields)
|
|
)
|
|
cursors = cursor.fetchall()
|
|
return [self.PostgresCursor._make(cursor) for cursor in cursors]
|
|
|
|
@contextmanager
|
|
def override_db_setting(self, **kwargs):
|
|
for setting in kwargs:
|
|
original_value = connection.settings_dict.get(setting)
|
|
if setting in connection.settings_dict:
|
|
self.addCleanup(
|
|
operator.setitem, connection.settings_dict, setting, original_value
|
|
)
|
|
else:
|
|
self.addCleanup(operator.delitem, connection.settings_dict, setting)
|
|
|
|
connection.settings_dict[setting] = kwargs[setting]
|
|
yield
|
|
|
|
def assertUsesCursor(self, queryset, num_expected=1):
|
|
next(queryset) # Open a server-side cursor
|
|
cursors = self.inspect_cursors()
|
|
self.assertEqual(len(cursors), num_expected)
|
|
for cursor in cursors:
|
|
self.assertIn("_django_curs_", cursor.name)
|
|
self.assertFalse(cursor.is_scrollable)
|
|
self.assertFalse(cursor.is_holdable)
|
|
self.assertFalse(cursor.is_binary)
|
|
|
|
def asserNotUsesCursor(self, queryset):
|
|
self.assertUsesCursor(queryset, num_expected=0)
|
|
|
|
def test_server_side_cursor(self):
|
|
self.assertUsesCursor(Person.objects.iterator())
|
|
|
|
def test_values(self):
|
|
self.assertUsesCursor(Person.objects.values("first_name").iterator())
|
|
|
|
def test_values_list(self):
|
|
self.assertUsesCursor(Person.objects.values_list("first_name").iterator())
|
|
|
|
def test_values_list_flat(self):
|
|
self.assertUsesCursor(
|
|
Person.objects.values_list("first_name", flat=True).iterator()
|
|
)
|
|
|
|
def test_values_list_fields_not_equal_to_names(self):
|
|
expr = models.Count("id")
|
|
self.assertUsesCursor(
|
|
Person.objects.annotate(id__count=expr)
|
|
.values_list(expr, "id__count")
|
|
.iterator()
|
|
)
|
|
|
|
def test_server_side_cursor_many_cursors(self):
|
|
persons = Person.objects.iterator()
|
|
persons2 = Person.objects.iterator()
|
|
next(persons) # Open a server-side cursor
|
|
self.assertUsesCursor(persons2, num_expected=2)
|
|
|
|
def test_closed_server_side_cursor(self):
|
|
persons = Person.objects.iterator()
|
|
next(persons) # Open a server-side cursor
|
|
del persons
|
|
cursors = self.inspect_cursors()
|
|
self.assertEqual(len(cursors), 0)
|
|
|
|
def test_server_side_cursors_setting(self):
|
|
with self.override_db_setting(DISABLE_SERVER_SIDE_CURSORS=False):
|
|
persons = Person.objects.iterator()
|
|
self.assertUsesCursor(persons)
|
|
del persons # Close server-side cursor
|
|
|
|
with self.override_db_setting(DISABLE_SERVER_SIDE_CURSORS=True):
|
|
self.asserNotUsesCursor(Person.objects.iterator())
|