|
| 1 | +# Copyright 2020 Google LLC |
| 2 | +# |
| 3 | +# Use of this source code is governed by a BSD-style |
| 4 | +# license that can be found in the LICENSE file or at |
| 5 | +# https://developers.google.com/open-source/licenses/bsd |
| 6 | + |
| 7 | +from google.cloud.spanner_dbapi.parse_utils import get_param_types |
| 8 | +from google.cloud.spanner_dbapi.parse_utils import parse_insert |
| 9 | +from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner |
| 10 | +from google.cloud.spanner_v1 import param_types |
| 11 | + |
| 12 | + |
| 13 | +SQL_LIST_TABLES = """ |
| 14 | + SELECT |
| 15 | + t.table_name |
| 16 | + FROM |
| 17 | + information_schema.tables AS t |
| 18 | + WHERE |
| 19 | + t.table_catalog = '' and t.table_schema = '' |
| 20 | + """ |
| 21 | + |
| 22 | +SQL_GET_TABLE_COLUMN_SCHEMA = """SELECT |
| 23 | + COLUMN_NAME, IS_NULLABLE, SPANNER_TYPE |
| 24 | + FROM |
| 25 | + INFORMATION_SCHEMA.COLUMNS |
| 26 | + WHERE |
| 27 | + TABLE_SCHEMA = '' |
| 28 | + AND |
| 29 | + TABLE_NAME = @table_name |
| 30 | + """ |
| 31 | + |
| 32 | +# This table maps spanner_types to Spanner's data type sizes as per |
| 33 | +# https://cloud.google.com/spanner/docs/data-types#allowable-types |
| 34 | +# It is used to map `display_size` to a known type for Cursor.description |
| 35 | +# after a row fetch. |
| 36 | +# Since ResultMetadata |
| 37 | +# https://cloud.google.com/spanner/docs/reference/rest/v1/ResultSetMetadata |
| 38 | +# does not send back the actual size, we have to lookup the respective size. |
| 39 | +# Some fields' sizes are dependent upon the dynamic data hence aren't sent back |
| 40 | +# by Cloud Spanner. |
| 41 | +code_to_display_size = { |
| 42 | + param_types.BOOL.code: 1, |
| 43 | + param_types.DATE.code: 4, |
| 44 | + param_types.FLOAT64.code: 8, |
| 45 | + param_types.INT64.code: 8, |
| 46 | + param_types.TIMESTAMP.code: 12, |
| 47 | +} |
| 48 | + |
| 49 | + |
| 50 | +def _execute_insert_heterogenous(transaction, sql_params_list): |
| 51 | + for sql, params in sql_params_list: |
| 52 | + sql, params = sql_pyformat_args_to_spanner(sql, params) |
| 53 | + param_types = get_param_types(params) |
| 54 | + res = transaction.execute_sql( |
| 55 | + sql, params=params, param_types=param_types |
| 56 | + ) |
| 57 | + # TODO: File a bug with Cloud Spanner and the Python client maintainers |
| 58 | + # about a lost commit when res isn't read from. |
| 59 | + _ = list(res) |
| 60 | + |
| 61 | + |
| 62 | +def _execute_insert_homogenous(transaction, parts): |
| 63 | + # Perform an insert in one shot. |
| 64 | + table = parts.get("table") |
| 65 | + columns = parts.get("columns") |
| 66 | + values = parts.get("values") |
| 67 | + return transaction.insert(table, columns, values) |
| 68 | + |
| 69 | + |
| 70 | +def handle_insert(connection, sql, params): |
| 71 | + parts = parse_insert(sql, params) |
| 72 | + |
| 73 | + # The split between the two styles exists because: |
| 74 | + # in the common case of multiple values being passed |
| 75 | + # with simple pyformat arguments, |
| 76 | + # SQL: INSERT INTO T (f1, f2) VALUES (%s, %s, %s) |
| 77 | + # Params: [(1, 2, 3, 4, 5, 6, 7, 8, 9, 10,)] |
| 78 | + # we can take advantage of a single RPC with: |
| 79 | + # transaction.insert(table, columns, values) |
| 80 | + # instead of invoking: |
| 81 | + # with transaction: |
| 82 | + # for sql, params in sql_params_list: |
| 83 | + # transaction.execute_sql(sql, params, param_types) |
| 84 | + # which invokes more RPCs and is more costly. |
| 85 | + |
| 86 | + if parts.get("homogenous"): |
| 87 | + # The common case of multiple values being passed in |
| 88 | + # non-complex pyformat args and need to be uploaded in one RPC. |
| 89 | + return connection.database.run_in_transaction( |
| 90 | + _execute_insert_homogenous, parts |
| 91 | + ) |
| 92 | + else: |
| 93 | + # All the other cases that are esoteric and need |
| 94 | + # transaction.execute_sql |
| 95 | + sql_params_list = parts.get("sql_params_list") |
| 96 | + return connection.database.run_in_transaction( |
| 97 | + _execute_insert_heterogenous, sql_params_list |
| 98 | + ) |
| 99 | + |
| 100 | + |
| 101 | +class ColumnInfo: |
| 102 | + """Row column description object.""" |
| 103 | + |
| 104 | + def __init__( |
| 105 | + self, |
| 106 | + name, |
| 107 | + type_code, |
| 108 | + display_size=None, |
| 109 | + internal_size=None, |
| 110 | + precision=None, |
| 111 | + scale=None, |
| 112 | + null_ok=False, |
| 113 | + ): |
| 114 | + self.name = name |
| 115 | + self.type_code = type_code |
| 116 | + self.display_size = display_size |
| 117 | + self.internal_size = internal_size |
| 118 | + self.precision = precision |
| 119 | + self.scale = scale |
| 120 | + self.null_ok = null_ok |
| 121 | + |
| 122 | + self.fields = ( |
| 123 | + self.name, |
| 124 | + self.type_code, |
| 125 | + self.display_size, |
| 126 | + self.internal_size, |
| 127 | + self.precision, |
| 128 | + self.scale, |
| 129 | + self.null_ok, |
| 130 | + ) |
| 131 | + |
| 132 | + def __repr__(self): |
| 133 | + return self.__str__() |
| 134 | + |
| 135 | + def __getitem__(self, index): |
| 136 | + return self.fields[index] |
| 137 | + |
| 138 | + def __str__(self): |
| 139 | + str_repr = ", ".join( |
| 140 | + filter( |
| 141 | + lambda part: part is not None, |
| 142 | + [ |
| 143 | + "name='%s'" % self.name, |
| 144 | + "type_code=%d" % self.type_code, |
| 145 | + "display_size=%d" % self.display_size |
| 146 | + if self.display_size |
| 147 | + else None, |
| 148 | + "internal_size=%d" % self.internal_size |
| 149 | + if self.internal_size |
| 150 | + else None, |
| 151 | + "precision='%s'" % self.precision |
| 152 | + if self.precision |
| 153 | + else None, |
| 154 | + "scale='%s'" % self.scale if self.scale else None, |
| 155 | + "null_ok='%s'" % self.null_ok if self.null_ok else None, |
| 156 | + ], |
| 157 | + ) |
| 158 | + ) |
| 159 | + return "ColumnInfo(%s)" % str_repr |
0 commit comments