def _GetQueryRequestBeta(sql,
query_mode,
session_ref=None,
enable_partitioned_dml=False):
"""Formats the request based on whether the statement contains DML.
Args:
sql: String, The SQL to execute.
query_mode: String, The mode in which to run the query. Must be one of
'NORMAL', 'PLAN', or 'PROFILE'
session_ref: Reference to the session.
enable_partitioned_dml: Boolean, whether partitioned dml is enabled.
Returns:
ExecuteSqlRequest parameters
"""
msgs = apis.GetMessagesModule('spanner', 'v1')
if enable_partitioned_dml is True:
transaction = _GetPartitionedDmlTransaction(session_ref)
elif QueryHasDml(sql):
transaction_options = msgs.TransactionOptions(readWrite=msgs.ReadWrite())
transaction = msgs.TransactionSelector(begin=transaction_options)
else:
transaction_options = msgs.TransactionOptions(
readOnly=msgs.ReadOnly(strong=True))
transaction = msgs.TransactionSelector(singleUse=transaction_options)
return msgs.ExecuteSqlRequest(
sql=sql,
queryMode=msgs.ExecuteSqlRequest.QueryModeValueValuesEnum(query_mode),
transaction=transaction)
def _GetQueryRequest(sql, query_mode):
"""Formats the request based on whether the statement contains DML.
Args:
sql: String, The SQL to execute.
query_mode: String, The mode in which to run the query. Must be one of
'NORMAL', 'PLAN', or 'PROFILE'
Returns:
ExecuteSqlRequest parameters
"""
msgs = apis.GetMessagesModule('spanner', 'v1')
if QueryHasDml(sql):
transaction_options = msgs.TransactionOptions(
readWrite=msgs.ReadWrite())
transaction = msgs.TransactionSelector(begin=transaction_options)
else:
transaction_options = msgs.TransactionOptions(readOnly=msgs.ReadOnly(
strong=True))
transaction = msgs.TransactionSelector(singleUse=transaction_options)
return msgs.ExecuteSqlRequest(
sql=sql,
queryMode=msgs.ExecuteSqlRequest.QueryModeValueValuesEnum(query_mode),
transaction=transaction)
def ExecuteSqlBeta(session_ref, sql, query_mode, enable_partitioned_dml=False):
"""Execute an SQL command.
Args:
session_ref: Session, Indicates that the repo should be created if it does
not exist.
sql: String, The SQL to execute.
query_mode: String, The mode in which to run the query. Must be one of
'NORMAL', 'PLAN', or 'PROFILE'
enable_partitioned_dml: Boolean, whether partitioned dml is enabled.
Returns:
(Repo) The capture repository.
"""
client = apis.GetClientInstance('spanner', 'v1')
msgs = apis.GetMessagesModule('spanner', 'v1')
_RegisterCustomMessageCodec(msgs)
execute_sql_request = _GetQueryRequestBeta(sql, query_mode, session_ref,
enable_partitioned_dml)
req = msgs.SpannerProjectsInstancesDatabasesSessionsExecuteSqlRequest(
session=session_ref.RelativeName(), executeSqlRequest=execute_sql_request)
resp = client.projects_instances_databases_sessions.ExecuteSql(req)
if QueryHasDml(sql) and enable_partitioned_dml is False:
result_set = msgs.ResultSet(metadata=resp.metadata)
Commit(session_ref, [], result_set.metadata.transaction.id)
return resp
def ExecuteSql(sql, query_mode, session_ref, read_only_options=None,
enable_partitioned_dml=False, http_timeout_sec=None):
"""Execute an SQL command.
Args:
sql: String, The SQL to execute.
query_mode: String, The mode in which to run the query. Must be one of
'NORMAL', 'PLAN', or 'PROFILE'
session_ref: Session, Indicates that the repo should be created if it does
not exist.
read_only_options: The ReadOnly message for a read-only request. It is
ignored in a DML request.
enable_partitioned_dml: Boolean, whether partitioned dml is enabled.
http_timeout_sec: int, Maximum time in seconds to wait for the SQL query to
complete.
Returns:
(Repo) The capture repository.
"""
client = _GetClientInstance('spanner', 'v1', http_timeout_sec)
msgs = apis.GetMessagesModule('spanner', 'v1')
_RegisterCustomMessageCodec(msgs)
execute_sql_request = _GetQueryRequest(
sql, query_mode, session_ref, read_only_options, enable_partitioned_dml)
req = msgs.SpannerProjectsInstancesDatabasesSessionsExecuteSqlRequest(
session=session_ref.RelativeName(), executeSqlRequest=execute_sql_request)
resp = client.projects_instances_databases_sessions.ExecuteSql(req)
if QueryHasDml(sql) and enable_partitioned_dml is False:
result_set = msgs.ResultSet(metadata=resp.metadata)
Commit(session_ref, [], result_set.metadata.transaction.id)
return resp
def ParseReadOnlyOptions(self, args):
"""Parses the options for a read-only request from command line arguments.
Args:
args: Command line arguments.
Returns:
A ReadOnly message if the query is read-only (not DML), otherwise None.
"""
if QueryHasDml(args.sql):
if args.IsSpecified('strong'):
raise c_exceptions.InvalidArgumentException(
'--strong',
'A timestamp bound cannot be specified for a DML statement.'
)
if args.IsSpecified('read_timestamp'):
raise c_exceptions.InvalidArgumentException(
'--read-timestamp',
'A timestamp bound cannot be specified for a DML statement.'
)
return None
else:
msgs = apis.GetMessagesModule('spanner', 'v1')
if args.IsSpecified('read_timestamp'):
return msgs.ReadOnly(readTimestamp=args.read_timestamp)
elif args.IsSpecified('strong'):
if not args.strong:
raise c_exceptions.InvalidArgumentException(
'--strong', '`--strong` cannot be set to false. '
'Instead specify a different type of timestamp bound.')
else:
return msgs.ReadOnly(strong=True)
else:
# The default timestamp bound is strong.
return msgs.ReadOnly(strong=True)
def ExecuteSql(session_ref, sql, query_mode):
"""Execute an SQL command.
Args:
session_ref: Session, Indicates that the repo should be created if
it does not exist.
sql: String, The SQL to execute.
query_mode: String, The mode in which to run the query. Must be one
of 'NORMAL', 'PLAN', or 'PROFILE'
Returns:
(Repo) The capture repository.
"""
client = apis.GetClientInstance('spanner', 'v1')
msgs = apis.GetMessagesModule('spanner', 'v1')
# TODO(b/33482229): remove this workaround
def _ToJson(msg):
return extra_types.JsonProtoEncoder(
extra_types.JsonArray(entries=msg.entry))
def _FromJson(data):
return msgs.ResultSet.RowsValueListEntry(
entry=extra_types.JsonProtoDecoder(data).entries)
encoding.RegisterCustomMessageCodec(encoder=_ToJson, decoder=_FromJson)(
msgs.ResultSet.RowsValueListEntry)
execute_sql_request = _GetQueryRequest(sql, query_mode)
req = msgs.SpannerProjectsInstancesDatabasesSessionsExecuteSqlRequest(
session=session_ref.RelativeName(),
executeSqlRequest=execute_sql_request)
resp = client.projects_instances_databases_sessions.ExecuteSql(req)
if QueryHasDml(sql):
result_set = msgs.ResultSet(metadata=resp.metadata)
Commit(session_ref, [], result_set.metadata.transaction.id)
return resp
def testQueryHasDmlWithHint(self):
sql = '@ { hint } \nUPDATE abc SET b = 1 WHERE a = 1'
self.assertEqual(QueryHasDml(sql), True)
def testQueryDoesNotHaveDml(self): sql = 'select * from abc' self.assertEqual(QueryHasDml(sql), False)
def testQueryHasDmlWithMultilineComment(self): sql = '/*\nComment.\n*/\n/* Comment. */ update abc SET b = 1 WHERE a = 1' self.assertEqual(QueryHasDml(sql), True)
def testQueryHasDmlWithHashComment(self): sql = '# Comment.\nupdate abc SET b = 1 WHERE a = 1' self.assertEqual(QueryHasDml(sql), True)
def testQueryHasDmlDelete(self): sql = ' DELETE abc WHERE a = 1' self.assertEqual(QueryHasDml(sql), True)
def testQueryHasDmlUpdate(self): sql = ' update abc SET b = 1 WHERE a = 1' self.assertEqual(QueryHasDml(sql), True)
def testQueryHasDmlInsert(self): sql = ' INSERT abc (a, b, c) VALUES (1, 2, 3)' self.assertEqual(QueryHasDml(sql), True)