def execute(self, operation, parameters=None, job_id=None):
"""Prepare and execute a database operation.
.. note::
When setting query parameters, values which are "text"
(``unicode`` in Python2, ``str`` in Python3) will use
the 'STRING' BigQuery type. Values which are "bytes" (``str`` in
Python2, ``bytes`` in Python3), will use using the 'BYTES' type.
A `~datetime.datetime` parameter without timezone information uses
the 'DATETIME' BigQuery type (example: Global Pi Day Celebration
March 14, 2017 at 1:59pm). A `~datetime.datetime` parameter with
timezone information uses the 'TIMESTAMP' BigQuery type (example:
a wedding on April 29, 2011 at 11am, British Summer Time).
For more information about BigQuery data types, see:
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types
``STRUCT``/``RECORD`` and ``REPEATED`` query parameters are not
yet supported. See:
https://github.com/GoogleCloudPlatform/google-cloud-python/issues/3524
:type operation: str
:param operation: A Google BigQuery query string.
:type parameters: Mapping[str, Any] or Sequence[Any]
:param parameters:
(Optional) dictionary or sequence of parameter values.
:type job_id: str
:param job_id: (Optional) The job_id to use. If not set, a job ID
is generated at random.
"""
self._query_data = None
self._query_job = None
client = self.connection._client
# The DB-API uses the pyformat formatting, since the way BigQuery does
# query parameters was not one of the standard options. Convert both
# the query and the parameters to the format expected by the client
# libraries.
formatted_operation = _format_operation(
operation, parameters=parameters)
query_parameters = _helpers.to_query_parameters(parameters)
config = job.QueryJobConfig()
config.query_parameters = query_parameters
config.use_legacy_sql = False
self._query_job = client.query(
formatted_operation, job_config=config, job_id=job_id)
# Wait for the query to finish.
try:
self._query_job.result()
except google.cloud.exceptions.GoogleCloudError as exc:
raise exceptions.DatabaseError(exc)
query_results = self._query_job._query_results
self._set_rowcount(query_results)
self._set_description(query_results.schema)
def test_bigquery_magic_w_maximum_bytes_billed_w_context_setter():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context._project = None
magics.context.default_query_job_config = job.QueryJobConfig(
maximum_bytes_billed=10203
)
project = "test-project"
job_reference = copy.deepcopy(JOB_REFERENCE_RESOURCE)
job_reference["projectId"] = project
query = "SELECT 17 AS num"
resource = copy.deepcopy(QUERY_RESOURCE)
resource["jobReference"] = job_reference
resource["configuration"]["query"]["query"] = query
data = {"jobReference": job_reference, "totalRows": 0, "rows": []}
credentials_mock = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)
default_patch = mock.patch(
"google.auth.default", return_value=(credentials_mock, "general-project")
)
conn = magics.context._connection = make_connection(resource, data)
list_rows_patch = mock.patch(
"google.cloud.bigquery.client.Client.list_rows",
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
)
with list_rows_patch, default_patch:
ip.run_cell_magic("bigquery", "", query)
_, req = conn.api_request.call_args_list[0]
sent_config = req["data"]["configuration"]["query"]
assert sent_config["maximumBytesBilled"] == "10203"
def test_execute_w_default_config(self):
from google.cloud.bigquery.dbapi import connect
from google.cloud.bigquery import job
default_config = job.QueryJobConfig(use_legacy_sql=False, flatten_results=True)
client = self._mock_client(
rows=[], num_dml_affected_rows=0, default_query_job_config=default_config
)
connection = connect(client)
cursor = connection.cursor()
cursor.execute("SELECT 1;", job_id="foo")
_, kwargs = client.query.call_args
used_config = kwargs["job_config"]
expected_config = job.QueryJobConfig(
use_legacy_sql=False, flatten_results=True, query_parameters=[]
)
self.assertEqual(used_config._properties, expected_config._properties)
def test_execute_custom_job_config_wo_default_config(self):
from google.cloud.bigquery.dbapi import connect
from google.cloud.bigquery import job
config = job.QueryJobConfig(use_legacy_sql=True)
client = self._mock_client(rows=[], num_dml_affected_rows=0)
connection = connect(client)
cursor = connection.cursor()
cursor.execute("SELECT 1;", job_id="foo", job_config=config)
args, kwargs = client.query.call_args
self.assertEqual(args[0], "SELECT 1;")
self.assertEqual(kwargs["job_id"], "foo")
self.assertEqual(kwargs["job_config"], config)
def create_from_query(cls, query, flatten_results=True):
"""
Load instances through a query job.
The job is asynchronous but this function will wait for the job to complete.
See https://cloud.google.com/bigquery/docs/writing-results
Note that this method must compile the sql query to a string.
It does so using sqlalchemy_query.statement.compile(compile_kwargs={"literal_binds": True}).
This will fail for certain queries and should not be used for queries which depend on untrusted input.
See https://docs.sqlalchemy.org/en/13/faq/sqlexpressions.html for more information.
Args:
query (BigQueryQuery): A query object whose results are
to be appended to the table.
flatten_results (Optional[bool]): If True, will flatten the query results.
Defaults to True.
"""
client = DatabaseContext.get_session().connection().connection._client
table_ref = _get_table_ref(cls.__table__.name, client)
job_config = bigquery_job.QueryJobConfig(
destination=table_ref,
create_disposition=bigquery_job.CreateDisposition.CREATE_NEVER,
write_disposition=bigquery_job.WriteDisposition.WRITE_APPEND,
flatten_results=flatten_results,
allow_large_results=not flatten_results,
)
dialect = DatabaseContext.get_engine().dialect
compiled_sql = query.sqlalchemy_query.statement.compile(
dialect=dialect, compile_kwargs={
'literal_binds': True,
})
raw_sql = str(compiled_sql)
query_job = client.query(raw_sql, job_config=job_config)
try:
query_job.result()
except Exception as e:
raise exceptions.DatabaseError('{}\n{}\n{}'.format(
query_job.errors,
'{}({})'.format(type(e), e),
query_job.error_result,
))
if ((query_job.error_result and len(query_job.error_result) > 0)
or (query_job.errors and len(query_job.errors) > 0)):
raise exceptions.DatabaseError('{}\n{}'.format(
query_job.errors, query_job.error_result))
def test__run_query_dry_run_without_errors_is_silent():
magics.context._credentials = None
sql = "SELECT 17"
client_patch = mock.patch("google.cloud.bigquery.magics.bigquery.Client",
autospec=True)
job_config = job.QueryJobConfig()
job_config.dry_run = True
with client_patch as client_mock, io.capture_output() as captured:
client_mock().query(sql).job_id = None
magics._run_query(client_mock(), sql, job_config=job_config)
assert len(captured.stderr) == 0
assert len(captured.stdout) == 0
def _execute(self, formatted_operation, parameters, job_id, job_config,
parameter_types):
self._query_data = None
self._query_job = None
client = self.connection._client
# The DB-API uses the pyformat formatting, since the way BigQuery does
# query parameters was not one of the standard options. Convert both
# the query and the parameters to the format expected by the client
# libraries.
query_parameters = _helpers.to_query_parameters(
parameters, parameter_types)
if client._default_query_job_config:
if job_config:
config = job_config._fill_from_default(
client._default_query_job_config)
else:
config = copy.deepcopy(client._default_query_job_config)
else:
config = job_config or job.QueryJobConfig(use_legacy_sql=False)
config.query_parameters = query_parameters
self._query_job = client.query(formatted_operation,
job_config=config,
job_id=job_id)
if self._query_job.dry_run:
self._set_description(schema=None)
self.rowcount = 0
return
# Wait for the query to finish.
try:
self._query_job.result()
except google.cloud.exceptions.GoogleCloudError as exc:
raise exceptions.DatabaseError(exc)
query_results = self._query_job._query_results
self._set_rowcount(query_results)
self._set_description(query_results.schema)
