def _execute_update_helper(self, count=0, query_options=None): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ( ResultSet, ResultSetStats, ) from google.cloud.spanner_v1 import TransactionSelector from google.cloud.spanner_v1._helpers import ( _make_value_pb, _merge_query_options, ) from google.cloud.spanner_v1 import ExecuteSqlRequest MODE = 2 # PROFILE stats_pb = ResultSetStats(row_count_exact=1) database = _Database() api = database.spanner_api = self._make_spanner_api() api.execute_sql.return_value = ResultSet(stats=stats_pb) session = _Session(database) transaction = self._make_one(session) transaction._transaction_id = self.TRANSACTION_ID transaction._execute_sql_count = count row_count = transaction.execute_update( DML_QUERY_WITH_PARAM, PARAMS, PARAM_TYPES, query_mode=MODE, query_options=query_options, ) self.assertEqual(row_count, 1) expected_transaction = TransactionSelector(id=self.TRANSACTION_ID) expected_params = Struct( fields={key: _make_value_pb(value) for (key, value) in PARAMS.items()} ) expected_query_options = database._instance._client._query_options if query_options: expected_query_options = _merge_query_options( expected_query_options, query_options ) expected_request = ExecuteSqlRequest( session=self.SESSION_NAME, sql=DML_QUERY_WITH_PARAM, transaction=expected_transaction, params=expected_params, param_types=PARAM_TYPES, query_mode=MODE, query_options=expected_query_options, seqno=count, ) api.execute_sql.assert_called_once_with( request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], ) self.assertEqual(transaction._execute_sql_count, count + 1)
def execute_update( self, dml, params=None, param_types=None, query_mode=None, query_options=None ): """Perform an ``ExecuteSql`` API request with DML. :type dml: str :param dml: SQL DML statement :type params: dict, {str -> column value} :param params: values for parameter replacement. Keys must match the names used in ``dml``. :type param_types: dict[str -> Union[dict, .types.Type]] :param param_types: (Optional) maps explicit types for one or more param values; required if parameters are passed. :type query_mode: :class:`~google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryMode` :param query_mode: Mode governing return of results / query plan. See: `QueryMode <https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode>`_. :type query_options: :class:`~google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions` or :class:`dict` :param query_options: (Optional) Options that are provided for query plan stability. :rtype: int :returns: Count of rows affected by the DML statement. """ params_pb = self._make_params_pb(params, param_types) database = self._session._database metadata = _metadata_with_prefix(database.name) transaction = self._make_txn_selector() api = database.spanner_api seqno, self._execute_sql_count = ( self._execute_sql_count, self._execute_sql_count + 1, ) # Query-level options have higher precedence than client-level and # environment-level options default_query_options = database._instance._client._query_options query_options = _merge_query_options(default_query_options, query_options) response = api.execute_sql( self._session.name, dml, transaction=transaction, params=params_pb, param_types=param_types, query_mode=query_mode, query_options=query_options, seqno=seqno, metadata=metadata, ) return response.stats.row_count_exact
def __init__( self, project=None, credentials=None, client_info=_CLIENT_INFO, user_agent=None, client_options=None, query_options=None, ): self._emulator_host = _get_spanner_emulator_host() if client_options and type(client_options) == dict: self._client_options = google.api_core.client_options.from_dict( client_options ) else: self._client_options = client_options if self._emulator_host: credentials = AnonymousCredentials() elif isinstance(credentials, AnonymousCredentials): self._emulator_host = self._client_options.api_endpoint # NOTE: This API has no use for the _http argument, but sending it # will have no impact since the _http() @property only lazily # creates a working HTTP object. super(Client, self).__init__( project=project, credentials=credentials, client_options=client_options, _http=None, ) self._client_info = client_info env_query_options = ExecuteSqlRequest.QueryOptions( optimizer_version=_get_spanner_optimizer_version() ) # Environment flag config has higher precedence than application config. self._query_options = _merge_query_options(query_options, env_query_options) if user_agent is not None: warnings.warn(_USER_AGENT_DEPRECATED, DeprecationWarning, stacklevel=2) self.user_agent = user_agent if self._emulator_host is not None and ( "http://" in self._emulator_host or "https://" in self._emulator_host ): warnings.warn(_EMULATOR_HOST_HTTP_SCHEME)
def execute_sql( self, sql, params=None, param_types=None, query_mode=None, query_options=None, request_options=None, partition=None, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, ): """Perform an ``ExecuteStreamingSql`` API request. :type sql: str :param sql: SQL query statement :type params: dict, {str -> column value} :param params: values for parameter replacement. Keys must match the names used in ``sql``. :type param_types: dict[str -> Union[dict, .types.Type]] :param param_types: (Optional) maps explicit types for one or more param values; required if parameters are passed. :type query_mode: :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryMode` :param query_mode: Mode governing return of results / query plan. See: `QueryMode <https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode>`_. :type query_options: :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions` or :class:`dict` :param query_options: (Optional) Query optimizer configuration to use for the given query. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.QueryOptions` :type request_options: :class:`google.cloud.spanner_v1.types.RequestOptions` :param request_options: (Optional) Common options for this request. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.RequestOptions`. :type partition: bytes :param partition: (Optional) one of the partition tokens returned from :meth:`partition_query`. :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet` :returns: a result set instance which can be used to consume rows. :type retry: :class:`~google.api_core.retry.Retry` :param retry: (Optional) The retry settings for this request. :type timeout: float :param timeout: (Optional) The timeout for this request. :raises ValueError: for reuse of single-use snapshots, or if a transaction ID is already pending for multiple-use snapshots. """ if self._read_request_count > 0: if not self._multi_use: raise ValueError("Cannot re-use single-use snapshot.") if self._transaction_id is None: raise ValueError("Transaction ID pending.") if params is not None: if param_types is None: raise ValueError( "Specify 'param_types' when passing 'params'.") params_pb = Struct(fields={ key: _make_value_pb(value) for key, value in params.items() }) else: params_pb = {} database = self._session._database metadata = _metadata_with_prefix(database.name) transaction = self._make_txn_selector() api = database.spanner_api # Query-level options have higher precedence than client-level and # environment-level options default_query_options = database._instance._client._query_options query_options = _merge_query_options(default_query_options, query_options) if type(request_options) == dict: request_options = RequestOptions(request_options) request = ExecuteSqlRequest( session=self._session.name, sql=sql, transaction=transaction, params=params_pb, param_types=param_types, query_mode=query_mode, partition_token=partition, seqno=self._execute_sql_count, query_options=query_options, request_options=request_options, ) restart = functools.partial( api.execute_streaming_sql, request=request, metadata=metadata, retry=retry, timeout=timeout, ) trace_attributes = {"db.statement": sql} iterator = _restart_on_unavailable( restart, request, "CloudSpanner.ReadWriteTransaction", self._session, trace_attributes, ) self._read_request_count += 1 self._execute_sql_count += 1 if self._multi_use: return StreamedResultSet(iterator, source=self) else: return StreamedResultSet(iterator)
def _callFUT(self, *args, **kw): from google.cloud.spanner_v1._helpers import _merge_query_options return _merge_query_options(*args, **kw)
def execute_update( self, dml, params=None, param_types=None, query_mode=None, query_options=None, request_options=None, *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT, ): """Perform an ``ExecuteSql`` API request with DML. :type dml: str :param dml: SQL DML statement :type params: dict, {str -> column value} :param params: values for parameter replacement. Keys must match the names used in ``dml``. :type param_types: dict[str -> Union[dict, .types.Type]] :param param_types: (Optional) maps explicit types for one or more param values; required if parameters are passed. :type query_mode: :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryMode` :param query_mode: Mode governing return of results / query plan. See: `QueryMode <https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode>`_. :type query_options: :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions` or :class:`dict` :param query_options: (Optional) Options that are provided for query plan stability. :type request_options: :class:`google.cloud.spanner_v1.types.RequestOptions` :param request_options: (Optional) Common options for this request. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.RequestOptions`. :type retry: :class:`~google.api_core.retry.Retry` :param retry: (Optional) The retry settings for this request. :type timeout: float :param timeout: (Optional) The timeout for this request. :rtype: int :returns: Count of rows affected by the DML statement. """ params_pb = self._make_params_pb(params, param_types) database = self._session._database metadata = _metadata_with_prefix(database.name) transaction = self._make_txn_selector() api = database.spanner_api seqno, self._execute_sql_count = ( self._execute_sql_count, self._execute_sql_count + 1, ) # Query-level options have higher precedence than client-level and # environment-level options default_query_options = database._instance._client._query_options query_options = _merge_query_options(default_query_options, query_options) if type(request_options) == dict: request_options = RequestOptions(request_options) trace_attributes = {"db.statement": dml} request = ExecuteSqlRequest( session=self._session.name, sql=dml, transaction=transaction, params=params_pb, param_types=param_types, query_mode=query_mode, query_options=query_options, seqno=seqno, request_options=request_options, ) with trace_call("CloudSpanner.ReadWriteTransaction", self._session, trace_attributes): response = api.execute_sql(request=request, metadata=metadata, retry=retry, timeout=timeout) return response.stats.row_count_exact
def generate_query_batches( self, sql, params=None, param_types=None, partition_size_bytes=None, max_partitions=None, query_options=None, ): """Start a partitioned query operation. Uses the ``PartitionQuery`` API request to start a partitioned query operation. Returns a list of batch information needed to peform the actual queries. :type sql: str :param sql: SQL query statement :type params: dict, {str -> column value} :param params: values for parameter replacement. Keys must match the names used in ``sql``. :type param_types: dict[str -> Union[dict, .types.Type]] :param param_types: (Optional) maps explicit types for one or more param values; required if parameters are passed. :type partition_size_bytes: int :param partition_size_bytes: (Optional) desired size for each partition generated. The service uses this as a hint, the actual partition size may differ. :type partition_size_bytes: int :param partition_size_bytes: (Optional) desired size for each partition generated. The service uses this as a hint, the actual partition size may differ. :type max_partitions: int :param max_partitions: (Optional) desired maximum number of partitions generated. The service uses this as a hint, the actual number of partitions may differ. :type query_options: :class:`~google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions` or :class:`dict` :param query_options: (Optional) Query optimizer configuration to use for the given query. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.QueryOptions` :rtype: iterable of dict :returns: mappings of information used peform actual partitioned reads via :meth:`process_read_batch`. """ partitions = self._get_snapshot().partition_query( sql=sql, params=params, param_types=param_types, partition_size_bytes=partition_size_bytes, max_partitions=max_partitions, ) query_info = {"sql": sql} if params: query_info["params"] = params query_info["param_types"] = param_types # Query-level options have higher precedence than client-level and # environment-level options default_query_options = self._database._instance._client._query_options query_info["query_options"] = _merge_query_options( default_query_options, query_options) for partition in partitions: yield {"partition": partition, "query": query_info}
def execute_partitioned_dml(self, dml, params=None, param_types=None, query_options=None): """Execute a partitionable DML statement. :type dml: str :param dml: DML statement :type params: dict, {str -> column value} :param params: values for parameter replacement. Keys must match the names used in ``dml``. :type param_types: dict[str -> Union[dict, .types.Type]] :param param_types: (Optional) maps explicit types for one or more param values; required if parameters are passed. :type query_options: :class:`~google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions` or :class:`dict` :param query_options: (Optional) Query optimizer configuration to use for the given query. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.QueryOptions` :rtype: int :returns: Count of rows affected by the DML statement. """ query_options = _merge_query_options( self._instance._client._query_options, query_options) if params is not None: if param_types is None: raise ValueError( "Specify 'param_types' when passing 'params'.") params_pb = Struct(fields={ key: _make_value_pb(value) for key, value in params.items() }) else: params_pb = None api = self.spanner_api txn_options = TransactionOptions( partitioned_dml=TransactionOptions.PartitionedDml()) metadata = _metadata_with_prefix(self.name) def execute_pdml(): with SessionCheckout(self._pool) as session: txn = api.begin_transaction(session.name, txn_options, metadata=metadata) txn_selector = TransactionSelector(id=txn.id) restart = functools.partial( api.execute_streaming_sql, session.name, dml, transaction=txn_selector, params=params_pb, param_types=param_types, query_options=query_options, metadata=metadata, ) iterator = _restart_on_unavailable(restart) result_set = StreamedResultSet(iterator) list(result_set) # consume all partials return result_set.stats.row_count_lower_bound retry_config = api._method_configs["ExecuteStreamingSql"].retry return _retry_on_aborted(execute_pdml, retry_config)()
def _execute_sql_helper( self, multi_use, first=True, count=0, partition=None, sql_count=0, query_options=None, timeout=google.api_core.gapic_v1.method.DEFAULT, retry=google.api_core.gapic_v1.method.DEFAULT, ): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1.proto.result_set_pb2 import ( PartialResultSet, ResultSetMetadata, ResultSetStats, ) from google.cloud.spanner_v1.proto.transaction_pb2 import ( TransactionSelector, TransactionOptions, ) from google.cloud.spanner_v1.proto.type_pb2 import Type, StructType from google.cloud.spanner_v1.proto.type_pb2 import STRING, INT64 from google.cloud.spanner_v1._helpers import ( _make_value_pb, _merge_query_options, ) VALUES = [[u"bharney", u"rhubbyl", 31], [u"phred", u"phlyntstone", 32]] VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES] MODE = 2 # PROFILE struct_type_pb = StructType(fields=[ StructType.Field(name="first_name", type=Type(code=STRING)), StructType.Field(name="last_name", type=Type(code=STRING)), StructType.Field(name="age", type=Type(code=INT64)), ]) metadata_pb = ResultSetMetadata(row_type=struct_type_pb) stats_pb = ResultSetStats(query_stats=Struct( fields={"rows_returned": _make_value_pb(2)})) result_sets = [ PartialResultSet(values=VALUE_PBS[0], metadata=metadata_pb), PartialResultSet(values=VALUE_PBS[1], stats=stats_pb), ] iterator = _MockIterator(*result_sets) database = _Database() api = database.spanner_api = self._make_spanner_api() api.execute_streaming_sql.return_value = iterator session = _Session(database) derived = self._makeDerived(session) derived._multi_use = multi_use derived._read_request_count = count derived._execute_sql_count = sql_count if not first: derived._transaction_id = TXN_ID result_set = derived.execute_sql( SQL_QUERY_WITH_PARAM, PARAMS, PARAM_TYPES, query_mode=MODE, query_options=query_options, partition=partition, retry=retry, timeout=timeout, ) self.assertEqual(derived._read_request_count, count + 1) if multi_use: self.assertIs(result_set._source, derived) else: self.assertIsNone(result_set._source) self.assertEqual(list(result_set), VALUES) self.assertEqual(result_set.metadata, metadata_pb) self.assertEqual(result_set.stats, stats_pb) txn_options = TransactionOptions(read_only=TransactionOptions.ReadOnly( strong=True)) if multi_use: if first: expected_transaction = TransactionSelector(begin=txn_options) else: expected_transaction = TransactionSelector(id=TXN_ID) else: expected_transaction = TransactionSelector(single_use=txn_options) expected_params = Struct(fields={ key: _make_value_pb(value) for (key, value) in PARAMS.items() }) expected_query_options = database._instance._client._query_options if query_options: expected_query_options = _merge_query_options( expected_query_options, query_options) api.execute_streaming_sql.assert_called_once_with( self.SESSION_NAME, SQL_QUERY_WITH_PARAM, transaction=expected_transaction, params=expected_params, param_types=PARAM_TYPES, query_mode=MODE, query_options=expected_query_options, partition_token=partition, seqno=sql_count, metadata=[("google-cloud-resource-prefix", database.name)], timeout=timeout, retry=retry, ) self.assertEqual(derived._execute_sql_count, sql_count + 1)
def execute_partitioned_dml( self, dml, params=None, param_types=None, query_options=None, request_options=None, ): """Execute a partitionable DML statement. :type dml: str :param dml: DML statement :type params: dict, {str -> column value} :param params: values for parameter replacement. Keys must match the names used in ``dml``. :type param_types: dict[str -> Union[dict, .types.Type]] :param param_types: (Optional) maps explicit types for one or more param values; required if parameters are passed. :type query_options: :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions` or :class:`dict` :param query_options: (Optional) Query optimizer configuration to use for the given query. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.QueryOptions` :type request_options: :class:`google.cloud.spanner_v1.types.RequestOptions` :param request_options: (Optional) Common options for this request. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.RequestOptions`. Please note, the `transactionTag` setting will be ignored as it is not supported for partitioned DML. :rtype: int :returns: Count of rows affected by the DML statement. """ query_options = _merge_query_options( self._instance._client._query_options, query_options) if request_options is None: request_options = RequestOptions() elif type(request_options) == dict: request_options = RequestOptions(request_options) request_options.transaction_tag = None if params is not None: from google.cloud.spanner_v1.transaction import Transaction if param_types is None: raise ValueError( "Specify 'param_types' when passing 'params'.") params_pb = Transaction._make_params_pb(params, param_types) else: params_pb = {} api = self.spanner_api txn_options = TransactionOptions( partitioned_dml=TransactionOptions.PartitionedDml()) metadata = _metadata_with_prefix(self.name) def execute_pdml(): with SessionCheckout(self._pool) as session: txn = api.begin_transaction(session=session.name, options=txn_options, metadata=metadata) txn_selector = TransactionSelector(id=txn.id) request = ExecuteSqlRequest( session=session.name, sql=dml, transaction=txn_selector, params=params_pb, param_types=param_types, query_options=query_options, request_options=request_options, ) method = functools.partial( api.execute_streaming_sql, metadata=metadata, ) iterator = _restart_on_unavailable(method, request) result_set = StreamedResultSet(iterator) list(result_set) # consume all partials return result_set.stats.row_count_lower_bound return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()