def _merge_query_options(base, merge):
"""Merge higher precedence QueryOptions with current QueryOptions.
:type base:
:class:`~google.cloud.spanner_v1.ExecuteSqlRequest.QueryOptions`
or :class:`dict` or None
:param base: The current QueryOptions that is intended for use.
:type merge:
:class:`~google.cloud.spanner_v1.ExecuteSqlRequest.QueryOptions`
or :class:`dict` or None
:param merge:
The QueryOptions that have a higher priority than base. These options
should overwrite the fields in base.
:rtype:
:class:`~google.cloud.spanner_v1.ExecuteSqlRequest.QueryOptions`
or None
:returns:
QueryOptions object formed by merging the two given QueryOptions.
If the resultant object only has empty fields, returns None.
"""
combined = base or ExecuteSqlRequest.QueryOptions()
if type(combined) == dict:
combined = ExecuteSqlRequest.QueryOptions(
optimizer_version=combined.get("optimizer_version", ""))
merge = merge or ExecuteSqlRequest.QueryOptions()
if type(merge) == dict:
merge = ExecuteSqlRequest.QueryOptions(
optimizer_version=merge.get("optimizer_version", ""))
type(combined).pb(combined).MergeFrom(type(merge).pb(merge))
if not combined.optimizer_version:
return None
return combined
def test_base_empty_and_merge_empty(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
base = ExecuteSqlRequest.QueryOptions()
merge = ExecuteSqlRequest.QueryOptions()
result = self._callFUT(base, merge)
self.assertIsNone(result)
def test_base_object_merge_dict(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
base = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
merge = {"optimizer_version": "3"}
expected = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
result = self._callFUT(base, merge)
self.assertEqual(result, expected)
def test_constructor_custom_query_options_client_config(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
from google.cloud.spanner_v1 import client as MUT
expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
creds = _make_credentials()
self._constructor_test_helper(
expected_scopes,
creds,
query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"),
expected_query_options=ExecuteSqlRequest.QueryOptions(
optimizer_version="1"
),
)
def test_base_object_merge_dict(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
base = ExecuteSqlRequest.QueryOptions(
optimizer_version="1",
optimizer_statistics_package="auto_20191128_14_47_22UTC",
)
merge = {"optimizer_version": "3"}
expected = ExecuteSqlRequest.QueryOptions(
optimizer_version="3",
optimizer_statistics_package="auto_20191128_14_47_22UTC",
)
result = self._callFUT(base, merge)
self.assertEqual(result, expected)
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 test_base_none_merge_object(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
base = None
merge = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
result = self._callFUT(base, merge)
self.assertEqual(result, merge)
def test_execute_sql_w_query_options(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
self._execute_sql_helper(
multi_use=False,
query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="3"),
)
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,
)
restart = functools.partial(
api.execute_streaming_sql, request=request, 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
def test_base_dict_and_merge_none(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
base = {"optimizer_version": "2"}
merge = None
expected = ExecuteSqlRequest.QueryOptions(optimizer_version="2")
result = self._callFUT(base, merge)
self.assertEqual(result, expected)
def ping(self):
"""Ping the session to keep it alive by executing "SELECT 1".
:raises: ValueError: if :attr:`session_id` is not already set.
"""
if self._session_id is None:
raise ValueError("Session ID not set by back-end")
api = self._database.spanner_api
metadata = _metadata_with_prefix(self._database.name)
request = ExecuteSqlRequest(session=self.name, sql="SELECT 1")
api.execute_sql(request=request, metadata=metadata)
def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats):
from google.cloud.spanner_v1 import ExecuteSqlRequest
from google.cloud.spanner_v1 import client as MUT
expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
creds = _make_credentials()
mock_ver.return_value = "2"
mock_stats.return_value = "auto_20191128_14_47_22UTC"
query_options = ExecuteSqlRequest.QueryOptions(
optimizer_version="1",
optimizer_statistics_package="auto_20191128_10_47_22UTC",
)
expected_query_options = ExecuteSqlRequest.QueryOptions(
optimizer_version="2",
optimizer_statistics_package="auto_20191128_14_47_22UTC",
)
self._constructor_test_helper(
expected_scopes,
creds,
query_options=query_options,
expected_query_options=expected_query_options,
)
def __init__(
self,
project=None,
credentials=None,
client_info=_CLIENT_INFO,
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(),
optimizer_statistics_package=
_get_spanner_optimizer_statistics_package(),
)
# Environment flag config has higher precedence than application config.
self._query_options = _merge_query_options(query_options,
env_query_options)
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 test_ping_hit(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
gax_api = self._make_spanner_api()
gax_api.execute_sql.return_value = "1"
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
session.ping()
request = ExecuteSqlRequest(
session=self.SESSION_NAME,
sql="SELECT 1",
)
gax_api.execute_sql.assert_called_once_with(
request=request,
metadata=[("google-cloud-resource-prefix", database.name)],
)
def test_ping_error(self):
from google.api_core.exceptions import Unknown
from google.cloud.spanner_v1 import ExecuteSqlRequest
gax_api = self._make_spanner_api()
gax_api.execute_sql.side_effect = Unknown("testing")
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
with self.assertRaises(Unknown):
session.ping()
request = ExecuteSqlRequest(
session=self.SESSION_NAME,
sql="SELECT 1",
)
gax_api.execute_sql.assert_called_once_with(
request=request,
metadata=[("google-cloud-resource-prefix", database.name)],
)
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 test_execute_update_w_query_options(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
self._execute_update_helper(
query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="3")
)
def __init__(self):
from google.cloud.spanner_v1 import ExecuteSqlRequest
self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
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 _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 import (
PartialResultSet,
ResultSetMetadata,
ResultSetStats,
)
from google.cloud.spanner_v1 import (
TransactionSelector,
TransactionOptions,
)
from google.cloud.spanner_v1 import ExecuteSqlRequest
from google.cloud.spanner_v1 import Type, StructType
from google.cloud.spanner_v1 import TypeCode
from google.cloud.spanner_v1._helpers import (
_make_value_pb,
_merge_query_options,
)
VALUES = [[u"bharney", u"rhubbyl", 31], [u"phred", u"phlyntstone", 32]]
MODE = 2 # PROFILE
struct_type_pb = StructType(
fields=[
StructType.Field(name="first_name", type_=Type(code=TypeCode.STRING)),
StructType.Field(name="last_name", type_=Type(code=TypeCode.STRING)),
StructType.Field(name="age", type_=Type(code=TypeCode.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(metadata=metadata_pb),
PartialResultSet(stats=stats_pb),
]
for i in range(len(result_sets)):
result_sets[i].values.extend(VALUES[i])
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
)
expected_request = ExecuteSqlRequest(
session=self.SESSION_NAME,
sql=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,
)
api.execute_streaming_sql.assert_called_once_with(
request=expected_request,
metadata=[("google-cloud-resource-prefix", database.name)],
timeout=timeout,
retry=retry,
)
self.assertEqual(derived._execute_sql_count, sql_count + 1)
self.assertSpanAttributes(
"CloudSpanner.ReadWriteTransaction",
status=StatusCanonicalCode.OK,
attributes=dict(BASE_ATTRIBUTES, **{"db.statement": SQL_QUERY_WITH_PARAM}),
)
