codekingpro/portable-devtools
114k
1# engine/cursor.py
2# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: allow-untyped-defs, allow-untyped-calls
8
9"""Define cursor-specific result set constructs including
10:class:`.CursorResult`."""
11
12
13from __future__ import annotations
14
15import collections
16import functools
17import operator
18import typing
19from typing import Any
20from typing import cast
21from typing import ClassVar
22from typing import Dict
23from typing import Iterator
24from typing import List
25from typing import Mapping
26from typing import NoReturn
27from typing import Optional
28from typing import Sequence
29from typing import Tuple
30from typing import TYPE_CHECKING
31from typing import TypeVar
32from typing import Union
33
34from .result import IteratorResult
35from .result import MergedResult
36from .result import Result
37from .result import ResultMetaData
38from .result import SimpleResultMetaData
39from .result import tuplegetter
40from .row import Row
41from .. import exc
42from .. import util
43from ..sql import elements
44from ..sql import sqltypes
45from ..sql import util as sql_util
46from ..sql.base import _generative
47from ..sql.compiler import ResultColumnsEntry
48from ..sql.compiler import RM_NAME
49from ..sql.compiler import RM_OBJECTS
50from ..sql.compiler import RM_RENDERED_NAME
51from ..sql.compiler import RM_TYPE
52from ..sql.type_api import TypeEngine
53from ..util import compat
54from ..util.typing import Literal
55from ..util.typing import Self
56
57
58if typing.TYPE_CHECKING:
59 from .base import Connection
60 from .default import DefaultExecutionContext
61 from .interfaces import _DBAPICursorDescription
62 from .interfaces import DBAPICursor
63 from .interfaces import Dialect
64 from .interfaces import ExecutionContext
65 from .result import _KeyIndexType
66 from .result import _KeyMapRecType
67 from .result import _KeyMapType
68 from .result import _KeyType
69 from .result import _ProcessorsType
70 from .result import _TupleGetterType
71 from ..sql.type_api import _ResultProcessorType
72
73
74_T = TypeVar("_T", bound=Any)
75
76
77# metadata entry tuple indexes.
78# using raw tuple is faster than namedtuple.
79# these match up to the positions in
80# _CursorKeyMapRecType
81MD_INDEX: Literal[0] = 0
82"""integer index in cursor.description
83
84"""
85
86MD_RESULT_MAP_INDEX: Literal[1] = 1
87"""integer index in compiled._result_columns"""
88
89MD_OBJECTS: Literal[2] = 2
90"""other string keys and ColumnElement obj that can match.
91
92This comes from compiler.RM_OBJECTS / compiler.ResultColumnsEntry.objects
93
94"""
95
96MD_LOOKUP_KEY: Literal[3] = 3
97"""string key we usually expect for key-based lookup
98
99this comes from compiler.RM_NAME / compiler.ResultColumnsEntry.name
100"""
101
102
103MD_RENDERED_NAME: Literal[4] = 4
104"""name that is usually in cursor.description
105
106this comes from compiler.RENDERED_NAME / compiler.ResultColumnsEntry.keyname
107"""
108
109
110MD_PROCESSOR: Literal[5] = 5
111"""callable to process a result value into a row"""
112
113MD_UNTRANSLATED: Literal[6] = 6
114"""raw name from cursor.description"""
115
116
117_CursorKeyMapRecType = Tuple[
118 Optional[int], # MD_INDEX, None means the record is ambiguously named
119 int, # MD_RESULT_MAP_INDEX
120 List[Any], # MD_OBJECTS
121 str, # MD_LOOKUP_KEY
122 str, # MD_RENDERED_NAME
123 Optional["_ResultProcessorType[Any]"], # MD_PROCESSOR
124 Optional[str], # MD_UNTRANSLATED
125]
126
127_CursorKeyMapType = Mapping["_KeyType", _CursorKeyMapRecType]
128
129# same as _CursorKeyMapRecType except the MD_INDEX value is definitely
130# not None
131_NonAmbigCursorKeyMapRecType = Tuple[
132 int,
133 int,
134 List[Any],
135 str,
136 str,
137 Optional["_ResultProcessorType[Any]"],
138 str,
139]
140
141
142class CursorResultMetaData(ResultMetaData):
143 """Result metadata for DBAPI cursors."""
144
145 __slots__ = (
146 "_keymap",
147 "_processors",
148 "_keys",
149 "_keymap_by_result_column_idx",
150 "_tuplefilter",
151 "_translated_indexes",
152 "_safe_for_cache",
153 "_unpickled",
154 "_key_to_index",
155 # don't need _unique_filters support here for now. Can be added
156 # if a need arises.
157 )
158
159 _keymap: _CursorKeyMapType
160 _processors: _ProcessorsType
161 _keymap_by_result_column_idx: Optional[Dict[int, _KeyMapRecType]]
162 _unpickled: bool
163 _safe_for_cache: bool
164 _translated_indexes: Optional[List[int]]
165
166 returns_rows: ClassVar[bool] = True
167
168 def _has_key(self, key: Any) -> bool:
169 return key in self._keymap
170
171 def _for_freeze(self) -> ResultMetaData:
172 return SimpleResultMetaData(
173 self._keys,
174 extra=[self._keymap[key][MD_OBJECTS] for key in self._keys],
175 )
176
177 def _make_new_metadata(
178 self,
179 *,
180 unpickled: bool,
181 processors: _ProcessorsType,
182 keys: Sequence[str],
183 keymap: _KeyMapType,
184 tuplefilter: Optional[_TupleGetterType],
185 translated_indexes: Optional[List[int]],
186 safe_for_cache: bool,
187 keymap_by_result_column_idx: Any,
188 ) -> CursorResultMetaData:
189 new_obj = self.__class__.__new__(self.__class__)
190 new_obj._unpickled = unpickled
191 new_obj._processors = processors
192 new_obj._keys = keys
193 new_obj._keymap = keymap
194 new_obj._tuplefilter = tuplefilter
195 new_obj._translated_indexes = translated_indexes
196 new_obj._safe_for_cache = safe_for_cache
197 new_obj._keymap_by_result_column_idx = keymap_by_result_column_idx
198 new_obj._key_to_index = self._make_key_to_index(keymap, MD_INDEX)
199 return new_obj
200
201 def _remove_processors(self) -> CursorResultMetaData:
202 assert not self._tuplefilter
203 return self._make_new_metadata(
204 unpickled=self._unpickled,
205 processors=[None] * len(self._processors),
206 tuplefilter=None,
207 translated_indexes=None,
208 keymap={
209 key: value[0:5] + (None,) + value[6:]
210 for key, value in self._keymap.items()
211 },
212 keys=self._keys,
213 safe_for_cache=self._safe_for_cache,
214 keymap_by_result_column_idx=self._keymap_by_result_column_idx,
215 )
216
217 def _splice_horizontally(
218 self, other: CursorResultMetaData
219 ) -> CursorResultMetaData:
220 assert not self._tuplefilter
221
222 keymap = dict(self._keymap)
223 offset = len(self._keys)
224 keymap.update(
225 {
226 key: (
227 # int index should be None for ambiguous key
228 (
229 value[0] + offset
230 if value[0] is not None and key not in keymap
231 else None
232 ),
233 value[1] + offset,
234 *value[2:],
235 )
236 for key, value in other._keymap.items()
237 }
238 )
239 return self._make_new_metadata(
240 unpickled=self._unpickled,
241 processors=self._processors + other._processors, # type: ignore
242 tuplefilter=None,
243 translated_indexes=None,
244 keys=self._keys + other._keys, # type: ignore
245 keymap=keymap,
246 safe_for_cache=self._safe_for_cache,
247 keymap_by_result_column_idx={
248 metadata_entry[MD_RESULT_MAP_INDEX]: metadata_entry
249 for metadata_entry in keymap.values()
250 },
251 )
252
253 def _reduce(self, keys: Sequence[_KeyIndexType]) -> ResultMetaData:
254 recs = list(self._metadata_for_keys(keys))
255
256 indexes = [rec[MD_INDEX] for rec in recs]
257 new_keys: List[str] = [rec[MD_LOOKUP_KEY] for rec in recs]
258
259 if self._translated_indexes:
260 indexes = [self._translated_indexes[idx] for idx in indexes]
261 tup = tuplegetter(*indexes)
262 new_recs = [(index,) + rec[1:] for index, rec in enumerate(recs)]
263
264 keymap = {rec[MD_LOOKUP_KEY]: rec for rec in new_recs}
265 # TODO: need unit test for:
266 # result = connection.execute("raw sql, no columns").scalars()
267 # without the "or ()" it's failing because MD_OBJECTS is None
268 keymap.update(
269 (e, new_rec)
270 for new_rec in new_recs
271 for e in new_rec[MD_OBJECTS] or ()
272 )
273
274 return self._make_new_metadata(
275 unpickled=self._unpickled,
276 processors=self._processors,
277 keys=new_keys,
278 tuplefilter=tup,
279 translated_indexes=indexes,
280 keymap=keymap, # type: ignore[arg-type]
281 safe_for_cache=self._safe_for_cache,
282 keymap_by_result_column_idx=self._keymap_by_result_column_idx,
283 )
284
285 def _adapt_to_context(self, context: ExecutionContext) -> ResultMetaData:
286 """When using a cached Compiled construct that has a _result_map,
287 for a new statement that used the cached Compiled, we need to ensure
288 the keymap has the Column objects from our new statement as keys.
289 So here we rewrite keymap with new entries for the new columns
290 as matched to those of the cached statement.
291
292 """
293
294 if not context.compiled or not context.compiled._result_columns:
295 return self
296
297 compiled_statement = context.compiled.statement
298 invoked_statement = context.invoked_statement
299
300 if TYPE_CHECKING:
301 assert isinstance(invoked_statement, elements.ClauseElement)
302
303 if compiled_statement is invoked_statement:
304 return self
305
306 assert invoked_statement is not None
307
308 # this is the most common path for Core statements when
309 # caching is used. In ORM use, this codepath is not really used
310 # as the _result_disable_adapt_to_context execution option is
311 # set by the ORM.
312
313 # make a copy and add the columns from the invoked statement
314 # to the result map.
315
316 keymap_by_position = self._keymap_by_result_column_idx
317
318 if keymap_by_position is None:
319 # first retrival from cache, this map will not be set up yet,
320 # initialize lazily
321 keymap_by_position = self._keymap_by_result_column_idx = {
322 metadata_entry[MD_RESULT_MAP_INDEX]: metadata_entry
323 for metadata_entry in self._keymap.values()
324 }
325
326 assert not self._tuplefilter
327 return self._make_new_metadata(
328 keymap=compat.dict_union(
329 self._keymap,
330 {
331 new: keymap_by_position[idx]
332 for idx, new in enumerate(
333 invoked_statement._all_selected_columns
334 )
335 if idx in keymap_by_position
336 },
337 ),
338 unpickled=self._unpickled,
339 processors=self._processors,
340 tuplefilter=None,
341 translated_indexes=None,
342 keys=self._keys,
343 safe_for_cache=self._safe_for_cache,
344 keymap_by_result_column_idx=self._keymap_by_result_column_idx,
345 )
346
347 def __init__(
348 self,
349 parent: CursorResult[Any],
350 cursor_description: _DBAPICursorDescription,
351 ):
352 context = parent.context
353 self._tuplefilter = None
354 self._translated_indexes = None
355 self._safe_for_cache = self._unpickled = False
356
357 if context.result_column_struct:
358 (
359 result_columns,
360 cols_are_ordered,
361 textual_ordered,
362 ad_hoc_textual,
363 loose_column_name_matching,
364 ) = context.result_column_struct
365 num_ctx_cols = len(result_columns)
366 else:
367 result_columns = cols_are_ordered = ( # type: ignore
368 num_ctx_cols
369 ) = ad_hoc_textual = loose_column_name_matching = (
370 textual_ordered
371 ) = False
372
373 # merge cursor.description with the column info
374 # present in the compiled structure, if any
375 raw = self._merge_cursor_description(
376 context,
377 cursor_description,
378 result_columns,
379 num_ctx_cols,
380 cols_are_ordered,
381 textual_ordered,
382 ad_hoc_textual,
383 loose_column_name_matching,
384 )
385
386 # processors in key order which are used when building up
387 # a row
388 self._processors = [
389 metadata_entry[MD_PROCESSOR] for metadata_entry in raw
390 ]
391
392 # this is used when using this ResultMetaData in a Core-only cache
393 # retrieval context. it's initialized on first cache retrieval
394 # when the _result_disable_adapt_to_context execution option
395 # (which the ORM generally sets) is not set.
396 self._keymap_by_result_column_idx = None
397
398 # for compiled SQL constructs, copy additional lookup keys into
399 # the key lookup map, such as Column objects, labels,
400 # column keys and other names
401 if num_ctx_cols:
402 # keymap by primary string...
403 by_key = {
404 metadata_entry[MD_LOOKUP_KEY]: metadata_entry
405 for metadata_entry in raw
406 }
407
408 if len(by_key) != num_ctx_cols:
409 # if by-primary-string dictionary smaller than
410 # number of columns, assume we have dupes; (this check
411 # is also in place if string dictionary is bigger, as
412 # can occur when '*' was used as one of the compiled columns,
413 # which may or may not be suggestive of dupes), rewrite
414 # dupe records with "None" for index which results in
415 # ambiguous column exception when accessed.
416 #
417 # this is considered to be the less common case as it is not
418 # common to have dupe column keys in a SELECT statement.
419 #
420 # new in 1.4: get the complete set of all possible keys,
421 # strings, objects, whatever, that are dupes across two
422 # different records, first.
423 index_by_key: Dict[Any, Any] = {}
424 dupes = set()
425 for metadata_entry in raw:
426 for key in (metadata_entry[MD_RENDERED_NAME],) + (
427 metadata_entry[MD_OBJECTS] or ()
428 ):
429 idx = metadata_entry[MD_INDEX]
430 # if this key has been associated with more than one
431 # positional index, it's a dupe
432 if index_by_key.setdefault(key, idx) != idx:
433 dupes.add(key)
434
435 # then put everything we have into the keymap excluding only
436 # those keys that are dupes.
437 self._keymap = {
438 obj_elem: metadata_entry
439 for metadata_entry in raw
440 if metadata_entry[MD_OBJECTS]
441 for obj_elem in metadata_entry[MD_OBJECTS]
442 if obj_elem not in dupes
443 }
444
445 # then for the dupe keys, put the "ambiguous column"
446 # record into by_key.
447 by_key.update(
448 {
449 key: (None, None, [], key, key, None, None)
450 for key in dupes
451 }
452 )
453
454 else:
455 # no dupes - copy secondary elements from compiled
456 # columns into self._keymap. this is the most common
457 # codepath for Core / ORM statement executions before the
458 # result metadata is cached
459 self._keymap = {
460 obj_elem: metadata_entry
461 for metadata_entry in raw
462 if metadata_entry[MD_OBJECTS]
463 for obj_elem in metadata_entry[MD_OBJECTS]
464 }
465 # update keymap with primary string names taking
466 # precedence
467 self._keymap.update(by_key)
468 else:
469 # no compiled objects to map, just create keymap by primary string
470 self._keymap = {
471 metadata_entry[MD_LOOKUP_KEY]: metadata_entry
472 for metadata_entry in raw
473 }
474
475 # update keymap with "translated" names. In SQLAlchemy this is a
476 # sqlite only thing, and in fact impacting only extremely old SQLite
477 # versions unlikely to be present in modern Python versions.
478 # however, the pyhive third party dialect is
479 # also using this hook, which means others still might use it as well.
480 # I dislike having this awkward hook here but as long as we need
481 # to use names in cursor.description in some cases we need to have
482 # some hook to accomplish this.
483 if not num_ctx_cols and context._translate_colname:
484 self._keymap.update(
485 {
486 metadata_entry[MD_UNTRANSLATED]: self._keymap[
487 metadata_entry[MD_LOOKUP_KEY]
488 ]
489 for metadata_entry in raw
490 if metadata_entry[MD_UNTRANSLATED]
491 }
492 )
493
494 self._key_to_index = self._make_key_to_index(self._keymap, MD_INDEX)
495
496 def _merge_cursor_description(
497 self,
498 context,
499 cursor_description,
500 result_columns,
501 num_ctx_cols,
502 cols_are_ordered,
503 textual_ordered,
504 ad_hoc_textual,
505 loose_column_name_matching,
506 ):
507 """Merge a cursor.description with compiled result column information.
508
509 There are at least four separate strategies used here, selected
510 depending on the type of SQL construct used to start with.
511
512 The most common case is that of the compiled SQL expression construct,
513 which generated the column names present in the raw SQL string and
514 which has the identical number of columns as were reported by
515 cursor.description. In this case, we assume a 1-1 positional mapping
516 between the entries in cursor.description and the compiled object.
517 This is also the most performant case as we disregard extracting /
518 decoding the column names present in cursor.description since we
519 already have the desired name we generated in the compiled SQL
520 construct.
521
522 The next common case is that of the completely raw string SQL,
523 such as passed to connection.execute(). In this case we have no
524 compiled construct to work with, so we extract and decode the
525 names from cursor.description and index those as the primary
526 result row target keys.
527
528 The remaining fairly common case is that of the textual SQL
529 that includes at least partial column information; this is when
530 we use a :class:`_expression.TextualSelect` construct.
531 This construct may have
532 unordered or ordered column information. In the ordered case, we
533 merge the cursor.description and the compiled construct's information
534 positionally, and warn if there are additional description names
535 present, however we still decode the names in cursor.description
536 as we don't have a guarantee that the names in the columns match
537 on these. In the unordered case, we match names in cursor.description
538 to that of the compiled construct based on name matching.
539 In both of these cases, the cursor.description names and the column
540 expression objects and names are indexed as result row target keys.
541
542 The final case is much less common, where we have a compiled
543 non-textual SQL expression construct, but the number of columns
544 in cursor.description doesn't match what's in the compiled
545 construct. We make the guess here that there might be textual
546 column expressions in the compiled construct that themselves include
547 a comma in them causing them to split. We do the same name-matching
548 as with textual non-ordered columns.
549
550 The name-matched system of merging is the same as that used by
551 SQLAlchemy for all cases up through the 0.9 series. Positional
552 matching for compiled SQL expressions was introduced in 1.0 as a
553 major performance feature, and positional matching for textual
554 :class:`_expression.TextualSelect` objects in 1.1.
555 As name matching is no longer
556 a common case, it was acceptable to factor it into smaller generator-
557 oriented methods that are easier to understand, but incur slightly
558 more performance overhead.
559
560 """
561
562 if (
563 num_ctx_cols
564 and cols_are_ordered
565 and not textual_ordered
566 and num_ctx_cols == len(cursor_description)
567 ):
568 self._keys = [elem[0] for elem in result_columns]
569 # pure positional 1-1 case; doesn't need to read
570 # the names from cursor.description
571
572 # most common case for Core and ORM
573
574 # this metadata is safe to cache because we are guaranteed
575 # to have the columns in the same order for new executions
576 self._safe_for_cache = True
577 return [
578 (
579 idx,
580 idx,
581 rmap_entry[RM_OBJECTS],
582 rmap_entry[RM_NAME],
583 rmap_entry[RM_RENDERED_NAME],
584 context.get_result_processor(
585 rmap_entry[RM_TYPE],
586 rmap_entry[RM_RENDERED_NAME],
587 cursor_description[idx][1],
588 ),
589 None,
590 )
591 for idx, rmap_entry in enumerate(result_columns)
592 ]
593 else:
594 # name-based or text-positional cases, where we need
595 # to read cursor.description names
596
597 if textual_ordered or (
598 ad_hoc_textual and len(cursor_description) == num_ctx_cols
599 ):
600 self._safe_for_cache = True
601 # textual positional case
602 raw_iterator = self._merge_textual_cols_by_position(
603 context, cursor_description, result_columns
604 )
605 elif num_ctx_cols:
606 # compiled SQL with a mismatch of description cols
607 # vs. compiled cols, or textual w/ unordered columns
608 # the order of columns can change if the query is
609 # against a "select *", so not safe to cache
610 self._safe_for_cache = False
611 raw_iterator = self._merge_cols_by_name(
612 context,
613 cursor_description,
614 result_columns,
615 loose_column_name_matching,
616 )
617 else:
618 # no compiled SQL, just a raw string, order of columns
619 # can change for "select *"
620 self._safe_for_cache = False
621 raw_iterator = self._merge_cols_by_none(
622 context, cursor_description
623 )
624
625 return [
626 (
627 idx,
628 ridx,
629 obj,
630 cursor_colname,
631 cursor_colname,
632 context.get_result_processor(
633 mapped_type, cursor_colname, coltype
634 ),
635 untranslated,
636 )
637 for (
638 idx,
639 ridx,
640 cursor_colname,
641 mapped_type,
642 coltype,
643 obj,
644 untranslated,
645 ) in raw_iterator
646 ]
647
648 def _colnames_from_description(self, context, cursor_description):
649 """Extract column names and data types from a cursor.description.
650
651 Applies unicode decoding, column translation, "normalization",
652 and case sensitivity rules to the names based on the dialect.
653
654 """
655
656 dialect = context.dialect
657 translate_colname = context._translate_colname
658 normalize_name = (
659 dialect.normalize_name if dialect.requires_name_normalize else None
660 )
661 untranslated = None
662
663 self._keys = []
664
665 for idx, rec in enumerate(cursor_description):
666 colname = rec[0]
667 coltype = rec[1]
668
669 if translate_colname:
670 colname, untranslated = translate_colname(colname)
671
672 if normalize_name:
673 colname = normalize_name(colname)
674
675 self._keys.append(colname)
676
677 yield idx, colname, untranslated, coltype
678
679 def _merge_textual_cols_by_position(
680 self, context, cursor_description, result_columns
681 ):
682 num_ctx_cols = len(result_columns)
683
684 if num_ctx_cols > len(cursor_description):
685 util.warn(
686 "Number of columns in textual SQL (%d) is "
687 "smaller than number of columns requested (%d)"
688 % (num_ctx_cols, len(cursor_description))
689 )
690 seen = set()
691
692 for (
693 idx,
694 colname,
695 untranslated,
696 coltype,
697 ) in self._colnames_from_description(context, cursor_description):
698 if idx < num_ctx_cols:
699 ctx_rec = result_columns[idx]
700 obj = ctx_rec[RM_OBJECTS]
701 ridx = idx
702 mapped_type = ctx_rec[RM_TYPE]
703 if obj[0] in seen:
704 raise exc.InvalidRequestError(
705 "Duplicate column expression requested "
706 "in textual SQL: %r" % obj[0]
707 )
708 seen.add(obj[0])
709 else:
710 mapped_type = sqltypes.NULLTYPE
711 obj = None
712 ridx = None
713 yield idx, ridx, colname, mapped_type, coltype, obj, untranslated
714
715 def _merge_cols_by_name(
716 self,
717 context,
718 cursor_description,
719 result_columns,
720 loose_column_name_matching,
721 ):
722 match_map = self._create_description_match_map(
723 result_columns, loose_column_name_matching
724 )
725 mapped_type: TypeEngine[Any]
726
727 for (
728 idx,
729 colname,
730 untranslated,
731 coltype,
732 ) in self._colnames_from_description(context, cursor_description):
733 try:
734 ctx_rec = match_map[colname]
735 except KeyError:
736 mapped_type = sqltypes.NULLTYPE
737 obj = None
738 result_columns_idx = None
739 else:
740 obj = ctx_rec[1]
741 mapped_type = ctx_rec[2]
742 result_columns_idx = ctx_rec[3]
743 yield (
744 idx,
745 result_columns_idx,
746 colname,
747 mapped_type,
748 coltype,
749 obj,
750 untranslated,
751 )
752
753 @classmethod
754 def _create_description_match_map(
755 cls,
756 result_columns: List[ResultColumnsEntry],
757 loose_column_name_matching: bool = False,
758 ) -> Dict[
759 Union[str, object], Tuple[str, Tuple[Any, ...], TypeEngine[Any], int]
760 ]:
761 """when matching cursor.description to a set of names that are present
762 in a Compiled object, as is the case with TextualSelect, get all the
763 names we expect might match those in cursor.description.
764 """
765
766 d: Dict[
767 Union[str, object],
768 Tuple[str, Tuple[Any, ...], TypeEngine[Any], int],
769 ] = {}
770 for ridx, elem in enumerate(result_columns):
771 key = elem[RM_RENDERED_NAME]
772 if key in d:
773 # conflicting keyname - just add the column-linked objects
774 # to the existing record. if there is a duplicate column
775 # name in the cursor description, this will allow all of those
776 # objects to raise an ambiguous column error
777 e_name, e_obj, e_type, e_ridx = d[key]
778 d[key] = e_name, e_obj + elem[RM_OBJECTS], e_type, ridx
779 else:
780 d[key] = (elem[RM_NAME], elem[RM_OBJECTS], elem[RM_TYPE], ridx)
781
782 if loose_column_name_matching:
783 # when using a textual statement with an unordered set
784 # of columns that line up, we are expecting the user
785 # to be using label names in the SQL that match to the column
786 # expressions. Enable more liberal matching for this case;
787 # duplicate keys that are ambiguous will be fixed later.
788 for r_key in elem[RM_OBJECTS]:
789 d.setdefault(
790 r_key,
791 (elem[RM_NAME], elem[RM_OBJECTS], elem[RM_TYPE], ridx),
792 )
793 return d
794
795 def _merge_cols_by_none(self, context, cursor_description):
796 for (
797 idx,
798 colname,
799 untranslated,
800 coltype,
801 ) in self._colnames_from_description(context, cursor_description):
802 yield (
803 idx,
804 None,
805 colname,
806 sqltypes.NULLTYPE,
807 coltype,
808 None,
809 untranslated,
810 )
811
812 if not TYPE_CHECKING:
813
814 def _key_fallback(
815 self, key: Any, err: Optional[Exception], raiseerr: bool = True
816 ) -> Optional[NoReturn]:
817 if raiseerr:
818 if self._unpickled and isinstance(key, elements.ColumnElement):
819 raise exc.NoSuchColumnError(
820 "Row was unpickled; lookup by ColumnElement "
821 "is unsupported"
822 ) from err
823 else:
824 raise exc.NoSuchColumnError(
825 "Could not locate column in row for column '%s'"
826 % util.string_or_unprintable(key)
827 ) from err
828 else:
829 return None
830
831 def _raise_for_ambiguous_column_name(self, rec):
832 raise exc.InvalidRequestError(
833 "Ambiguous column name '%s' in "
834 "result set column descriptions" % rec[MD_LOOKUP_KEY]
835 )
836
837 def _index_for_key(self, key: Any, raiseerr: bool = True) -> Optional[int]:
838 # TODO: can consider pre-loading ints and negative ints
839 # into _keymap - also no coverage here
840 if isinstance(key, int):
841 key = self._keys[key]
842
843 try:
844 rec = self._keymap[key]
845 except KeyError as ke:
846 x = self._key_fallback(key, ke, raiseerr)
847 assert x is None
848 return None
849
850 index = rec[0]
851
852 if index is None:
853 self._raise_for_ambiguous_column_name(rec)
854 return index
855
856 def _indexes_for_keys(self, keys):
857 try:
858 return [self._keymap[key][0] for key in keys]
859 except KeyError as ke:
860 # ensure it raises
861 CursorResultMetaData._key_fallback(self, ke.args[0], ke)
862
863 def _metadata_for_keys(
864 self, keys: Sequence[Any]
865 ) -> Iterator[_NonAmbigCursorKeyMapRecType]:
866 for key in keys:
867 if int in key.__class__.__mro__:
868 key = self._keys[key]
869
870 try:
871 rec = self._keymap[key]
872 except KeyError as ke:
873 # ensure it raises
874 CursorResultMetaData._key_fallback(self, ke.args[0], ke)
875
876 index = rec[MD_INDEX]
877
878 if index is None:
879 self._raise_for_ambiguous_column_name(rec)
880
881 yield cast(_NonAmbigCursorKeyMapRecType, rec)
882
883 def __getstate__(self):
884 # TODO: consider serializing this as SimpleResultMetaData
885 return {
886 "_keymap": {
887 key: (
888 rec[MD_INDEX],
889 rec[MD_RESULT_MAP_INDEX],
890 [],
891 key,
892 rec[MD_RENDERED_NAME],
893 None,
894 None,
895 )
896 for key, rec in self._keymap.items()
897 if isinstance(key, (str, int))
898 },
899 "_keys": self._keys,
900 "_translated_indexes": self._translated_indexes,
901 }
902
903 def __setstate__(self, state):
904 self._processors = [None for _ in range(len(state["_keys"]))]
905 self._keymap = state["_keymap"]
906 self._keymap_by_result_column_idx = None
907 self._key_to_index = self._make_key_to_index(self._keymap, MD_INDEX)
908 self._keys = state["_keys"]
909 self._unpickled = True
910 if state["_translated_indexes"]:
911 self._translated_indexes = cast(
912 "List[int]", state["_translated_indexes"]
913 )
914 self._tuplefilter = tuplegetter(*self._translated_indexes)
915 else:
916 self._translated_indexes = self._tuplefilter = None
917
918
919class ResultFetchStrategy:
920 """Define a fetching strategy for a result object.
921
922
923 .. versionadded:: 1.4
924
925 """
926
927 __slots__ = ()
928
929 alternate_cursor_description: Optional[_DBAPICursorDescription] = None
930
931 def soft_close(
932 self, result: CursorResult[Any], dbapi_cursor: Optional[DBAPICursor]
933 ) -> None:
934 raise NotImplementedError()
935
936 def hard_close(
937 self, result: CursorResult[Any], dbapi_cursor: Optional[DBAPICursor]
938 ) -> None:
939 raise NotImplementedError()
940
941 def yield_per(
942 self,
943 result: CursorResult[Any],
944 dbapi_cursor: Optional[DBAPICursor],
945 num: int,
946 ) -> None:
947 return
948
949 def fetchone(
950 self,
951 result: CursorResult[Any],
952 dbapi_cursor: DBAPICursor,
953 hard_close: bool = False,
954 ) -> Any:
955 raise NotImplementedError()
956
957 def fetchmany(
958 self,
959 result: CursorResult[Any],
960 dbapi_cursor: DBAPICursor,
961 size: Optional[int] = None,
962 ) -> Any:
963 raise NotImplementedError()
964
965 def fetchall(
966 self,
967 result: CursorResult[Any],
968 dbapi_cursor: DBAPICursor,
969 ) -> Any:
970 raise NotImplementedError()
971
972 def handle_exception(
973 self,
974 result: CursorResult[Any],
975 dbapi_cursor: Optional[DBAPICursor],
976 err: BaseException,
977 ) -> NoReturn:
978 raise err
979
980
981class NoCursorFetchStrategy(ResultFetchStrategy):
982 """Cursor strategy for a result that has no open cursor.
983
984 There are two varieties of this strategy, one for DQL and one for
985 DML (and also DDL), each of which represent a result that had a cursor
986 but no longer has one.
987
988 """
989
990 __slots__ = ()
991
992 def soft_close(self, result, dbapi_cursor):
993 pass
994
995 def hard_close(self, result, dbapi_cursor):
996 pass
997
998 def fetchone(self, result, dbapi_cursor, hard_close=False):
999 return self._non_result(result, None)
1000
1001 def fetchmany(self, result, dbapi_cursor, size=None):
1002 return self._non_result(result, [])
1003
1004 def fetchall(self, result, dbapi_cursor):
1005 return self._non_result(result, [])
1006
1007 def _non_result(self, result, default, err=None):
1008 raise NotImplementedError()
1009
1010
1011class NoCursorDQLFetchStrategy(NoCursorFetchStrategy):
1012 """Cursor strategy for a DQL result that has no open cursor.
1013
1014 This is a result set that can return rows, i.e. for a SELECT, or for an
1015 INSERT, UPDATE, DELETE that includes RETURNING. However it is in the state
1016 where the cursor is closed and no rows remain available. The owning result
1017 object may or may not be "hard closed", which determines if the fetch
1018 methods send empty results or raise for closed result.
1019
1020 """
1021
1022 __slots__ = ()
1023
1024 def _non_result(self, result, default, err=None):
1025 if result.closed:
1026 raise exc.ResourceClosedError(
1027 "This result object is closed."
1028 ) from err
1029 else:
1030 return default
1031
1032
1033_NO_CURSOR_DQL = NoCursorDQLFetchStrategy()
1034
1035
1036class NoCursorDMLFetchStrategy(NoCursorFetchStrategy):
1037 """Cursor strategy for a DML result that has no open cursor.
1038
1039 This is a result set that does not return rows, i.e. for an INSERT,
1040 UPDATE, DELETE that does not include RETURNING.
1041
1042 """
1043
1044 __slots__ = ()
1045
1046 def _non_result(self, result, default, err=None):
1047 # we only expect to have a _NoResultMetaData() here right now.
1048 assert not result._metadata.returns_rows
1049 result._metadata._we_dont_return_rows(err)
1050
1051
1052_NO_CURSOR_DML = NoCursorDMLFetchStrategy()
1053
1054
1055class CursorFetchStrategy(ResultFetchStrategy):
1056 """Call fetch methods from a DBAPI cursor.
1057
1058 Alternate versions of this class may instead buffer the rows from
1059 cursors or not use cursors at all.
1060
1061 """
1062
1063 __slots__ = ()
1064
1065 def soft_close(
1066 self, result: CursorResult[Any], dbapi_cursor: Optional[DBAPICursor]
1067 ) -> None:
1068 result.cursor_strategy = _NO_CURSOR_DQL
1069
1070 def hard_close(
1071 self, result: CursorResult[Any], dbapi_cursor: Optional[DBAPICursor]
1072 ) -> None:
1073 result.cursor_strategy = _NO_CURSOR_DQL
1074
1075 def handle_exception(
1076 self,
1077 result: CursorResult[Any],
1078 dbapi_cursor: Optional[DBAPICursor],
1079 err: BaseException,
1080 ) -> NoReturn:
1081 result.connection._handle_dbapi_exception(
1082 err, None, None, dbapi_cursor, result.context
1083 )
1084
1085 def yield_per(
1086 self,
1087 result: CursorResult[Any],
1088 dbapi_cursor: Optional[DBAPICursor],
1089 num: int,
1090 ) -> None:
1091 result.cursor_strategy = BufferedRowCursorFetchStrategy(
1092 dbapi_cursor,
1093 {"max_row_buffer": num},
1094 initial_buffer=collections.deque(),
1095 growth_factor=0,
1096 )
1097
1098 def fetchone(
1099 self,
1100 result: CursorResult[Any],
1101 dbapi_cursor: DBAPICursor,
1102 hard_close: bool = False,
1103 ) -> Any:
1104 try:
1105 row = dbapi_cursor.fetchone()
1106 if row is None:
1107 result._soft_close(hard=hard_close)
1108 return row
1109 except BaseException as e:
1110 self.handle_exception(result, dbapi_cursor, e)
1111
1112 def fetchmany(
1113 self,
1114 result: CursorResult[Any],
1115 dbapi_cursor: DBAPICursor,
1116 size: Optional[int] = None,
1117 ) -> Any:
1118 try:
1119 if size is None:
1120 l = dbapi_cursor.fetchmany()
1121 else:
1122 l = dbapi_cursor.fetchmany(size)
1123
1124 if not l:
1125 result._soft_close()
1126 return l
1127 except BaseException as e:
1128 self.handle_exception(result, dbapi_cursor, e)
1129
1130 def fetchall(
1131 self,
1132 result: CursorResult[Any],
1133 dbapi_cursor: DBAPICursor,
1134 ) -> Any:
1135 try:
1136 rows = dbapi_cursor.fetchall()
1137 result._soft_close()
1138 return rows
1139 except BaseException as e:
1140 self.handle_exception(result, dbapi_cursor, e)
1141
1142
1143_DEFAULT_FETCH = CursorFetchStrategy()
1144
1145
1146class BufferedRowCursorFetchStrategy(CursorFetchStrategy):
1147 """A cursor fetch strategy with row buffering behavior.
1148
1149 This strategy buffers the contents of a selection of rows
1150 before ``fetchone()`` is called. This is to allow the results of
1151 ``cursor.description`` to be available immediately, when
1152 interfacing with a DB-API that requires rows to be consumed before
1153 this information is available (currently psycopg2, when used with
1154 server-side cursors).
1155
1156 The pre-fetching behavior fetches only one row initially, and then
1157 grows its buffer size by a fixed amount with each successive need
1158 for additional rows up the ``max_row_buffer`` size, which defaults
1159 to 1000::
1160
1161 with psycopg2_engine.connect() as conn:
1162
1163 result = conn.execution_options(
1164 stream_results=True, max_row_buffer=50
1165 ).execute(text("select * from table"))
1166
1167 .. versionadded:: 1.4 ``max_row_buffer`` may now exceed 1000 rows.
1168
1169 .. seealso::
1170
1171 :ref:`psycopg2_execution_options`
1172 """
1173
1174 __slots__ = ("_max_row_buffer", "_rowbuffer", "_bufsize", "_growth_factor")
1175
1176 def __init__(
1177 self,
1178 dbapi_cursor,
1179 execution_options,
1180 growth_factor=5,
1181 initial_buffer=None,
1182 ):
1183 self._max_row_buffer = execution_options.get("max_row_buffer", 1000)
1184
1185 if initial_buffer is not None:
1186 self._rowbuffer = initial_buffer
1187 else:
1188 self._rowbuffer = collections.deque(dbapi_cursor.fetchmany(1))
1189 self._growth_factor = growth_factor
1190
1191 if growth_factor:
1192 self._bufsize = min(self._max_row_buffer, self._growth_factor)
1193 else:
1194 self._bufsize = self._max_row_buffer
1195
1196 @classmethod
1197 def create(cls, result):
1198 return BufferedRowCursorFetchStrategy(
1199 result.cursor,
1200 result.context.execution_options,