This repository was archived by the owner on Mar 31, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 104
Expand file tree
/
Copy pathcursor.py
More file actions
634 lines (556 loc) · 23.3 KB
/
cursor.py
File metadata and controls
634 lines (556 loc) · 23.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
# Copyright 2020 Google LLC All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Database cursor for Google Cloud Spanner DB API."""
from collections import namedtuple
from google.api_core.exceptions import (
Aborted,
AlreadyExists,
FailedPrecondition,
InternalServerError,
InvalidArgument,
OutOfRange,
)
import sqlparse
from google.cloud import spanner_v1 as spanner
from google.cloud.spanner_dbapi import (
_helpers,
batch_dml_executor,
client_side_statement_executor,
parse_utils,
)
from google.cloud.spanner_dbapi._helpers import CODE_TO_DISPLAY_SIZE, ColumnInfo
from google.cloud.spanner_dbapi.batch_dml_executor import BatchMode
from google.cloud.spanner_dbapi.exceptions import (
IntegrityError,
InterfaceError,
OperationalError,
ProgrammingError,
)
from google.cloud.spanner_dbapi.parse_utils import get_param_types
from google.cloud.spanner_dbapi.parsed_statement import (
AutocommitDmlMode,
ParsedStatement,
Statement,
StatementType,
)
from google.cloud.spanner_dbapi.transaction_helper import CursorStatementType
from google.cloud.spanner_dbapi.utils import PeekIterator, StreamedManyResultSets
from google.cloud.spanner_v1 import RequestOptions
from google.cloud.spanner_v1.merged_result_set import MergedResultSet
ColumnDetails = namedtuple("column_details", ["null_ok", "spanner_type"])
def check_not_closed(function):
"""`Cursor` class methods decorator.
Raise an exception if the cursor is closed, or not bound to a
connection, or the parent connection is closed.
:raises: :class:`InterfaceError` if this cursor is closed.
:raises: :class:`ProgrammingError` if this cursor is not bound to a connection.
"""
def wrapper(cursor, *args, **kwargs):
if not cursor.connection:
raise ProgrammingError("Cursor is not connected to the database")
if cursor.is_closed:
raise InterfaceError("Cursor and/or connection is already closed.")
return function(cursor, *args, **kwargs)
return wrapper
class Cursor(object):
"""Database cursor to manage the context of a fetch operation.
:type connection: :class:`~google.cloud.spanner_dbapi.connection.Connection`
:param connection: A DB-API connection to Google Cloud Spanner.
"""
def __init__(self, connection):
self._itr = None
self._result_set = None
self._row_count = None
self.lastrowid = None
self.connection = connection
self.transaction_helper = self.connection._transaction_helper
self._is_closed = False
# the number of rows to fetch at a time with fetchmany()
self.arraysize = 1
self._parsed_statement: ParsedStatement = None
self._in_retry_mode = False
self._batch_dml_rows_count = None
self._request_tag = None
@property
def request_tag(self):
"""The request tag that will be applied to the next statement on this
cursor. This property is automatically cleared when a statement is
executed.
Returns:
str: The request tag that will be applied to the next statement on
this cursor.
"""
return self._request_tag
@request_tag.setter
def request_tag(self, value):
"""Sets the request tag for the next statement on this cursor. This
property is automatically cleared when a statement is executed.
Args:
value (str): The request tag for the statement.
"""
self._request_tag = value
@property
def request_options(self):
options = self.connection.request_options
if self._request_tag:
if not options:
options = RequestOptions()
options.request_tag = self._request_tag
self._request_tag = None
return options
@property
def is_closed(self):
"""The cursor close indicator.
:rtype: bool
:returns: True if the cursor or the parent connection is closed,
otherwise False.
"""
return self._is_closed or self.connection.is_closed
@property
def description(self):
"""
Read-only attribute containing the result columns description
of a form:
- ``name``
- ``type_code``
- ``display_size``
- ``internal_size``
- ``precision``
- ``scale``
- ``null_ok``
:rtype: tuple
:returns: The result columns' description.
"""
if (
self._result_set is None
or not getattr(self._result_set, "metadata", None)
or self._result_set.metadata.row_type is None
or self._result_set.metadata.row_type.fields is None
or len(self._result_set.metadata.row_type.fields) == 0
):
return
columns = []
for field in self._result_set.metadata.row_type.fields:
columns.append(
ColumnInfo(
name=field.name,
type_code=field.type_.code,
# Size of the SQL type of the column.
display_size=CODE_TO_DISPLAY_SIZE.get(field.type_.code),
# Client perceived size of the column.
internal_size=field._pb.ByteSize(),
)
)
return tuple(columns)
@property
def rowcount(self):
"""The number of rows updated by the last INSERT, UPDATE, DELETE request's `execute()` call.
For SELECT requests the rowcount returns -1.
:rtype: int
:returns: The number of rows updated by the last INSERT, UPDATE, DELETE request's .execute*() call.
"""
if self._row_count is not None or self._result_set is None:
return self._row_count
stats = getattr(self._result_set, "stats", None)
if stats is not None and "row_count_exact" in stats:
return stats.row_count_exact
return -1
@check_not_closed
def callproc(self, procname, args=None):
"""A no-op, raising an error if the cursor or connection is closed."""
pass
@check_not_closed
def nextset(self):
"""A no-op, raising an error if the cursor or connection is closed."""
pass
@check_not_closed
def setinputsizes(self, sizes):
"""A no-op, raising an error if the cursor or connection is closed."""
pass
@check_not_closed
def setoutputsize(self, size, column=None):
"""A no-op, raising an error if the cursor or connection is closed."""
pass
def close(self):
"""Closes this cursor."""
self._is_closed = True
def _do_execute_update_in_autocommit(self, transaction, sql, params):
"""This function should only be used in autocommit mode."""
self.connection._transaction = transaction
self.connection._snapshot = None
self._result_set = transaction.execute_sql(
sql,
params=params,
param_types=get_param_types(params),
last_statement=True,
)
self._itr = PeekIterator(self._result_set)
self._row_count = None
def _batch_DDLs(self, sql):
"""
Check that the given operation contains only DDL
statements and batch them into an internal list.
:type sql: str
:param sql: A SQL query statement.
:raises: :class:`ValueError` in case not a DDL statement
present in the operation.
"""
statements = []
for ddl in sqlparse.split(sql):
if ddl:
ddl = ddl.rstrip(";")
if (
parse_utils.classify_statement(ddl).statement_type
!= StatementType.DDL
):
raise ValueError("Only DDL statements may be batched.")
statements.append(ddl)
# Only queue DDL statements if they are all correctly classified.
self.connection._ddl_statements.extend(statements)
def _reset(self):
if self.connection.database is None:
raise ValueError("Database needs to be passed for this operation")
self._itr = None
self._result_set = None
self._row_count = None
self._batch_dml_rows_count = None
@check_not_closed
def execute(self, sql, args=None):
self._execute(sql, args, False)
def _execute(self, sql, args=None, call_from_execute_many=False):
"""Prepares and executes a Spanner database operation.
:type sql: str
:param sql: A SQL query statement.
:type args: list
:param args: Additional parameters to supplement the SQL query.
"""
self._reset()
exception = None
try:
self._parsed_statement = parse_utils.classify_statement(sql, args)
if self._parsed_statement is None:
raise ProgrammingError("Invalid Statement.")
if self._parsed_statement.statement_type == StatementType.CLIENT_SIDE:
self._result_set = client_side_statement_executor.execute(
self, self._parsed_statement
)
if self._result_set is not None:
if isinstance(
self._result_set, StreamedManyResultSets
) or isinstance(self._result_set, MergedResultSet):
self._itr = self._result_set
else:
self._itr = PeekIterator(self._result_set)
elif self.connection._batch_mode == BatchMode.DML:
self.connection.execute_batch_dml_statement(self._parsed_statement)
elif self.connection.read_only or (
not self.connection._client_transaction_started
and self._parsed_statement.statement_type == StatementType.QUERY
):
self._handle_DQL(sql, args or None)
elif self._parsed_statement.statement_type == StatementType.DDL:
self._batch_DDLs(sql)
if not self.connection._client_transaction_started:
self.connection.run_prior_DDL_statements()
elif (
self.connection.autocommit_dml_mode
is AutocommitDmlMode.PARTITIONED_NON_ATOMIC
):
self._row_count = self.connection.database.execute_partitioned_dml(
sql,
params=args,
param_types=self._parsed_statement.statement.param_types,
request_options=self.request_options,
)
self._result_set = None
else:
self._execute_in_rw_transaction()
except (AlreadyExists, FailedPrecondition, OutOfRange) as e:
exception = IntegrityError(getattr(e, "details", e))
exception.__cause__ = e
raise exception
except InvalidArgument as e:
exception = ProgrammingError(getattr(e, "details", e))
exception.__cause__ = e
raise exception
except InternalServerError as e:
exception = OperationalError(getattr(e, "details", e))
exception.__cause__ = e
raise exception
except Exception as e:
exception = e
raise
finally:
if not self._in_retry_mode and not call_from_execute_many:
self.transaction_helper.add_execute_statement_for_retry(
self, sql, args, exception, False
)
if self.connection._client_transaction_started is False:
self.connection._spanner_transaction_started = False
def _execute_in_rw_transaction(self):
# For every other operation, we've got to ensure that
# any prior DDL statements were run.
self.connection.run_prior_DDL_statements()
statement = self._parsed_statement.statement
if self.connection._client_transaction_started:
while True:
try:
self._result_set = self.connection.run_statement(
statement, self.request_options
)
self._itr = PeekIterator(self._result_set)
return
except Aborted:
# We are raising it so it could be handled in transaction_helper.py and is retried
if self._in_retry_mode:
raise
else:
self.transaction_helper.retry_transaction()
except Exception as ex:
# In case of inline-begin failure, the transaction isn't started.
# We immediately retry with an explicit BeginTransaction.
transaction = getattr(self.connection, "_transaction", None)
if transaction and not transaction._transaction_id:
transaction._reset_and_begin()
# Let the existing retry loop handle the retry of the statement
continue
raise ex
else:
self.connection.database.run_in_transaction(
self._do_execute_update_in_autocommit,
statement.sql,
statement.params or None,
)
@check_not_closed
def executemany(self, operation, seq_of_params):
"""Execute the given SQL with every parameters set
from the given sequence of parameters.
:type operation: str
:param operation: SQL code to execute.
:type seq_of_params: list
:param seq_of_params: Sequence of additional parameters to run
the query with.
"""
self._reset()
exception = None
try:
self._parsed_statement = parse_utils.classify_statement(operation)
if self._parsed_statement.statement_type == StatementType.DDL:
raise ProgrammingError(
"Executing DDL statements with executemany() method is not allowed."
)
if self._parsed_statement.statement_type == StatementType.CLIENT_SIDE:
raise ProgrammingError(
"Executing the following operation: "
+ operation
+ ", with executemany() method is not allowed."
)
# For every operation, we've got to ensure that any prior DDL
# statements were run.
self.connection.run_prior_DDL_statements()
# Treat UNKNOWN statements as if they are DML and let the server
# determine what is wrong with it.
if self._parsed_statement.statement_type in (
StatementType.INSERT,
StatementType.UPDATE,
StatementType.UNKNOWN,
):
statements = []
for params in seq_of_params:
sql, params = parse_utils.sql_pyformat_args_to_spanner(
operation, params
)
statements.append(Statement(sql, params, get_param_types(params)))
many_result_set = batch_dml_executor.run_batch_dml(self, statements)
else:
many_result_set = StreamedManyResultSets()
for params in seq_of_params:
self._execute(operation, params, True)
many_result_set.add_iter(self._itr)
self._result_set = many_result_set
self._itr = many_result_set
except (AlreadyExists, FailedPrecondition, OutOfRange) as e:
exception = IntegrityError(getattr(e, "details", e))
exception.__cause__ = e
raise exception
except InvalidArgument as e:
exception = ProgrammingError(getattr(e, "details", e))
exception.__cause__ = e
raise exception
except InternalServerError as e:
exception = OperationalError(getattr(e, "details", e))
exception.__cause__ = e
raise exception
except Exception as e:
exception = e
raise
finally:
if not self._in_retry_mode:
self.transaction_helper.add_execute_statement_for_retry(
self,
operation,
seq_of_params,
exception,
True,
)
if self.connection._client_transaction_started is False:
self.connection._spanner_transaction_started = False
@check_not_closed
def fetchone(self):
"""Fetch the next row of a query result set, returning a single
sequence, or None when no more data is available."""
rows = self._fetch(CursorStatementType.FETCH_ONE)
if not rows:
return
return rows[0]
@check_not_closed
def fetchall(self):
"""Fetch all (remaining) rows of a query result, returning them as
a sequence of sequences.
"""
return self._fetch(CursorStatementType.FETCH_ALL)
@check_not_closed
def fetchmany(self, size=None):
"""Fetch the next set of rows of a query result, returning a sequence
of sequences. An empty sequence is returned when no more rows are available.
:type size: int
:param size: (Optional) The maximum number of results to fetch.
:raises InterfaceError:
if the previous call to .execute*() did not produce any result set
or if no call was issued yet.
"""
if size is None:
size = self.arraysize
return self._fetch(CursorStatementType.FETCH_MANY, size)
def _fetch(self, cursor_statement_type, size=None):
exception = None
rows = []
is_fetch_all = False
try:
while True:
rows = []
try:
if cursor_statement_type == CursorStatementType.FETCH_ALL:
is_fetch_all = True
for row in self:
rows.append(row)
elif cursor_statement_type == CursorStatementType.FETCH_MANY:
for _ in range(size):
try:
row = next(self)
rows.append(row)
except StopIteration:
break
elif cursor_statement_type == CursorStatementType.FETCH_ONE:
try:
row = next(self)
rows.append(row)
except StopIteration:
return
break
except Aborted:
if not self.connection.read_only:
if self._in_retry_mode:
raise
else:
self.transaction_helper.retry_transaction()
except Exception as e:
exception = e
finally:
if not self._in_retry_mode:
self.transaction_helper.add_fetch_statement_for_retry(
self, rows, exception, is_fetch_all
)
return rows
def _handle_DQL_with_snapshot(self, snapshot, sql, params):
self._result_set = snapshot.execute_sql(
sql,
params,
get_param_types(params),
request_options=self.request_options,
)
# Read the first element so that the StreamedResultSet can
# return the metadata after a DQL statement.
self._itr = PeekIterator(self._result_set)
# Unfortunately, Spanner doesn't seem to send back
# information about the number of rows available.
self._row_count = None
if self._result_set.metadata.transaction.read_timestamp is not None:
snapshot._transaction_read_timestamp = (
self._result_set.metadata.transaction.read_timestamp
)
def _handle_DQL(self, sql, params):
if self.connection.database is None:
raise ValueError("Database needs to be passed for this operation")
sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params)
if self.connection.read_only and self.connection._client_transaction_started:
# initiate or use the existing multi-use snapshot
self._handle_DQL_with_snapshot(
self.connection.snapshot_checkout(), sql, params
)
else:
# execute with single-use snapshot
with self.connection.database.snapshot(
**self.connection.staleness
) as snapshot:
self.connection._snapshot = snapshot
self.connection._transaction = None
self._handle_DQL_with_snapshot(snapshot, sql, params)
def __enter__(self):
return self
def __exit__(self, etype, value, traceback):
self.close()
def __next__(self):
if self._itr is None:
raise ProgrammingError("no results to return")
return next(self._itr)
def __iter__(self):
if self._itr is None:
raise ProgrammingError("no results to return")
return self._itr
def list_tables(self, schema_name="", include_views=True):
"""List the tables of the linked Database.
:rtype: list
:returns: The list of tables within the Database.
"""
return self.run_sql_in_snapshot(
sql=_helpers.SQL_LIST_TABLES_AND_VIEWS
if include_views
else _helpers.SQL_LIST_TABLES,
params={"table_schema": schema_name},
param_types={"table_schema": spanner.param_types.STRING},
)
def run_sql_in_snapshot(self, sql, params=None, param_types=None):
# Some SQL e.g. for INFORMATION_SCHEMA cannot be run in read-write transactions
# hence this method exists to circumvent that limit.
if self.connection.database is None:
raise ValueError("Database needs to be passed for this operation")
self.connection.run_prior_DDL_statements()
with self.connection.database.snapshot() as snapshot:
return list(snapshot.execute_sql(sql, params, param_types))
def get_table_column_schema(self, table_name, schema_name=""):
rows = self.run_sql_in_snapshot(
sql=_helpers.SQL_GET_TABLE_COLUMN_SCHEMA,
params={"schema_name": schema_name, "table_name": table_name},
param_types={
"schema_name": spanner.param_types.STRING,
"table_name": spanner.param_types.STRING,
},
)
column_details = {}
for column_name, is_nullable, spanner_type in rows:
column_details[column_name] = ColumnDetails(
null_ok=is_nullable == "YES", spanner_type=spanner_type
)
return column_details