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)