1774297745

Build Type

github

Committed by GitHub

Commit Message

Merge 3b7db5958 into 6ea264820

Pull Request Pull Request #99: better updatedata command

Run Details

610 of 688 branches covered (88.66%)

Branch coverage included in aggregate %.

140 of 209 new or added lines in 9 files covered. (66.99%)

24 existing lines in 4 files now uncovered.

1451 of 1564 relevant lines covered (92.77%)

12.92 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.47

/computedfields/fast_update.py

"""
Drop-in to avoid bad update performance with ``bulk_update``.

The update is based on `UPDATE FROM VALUES` variants, which performs much better
for bigger changesets.

Currently supported DBMS:

- sqlite 3.33+ (3.33, 3.37 and 3.38 tested)
- postgresql (14 tested, should work with all versions 9.1+)
- mariabd 10.3+ (10.3 to 10.6 tested)
- mysql 8
"""

from django.db.models.functions import Cast
from django.db.models.expressions import Col
from django.db import transaction
from django.db.utils import ProgrammingError
import logging

logger = logging.getLogger(__name__)

# typing imports
from django.db.models import Field, QuerySet, Model
from django.db.models.sql.compiler import SQLCompiler
from typing import Any, Dict, Iterable, Sequence, Type
from django.db.backends.base.base import BaseDatabaseWrapper


def _cast_col_postgres(tname: str, field: Field, compiler: SQLCompiler, connection: Any) -> str:
    return Cast(Col(tname, field), output_field=field).as_postgresql(compiler, connection)[0]


def as_dummy(
    tname: str,
    pkname: str,
    fields: Sequence[Field],
    count: int,
    compiler: SQLCompiler,
    connection: Any
) -> str:
    return ''


def as_postgresql(
    tname: str,
    pkname: str,
    fields: Sequence[Field],
    count: int,
    compiler: SQLCompiler,
    connection: BaseDatabaseWrapper
) -> str:
    dname = 'd' if tname != 'd' else 'c'
    cols = ','.join(f'"{f.column}"={_cast_col_postgres(dname, f, compiler, connection)}' for f in fields)
    value = f'({",".join(["%s"] * (len(fields) + 1))})'
    values = ','.join([value] * count)
    dcols = f'"{pkname}",' + ','.join(f'"{f.column}"' for f in fields)
    where = f'"{tname}"."{pkname}"="{dname}"."{pkname}"'
    return f'UPDATE "{tname}" SET {cols} FROM (VALUES {values}) AS "{dname}" ({dcols}) WHERE {where}'


def as_sqlite(
    tname: str,
    pkname: str,
    fields: Sequence[Field],
    count: int,
    compiler: SQLCompiler,
    connection: BaseDatabaseWrapper
) -> str:
    dname = 'd' if tname != 'd' else 'c'
    cols = ','.join(f'"{f.column}"="{dname}"."column{i + 2}"' for i, f in enumerate(fields))
    value = f'({",".join(["%s"] * (len(fields) + 1))})'
    values = ','.join([value] * count)
    where = f'"{tname}"."{pkname}"="{dname}"."column1"'
    return f'UPDATE "{tname}" SET {cols} FROM (VALUES {values}) AS "{dname}" WHERE {where}'


def as_mysql(
    tname: str,
    pkname: str,
    fields: Sequence[Field],
    count: int,
    compiler: SQLCompiler,
    connection: BaseDatabaseWrapper
) -> str:
    dname = 'd' if tname != 'd' else 'c'
    cols = ','.join(f'`{f.column}`={dname}.{i+1}' for i, f in enumerate(fields))
    value = f'({",".join(["%s"] * (len(fields) + 1))})'
    values = ",".join([value] * (count + 1))
    on = f'`{tname}`.`{pkname}` = {dname}.0'
    return f'UPDATE `{tname}` INNER JOIN (VALUES {values}) AS {dname} ON {on} SET {cols}'


def as_mysql8(
    tname: str,
    pkname: str,
    fields: Sequence[Field],
    count: int,
    compiler: SQLCompiler,
    connection: BaseDatabaseWrapper
) -> str:
    dname = 'd' if tname != 'd' else 'c'
    cols = ','.join(f'`{f.column}`={dname}.column_{i+1}' for i, f in enumerate(fields))
    value = f'ROW({",".join(["%s"] * (len(fields) + 1))})'
    values = ",".join([value] * count)
    on = f'`{tname}`.`{pkname}` = {dname}.column_0'
    return f'UPDATE `{tname}` INNER JOIN (VALUES {values}) AS {dname} ON {on} SET {cols}'


QUERY = {
    'sqlite': as_sqlite,
    'postgresql': as_postgresql,
    'mysql': as_mysql,
    'mysql8': as_mysql8
}

# decide at runtine on connection level, which mysql impl to use
CONNECTION_HASHES: Dict[int, str] = {}

def _adjust_mysql(connection: BaseDatabaseWrapper) -> str:
    if connection.connection:
        vendor = CONNECTION_HASHES.get(hash(connection.connection), None)
        if vendor is not None:
            return vendor
    try:
        with transaction.atomic():
            with connection.cursor() as cur:
                cur.execute("SELECT foo.0 FROM (VALUES (0, 1), (1, 'zzz'),(2, 'yyy')) as foo")
            CONNECTION_HASHES[hash(connection.connection)] = 'mysql'
            return 'mysql'
    except ProgrammingError:
        pass
    try:
        with transaction.atomic():
            with connection.cursor() as cur:
                cur.execute("SELECT column_1 FROM (VALUES ROW(1, 'zzz'), ROW(2, 'yyy')) as foo")
            CONNECTION_HASHES[hash(connection.connection)] = 'mysql8'
            return 'mysql8'
    except ProgrammingError:
        pass
    CONNECTION_HASHES[hash(connection.connection)] = ''
    logger.warning('mysql backend without UPDATE FROM VALUES support, falling back to bulk_update')
    return ''


def fast_update(qs: QuerySet, objs: Iterable[Any], fieldnames: Sequence[str], batch_size: int = 1000) -> None:
    """
    Drop-in for `bulk_update` with much better update performance for model local fields (10 - 25x faster).
    Note that non-local fields (e.g. from multi table inheritance) still will be updated with `bulk_update`.

    Example usage:

        >>> # change big stack of instances
        >>> changed = []
        >>> queryset = MyModel.objects.filter(...)
        >>> for instance in queryset:
        ...     instance.name = ...
        ...     instance.age = ...
        ...     changed.append(instance)
        >>> # instead of:
        >>> queryset.bulk_update(changed, ['name', 'age'])
        >>> # update with:
        >>> fast_update(queryset, changed, ['name', 'age'])

    `qs` is used to derive the model class and db connection settings. It is **not** used to filter or
    narrow the records to be updated.

    `objs` is a sequence of altered model instances. The objects **must** be model instances the queryset
    would return. This is not explicitly tested, giving other model types is undefined behavior.

    `fieldnames` denoted the fields, that should be updated.

    `batch_size` controls how many records a single SQL statement will update. For tiny data loads per record
    (less and/or smaller fields), increase the batch size. For more or bigger fields, lower the batchsize
    or increase the database temp buffer before doing the big updates. Defaults to 1000 records.
    """
    model: Type[Model] = qs.model

    # filter all non model local fields --> still handled by bulk_update
    non_local_fieldnames = []
    local_fieldnames = []
    for f in fieldnames:
        if model._meta.get_field(f) not in model._meta.local_fields:
            non_local_fieldnames.append(f)
        else:
            local_fieldnames.append(f)
        
    # avoid more expensive doubled updates
    if non_local_fieldnames and len(local_fieldnames) < 2:
        return model.objects.bulk_update(objs, fieldnames, batch_size)
    
    if local_fieldnames:
        from django.db import connections

        tablename = model._meta.db_table
        pk_field = model._meta.pk
        if not pk_field:
            return model.objects.bulk_update(objs, fieldnames, batch_size)
        fields = [model._meta.get_field(f) for f in local_fieldnames]
        compiler = qs.query.get_compiler(qs.db)
        connection = connections[qs.db]

        # construct update data
        data = []
        counter = 0
        batches = []
        for o in objs:
            counter += 1
            # pk as first value to "join" on
            sub = [pk_field.get_db_prep_save(getattr(o, pk_field.attname), connection)]
            for field in fields:
                sub.append(field.get_db_prep_save(getattr(o, field.attname), connection))
            data += sub
            if counter >= batch_size:
                batches.append((counter, data))
                data = []
                counter = 0
        if data:
            batches.append((counter, data))

        sql = ''
        last_counter = -1
        vendor = connection.vendor
        if vendor == 'mysql':
            vendor = _adjust_mysql(connection)
        for counter, data in batches:
            # construct update string
            if counter != last_counter:
                sql = QUERY.get(vendor, as_dummy)(
                    tablename, pk_field.column, fields, counter, compiler, connection)
                if not sql:
                    # exist with bulk_update for non supported db backends
                    return model.objects.bulk_update(objs, fieldnames, batch_size)
            
            if vendor == 'mysql':
                # mysql needs data patch with (0,1,2,...) as first VALUES entry
                data = list(range(len(fields) + 1)) + data

            with connection.cursor() as cur:
                cur.execute(sql, data)
    
    if non_local_fieldnames:
        model.objects.bulk_update(objs, non_local_fieldnames, batch_size)


def check_support(using: str = 'default') -> bool:
    """
    Check support for `fast_update` with your current database backend.
    To test it on-the-fly, run these in `./manage.py shell`:

        >>> from computedfields.fast_update import check_support
        >>> check_support()
        True

    `using` denotes the database connection to be tested.
    """
    from django.db import connections
    connection = connections[using]
    if connection.vendor == 'postgresql':
        return True
    elif connection.vendor == 'sqlite':
        if not connection.connection:
            with connection.cursor():
                pass
        # grab the module to also work with pysqlite3
        import importlib
        _mod = importlib.import_module(connection.connection.__class__.__module__)
        major, minor, _ = _mod.sqlite_version_info
        if major >= 3 and minor > 32:
            return True
    elif connection.vendor == 'mysql':
        if _adjust_mysql(connection):
            return True
    return False

1	"""
2	Drop-in to avoid bad update performance with ``bulk_update``.
3
4	The update is based on `UPDATE FROM VALUES` variants, which performs much better
5	for bigger changesets.
6
7	Currently supported DBMS:
8
9	- sqlite 3.33+ (3.33, 3.37 and 3.38 tested)
10	- postgresql (14 tested, should work with all versions 9.1+)
11	- mariabd 10.3+ (10.3 to 10.6 tested)
12	- mysql 8
13	"""
14
15	from django.db.models.functions import Cast	15✔
16	from django.db.models.expressions import Col	15✔
17	from django.db import transaction	15✔
18	from django.db.utils import ProgrammingError	15✔
19	import logging	15✔
20
21	logger = logging.getLogger(__name__)	15✔
22
23	# typing imports
24	from django.db.models import Field, QuerySet, Model	15✔
25	from django.db.models.sql.compiler import SQLCompiler	15✔
26	from typing import Any, Dict, Iterable, Sequence, Type	15✔
27	from django.db.backends.base.base import BaseDatabaseWrapper	15✔
28
29
30	def _cast_col_postgres(tname: str, field: Field, compiler: SQLCompiler, connection: Any) -> str:	15✔
31	return Cast(Col(tname, field), output_field=field).as_postgresql(compiler, connection)[0]	1✔
32
33
34	def as_dummy(	15✔
35	tname: str,
36	pkname: str,
37	fields: Sequence[Field],
38	count: int,
39	compiler: SQLCompiler,
40	connection: Any
41	) -> str:
UNCOV 42	return ''	×
43
44
45	def as_postgresql(	15✔
46	tname: str,
47	pkname: str,
48	fields: Sequence[Field],
49	count: int,
50	compiler: SQLCompiler,
51	connection: BaseDatabaseWrapper
52	) -> str:
53	dname = 'd' if tname != 'd' else 'c'	1✔
54	cols = ','.join(f'"{f.column}"={_cast_col_postgres(dname, f, compiler, connection)}' for f in fields)	1✔
55	value = f'({",".join(["%s"] * (len(fields) + 1))})'	1✔
56	values = ','.join([value] * count)	1✔
57	dcols = f'"{pkname}",' + ','.join(f'"{f.column}"' for f in fields)	1✔
58	where = f'"{tname}"."{pkname}"="{dname}"."{pkname}"'	1✔
59	return f'UPDATE "{tname}" SET {cols} FROM (VALUES {values}) AS "{dname}" ({dcols}) WHERE {where}'	1✔
60
61
62	def as_sqlite(	15✔
63	tname: str,
64	pkname: str,
65	fields: Sequence[Field],
66	count: int,
67	compiler: SQLCompiler,
68	connection: BaseDatabaseWrapper
69	) -> str:
70	dname = 'd' if tname != 'd' else 'c'	1✔
71	cols = ','.join(f'"{f.column}"="{dname}"."column{i + 2}"' for i, f in enumerate(fields))	1✔
72	value = f'({",".join(["%s"] * (len(fields) + 1))})'	1✔
73	values = ','.join([value] * count)	1✔
74	where = f'"{tname}"."{pkname}"="{dname}"."column1"'	1✔
75	return f'UPDATE "{tname}" SET {cols} FROM (VALUES {values}) AS "{dname}" WHERE {where}'	1✔
76
77
78	def as_mysql(	15✔
79	tname: str,
80	pkname: str,
81	fields: Sequence[Field],
82	count: int,
83	compiler: SQLCompiler,
84	connection: BaseDatabaseWrapper
85	) -> str:
86	dname = 'd' if tname != 'd' else 'c'	1✔
87	cols = ','.join(f'`{f.column}`={dname}.{i+1}' for i, f in enumerate(fields))	1✔
88	value = f'({",".join(["%s"] * (len(fields) + 1))})'	1✔
89	values = ",".join([value] * (count + 1))	1✔
90	on = f'`{tname}`.`{pkname}` = {dname}.0'	1✔
91	return f'UPDATE `{tname}` INNER JOIN (VALUES {values}) AS {dname} ON {on} SET {cols}'	1✔
92
93
94	def as_mysql8(	15✔
95	tname: str,
96	pkname: str,
97	fields: Sequence[Field],
98	count: int,
99	compiler: SQLCompiler,
100	connection: BaseDatabaseWrapper
101	) -> str:
102	dname = 'd' if tname != 'd' else 'c'	1✔
103	cols = ','.join(f'`{f.column}`={dname}.column_{i+1}' for i, f in enumerate(fields))	1✔
104	value = f'ROW({",".join(["%s"] * (len(fields) + 1))})'	1✔
105	values = ",".join([value] * count)	1✔
106	on = f'`{tname}`.`{pkname}` = {dname}.column_0'	1✔
107	return f'UPDATE `{tname}` INNER JOIN (VALUES {values}) AS {dname} ON {on} SET {cols}'	1✔
108
109
110	QUERY = {	15✔
111	'sqlite': as_sqlite,
112	'postgresql': as_postgresql,
113	'mysql': as_mysql,
114	'mysql8': as_mysql8
115	}
116
117	# decide at runtine on connection level, which mysql impl to use
118	CONNECTION_HASHES: Dict[int, str] = {}	15✔
119
120	def _adjust_mysql(connection: BaseDatabaseWrapper) -> str:	15✔
121	if connection.connection:	2!
122	vendor = CONNECTION_HASHES.get(hash(connection.connection), None)	2✔
123	if vendor is not None:	2✔
124	return vendor	2✔
125	try:	2✔
126	with transaction.atomic():	2✔
127	with connection.cursor() as cur:	2✔
128	cur.execute("SELECT foo.0 FROM (VALUES (0, 1), (1, 'zzz'),(2, 'yyy')) as foo")	2✔
129	CONNECTION_HASHES[hash(connection.connection)] = 'mysql'	1✔
130	return 'mysql'	1✔
131	except ProgrammingError:	1✔
132	pass	1✔
133	try:	1✔
134	with transaction.atomic():	1✔
135	with connection.cursor() as cur:	1✔
136	cur.execute("SELECT column_1 FROM (VALUES ROW(1, 'zzz'), ROW(2, 'yyy')) as foo")	1✔
137	CONNECTION_HASHES[hash(connection.connection)] = 'mysql8'	1✔
138	return 'mysql8'	1✔
139	except ProgrammingError:	×
140	pass	×
141	CONNECTION_HASHES[hash(connection.connection)] = ''	×
142	logger.warning('mysql backend without UPDATE FROM VALUES support, falling back to bulk_update')	×
UNCOV 143	return ''	×
144
145
146	def fast_update(qs: QuerySet, objs: Iterable[Any], fieldnames: Sequence[str], batch_size: int = 1000) -> None:	15✔
147	"""
148	Drop-in for `bulk_update` with much better update performance for model local fields (10 - 25x faster).
149	Note that non-local fields (e.g. from multi table inheritance) still will be updated with `bulk_update`.
150
151	Example usage:
152
153	>>> # change big stack of instances
154	>>> changed = []
155	>>> queryset = MyModel.objects.filter(...)
156	>>> for instance in queryset:
157	... instance.name = ...
158	... instance.age = ...
159	... changed.append(instance)
160	>>> # instead of:
161	>>> queryset.bulk_update(changed, ['name', 'age'])
162	>>> # update with:
163	>>> fast_update(queryset, changed, ['name', 'age'])
164
165	`qs` is used to derive the model class and db connection settings. It is not used to filter or
166	narrow the records to be updated.
167
168	`objs` is a sequence of altered model instances. The objects must be model instances the queryset
169	would return. This is not explicitly tested, giving other model types is undefined behavior.
170
171	`fieldnames` denoted the fields, that should be updated.
172
173	`batch_size` controls how many records a single SQL statement will update. For tiny data loads per record
174	(less and/or smaller fields), increase the batch size. For more or bigger fields, lower the batchsize
175	or increase the database temp buffer before doing the big updates. Defaults to 1000 records.
176	"""
177	model: Type[Model] = qs.model	4✔
178
179	# filter all non model local fields --> still handled by bulk_update
180	non_local_fieldnames = []	4✔
181	local_fieldnames = []	4✔
182	for f in fieldnames:	4✔
183	if model._meta.get_field(f) not in model._meta.local_fields:	4✔
184	non_local_fieldnames.append(f)	4✔
185	else:
186	local_fieldnames.append(f)	4✔
187
188	# avoid more expensive doubled updates
189	if non_local_fieldnames and len(local_fieldnames) < 2:	4✔
190	return model.objects.bulk_update(objs, fieldnames, batch_size)	4✔
191
192	if local_fieldnames:	4!
193	from django.db import connections	4✔
194
195	tablename = model._meta.db_table	4✔
196	pk_field = model._meta.pk	4✔
197	if not pk_field:	4!
UNCOV 198	return model.objects.bulk_update(objs, fieldnames, batch_size)	×
199	fields = [model._meta.get_field(f) for f in local_fieldnames]	4✔
200	compiler = qs.query.get_compiler(qs.db)	4✔
201	connection = connections[qs.db]	4✔
202
203	# construct update data
204	data = []	4✔
205	counter = 0	4✔
206	batches = []	4✔
207	for o in objs:	4✔
208	counter += 1	4✔
209	# pk as first value to "join" on
210	sub = [pk_field.get_db_prep_save(getattr(o, pk_field.attname), connection)]	4✔
211	for field in fields:	4✔
212	sub.append(field.get_db_prep_save(getattr(o, field.attname), connection))	4✔
213	data += sub	4✔
214	if counter >= batch_size:	4✔
215	batches.append((counter, data))	4✔
216	data = []	4✔
217	counter = 0	4✔
218	if data:	4✔
219	batches.append((counter, data))	4✔
220
221	sql = ''	4✔
222	last_counter = -1	4✔
223	vendor = connection.vendor	4✔
224	if vendor == 'mysql':	4✔
225	vendor = _adjust_mysql(connection)	2✔
226	for counter, data in batches:	4✔
227	# construct update string
228	if counter != last_counter:	4!
229	sql = QUERY.get(vendor, as_dummy)(	4✔
230	tablename, pk_field.column, fields, counter, compiler, connection)
231	if not sql:	4!
232	# exist with bulk_update for non supported db backends
UNCOV 233	return model.objects.bulk_update(objs, fieldnames, batch_size)	×
234
235	if vendor == 'mysql':	4✔
236	# mysql needs data patch with (0,1,2,...) as first VALUES entry
237	data = list(range(len(fields) + 1)) + data	1✔
238
239	with connection.cursor() as cur:	4✔
240	cur.execute(sql, data)	4✔
241
242	if non_local_fieldnames:	4!
UNCOV 243	model.objects.bulk_update(objs, non_local_fieldnames, batch_size)	×
244
245
246	def check_support(using: str = 'default') -> bool:	15✔
247	"""
248	Check support for `fast_update` with your current database backend.
249	To test it on-the-fly, run these in `./manage.py shell`:
250
251	>>> from computedfields.fast_update import check_support
252	>>> check_support()
253	True
254
255	`using` denotes the database connection to be tested.
256	"""
257	from django.db import connections	15✔
258	connection = connections[using]	15✔
259	if connection.vendor == 'postgresql':	15✔
260	return True	1✔
261	elif connection.vendor == 'sqlite':	14✔
262	if not connection.connection:	12!
263	with connection.cursor():	×
UNCOV 264	pass	×
265	# grab the module to also work with pysqlite3
266	import importlib	12✔
267	_mod = importlib.import_module(connection.connection.__class__.__module__)	12✔
268	major, minor, _ = _mod.sqlite_version_info	12✔
269	if major >= 3 and minor > 32:	12✔
270	return True	1✔
271	elif connection.vendor == 'mysql':	2!
272	if _adjust_mysql(connection):	2!
273	return True	2✔
274	return False	11✔

netzkolchose / django-computedfields / 1774297745

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous