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