def run_async_query(
self,
query: str,
job_id_prefix: str = None,
ems_query_job_config: EmsQueryJobConfig = EmsQueryJobConfig(
priority=EmsJobPriority.INTERACTIVE)
) -> str:
return self.__execute_query_job(
query=query,
ems_query_job_config=ems_query_job_config,
job_id_prefix=job_id_prefix).job_id
def setUp(self):
self.client_mock = Mock()
self.query_job_mock = Mock(QueryJob)
self.query_job_mock.priority = "INTERACTIVE"
self.query_job_mock.create_disposition = None
self.query_job_mock.write_disposition = None
self.query_job_mock.time_partitioning = TimePartitioning(
"DAY", "a", None, None)
self.query_config = EmsQueryJobConfig(
destination_project_id="some_destination_project_id",
destination_dataset="some_dataset",
destination_table="some_table",
labels={"label1": "label1_value"})
def test_run_sync_query_withDestinationSet(self):
ems_query_job_config = EmsQueryJobConfig(
destination_dataset=ItEmsBigqueryClient.DATASET.dataset_id,
destination_table=self.test_table.table_id)
query_with_destination_result = list(
self.client.run_sync_query(
self.DUMMY_SELECT_TO_TABLE,
ems_query_job_config=ems_query_job_config))
query_result = list(
self.client.run_sync_query(
self.SELECT_TEMPLATE.format(self.__get_table_path())))
assert [{"int_data": 1, "str_data": "hello"}] == query_result
assert query_with_destination_result == query_result
def test_run_async_query_setsDestinationProjectIdToDefaultIfNotGiven(
self, bigquery_module_patch: bigquery):
ems_bigquery_client = self.__setup_client(bigquery_module_patch)
query_config = EmsQueryJobConfig(destination_project_id=None,
destination_dataset="some_dataset",
destination_table="some_table")
ems_bigquery_client.run_async_query(self.QUERY,
ems_query_job_config=query_config)
bigquery_module_patch.Client.assert_called_once_with("some-project-id",
location="EU")
arguments = self.client_mock.query.call_args_list[0][1]
self.assertEqual(arguments["job_config"].destination.project,
"some-project-id")
def run_sync_query(
self,
query: str,
ems_query_job_config: EmsQueryJobConfig = EmsQueryJobConfig(
priority=EmsJobPriority.INTERACTIVE),
job_id_prefix: str = None) -> Iterable:
LOGGER.info("Sync query executed with priority: %s",
ems_query_job_config.priority)
try:
return self.__get_mapped_iterator(
self.__execute_query_job(
query=query,
ems_query_job_config=ems_query_job_config,
job_id_prefix=job_id_prefix).result())
except GoogleAPIError as e:
raise EmsApiError(
"Error caused while running query | {} |: {}!".format(
query, e.args[0]))
def __convert_to_ems_job(job):
if isinstance(job, QueryJob):
destination = job.destination
table_id, dataset_id, project_id = \
(destination.table_id, destination.dataset_id, destination.project) \
if destination is not None else (None, None, None)
config = EmsQueryJobConfig(
priority=EmsJobPriority[job.priority],
destination_project_id=project_id,
destination_dataset=dataset_id,
destination_table=table_id,
create_disposition=EmsBigqueryClient.
__convert_to_ems_create_disposition(job.create_disposition),
write_disposition=EmsBigqueryClient.
__convert_to_ems_write_disposition(job.write_disposition),
time_partitioning=EmsBigqueryClient.
__convert_to_ems_time_partitioning(job.time_partitioning),
labels=job.labels)
return EmsQueryJob(job.job_id, job.query, config,
EmsJobState(job.state), job.error_result,
job.created)
elif isinstance(job, LoadJob):
destination = job.destination
table_id, dataset_id, project_id = destination.table_id, destination.dataset_id, destination.project
schema = {
"fields": _build_schema_resource(job.schema)
} if job.schema else []
config = EmsLoadJobConfig(
schema=schema,
source_uri_template=job.source_uris[0]
if job.source_uris else None,
destination_project_id=project_id,
destination_dataset=dataset_id,
destination_table=table_id,
create_disposition=EmsBigqueryClient.
__convert_to_ems_create_disposition(job.create_disposition),
write_disposition=EmsBigqueryClient.
__convert_to_ems_write_disposition(job.write_disposition),
labels=job.labels)
return EmsLoadJob(job_id=job.job_id,
load_config=config,
state=EmsJobState(job.state),
error_result=None,
created=job.created)
elif isinstance(job, ExtractJob):
table = f'{job.source.project}.{job.source.dataset_id}.{job.source.table_id}'
destination_uris = job.destination_uris
job_config = EmsExtractJobConfig(
compression=Compression(job.compression)
if job.compression else Compression.NONE,
destination_format=DestinationFormat(job.destination_format)
if job.destination_format else DestinationFormat.CSV,
field_delimiter=job.field_delimiter,
print_header=job.print_header,
labels=job.labels)
return EmsExtractJob(job_id=job.job_id,
table=table,
destination_uris=destination_uris,
job_config=job_config,
state=EmsJobState(job.state),
error_result=job.error_result,
created=job.created)
else:
LOGGER.warning(
f"Unexpected job type for : {job.job_id}, with type class: {job.__class__}"
)
return None
def setUp(self):
self.query_config = EmsQueryJobConfig()
self.expected_error_result = {"some": "error", "happened": "here"}
self.ems_query_job = EmsQueryJob("test-job-id", "query",
self.query_config, EmsJobState.DONE,
self.expected_error_result)