def test_compose_simple(self) -> None:
duration1 = duration_pb2.Duration(seconds=1, nanos=2)
duration2 = duration_pb2.Duration(nanos=3)
duration = (ProtobufBuilder(duration_pb2.Duration).compose(
duration1).compose(duration2).build())
self.assertEqual(duration, duration_pb2.Duration(seconds=1, nanos=3))
duration = (ProtobufBuilder(duration_pb2.Duration).compose(
duration2).compose(duration1).build())
self.assertEqual(duration, duration_pb2.Duration(seconds=1, nanos=2))
def create_task(self,
*,
task_id: str,
queue_name: str,
relative_uri: str,
body: Optional[Dict[str, str]] = None,
schedule_delay_seconds: int = 0,
http_method: HttpMethod = HttpMethod.POST):
"""Creates a task with the given details.
Args:
task_id: Id of the task to include in the task name
queue_name: The queue on which to schedule the task
schedule_delay_seconds: The number of seconds by which to delay the
scheduling of the given task.
relative_uri: The relative uri to hit.
body: Dictionary of values that will be converted to JSON and
included in the request.
http_method: The method for this request (i.e. GET or POST)
"""
task_name = self.format_task_path(queue_name, task_id)
schedule_timestamp = None
if schedule_delay_seconds > 0:
schedule_time_sec = \
int(datetime.utcnow().timestamp()) + schedule_delay_seconds
schedule_timestamp = \
timestamp_pb2.Timestamp(seconds=schedule_time_sec)
task_builder = ProtobufBuilder(task_pb2.Task).update_args(
name=task_name,
app_engine_http_request={
'relative_uri': relative_uri,
},
)
if schedule_timestamp:
task_builder.update_args(schedule_time=schedule_timestamp, )
if http_method is not None:
task_builder.update_args(app_engine_http_request={
'http_method': http_method.value,
}, )
if http_method in (HttpMethod.POST, HttpMethod.PUT, HttpMethod.PATCH):
if body is None:
body = {}
task_builder.update_args(
app_engine_http_request={'body': json.dumps(body).encode()}, )
task = task_builder.build()
logging.info("Queueing task to queue [%s]: [%s]", queue_name,
task.name)
retry_grpc(self.NUM_GRPC_RETRIES, self.client.create_task,
self.format_queue_path(queue_name), task)
def test_build_multiple_ways(self) -> None:
duration1 = duration_pb2.Duration(seconds=1, nanos=2)
duration = (ProtobufBuilder(
duration_pb2.Duration).compose(duration1).update_args(
nanos=3, ).build())
self.assertEqual(duration, duration_pb2.Duration(seconds=1, nanos=3))
duration2 = duration_pb2.Duration(nanos=2)
duration = (ProtobufBuilder(duration_pb2.Duration).update_args(
seconds=1, ).compose(duration2).build())
self.assertEqual(duration, duration_pb2.Duration(seconds=1, nanos=2))
def test_update_args(self) -> None:
duration = (ProtobufBuilder(duration_pb2.Duration).update_args(
seconds=1,
nanos=2,
).update_args(nanos=3).build())
self.assertEqual(duration, duration_pb2.Duration(seconds=1, nanos=3))
duration = (ProtobufBuilder(
duration_pb2.Duration).update_args(nanos=3).update_args(
seconds=1,
nanos=2,
).build())
self.assertEqual(duration, duration_pb2.Duration(seconds=1, nanos=2))
def _queue_config_with_name(
client_wrapper: GoogleCloudTasksClientWrapper,
base_config: queue_pb2.Queue,
queue_name: str,
) -> queue_pb2.Queue:
return (ProtobufBuilder(queue_pb2.Queue).compose(base_config).update_args(
name=client_wrapper.format_queue_path(queue_name), ).build())
def test_multilevel_compose(self) -> None:
base_queue = queue_pb2.Queue(
rate_limits=queue_pb2.RateLimits(
max_dispatches_per_second=100,
max_concurrent_dispatches=1,
),
retry_config=queue_pb2.RetryConfig(max_attempts=5, ),
)
queue = (ProtobufBuilder(queue_pb2.Queue).compose(base_queue).compose(
queue_pb2.Queue(rate_limits=queue_pb2.RateLimits(
max_dispatches_per_second=50, ), )).update_args(
retry_config=queue_pb2.RetryConfig(max_doublings=4,
)).build())
expected_queue = queue_pb2.Queue(
rate_limits=queue_pb2.RateLimits(
max_dispatches_per_second=50,
max_concurrent_dispatches=1,
),
retry_config=queue_pb2.RetryConfig(
max_attempts=5,
max_doublings=4,
),
)
self.assertEqual(queue, expected_queue)
def _build_cloud_task_queue_configs(
client_wrapper: GoogleCloudTasksClientWrapper
) -> List[queue_pb2.Queue]:
"""Builds a list of configurations for all Google Cloud Tasks queues that
should be deployed in this environment.
"""
queues = []
# Direct ingest queues for handling /process_job requests
for queue_name in [
DIRECT_INGEST_STATE_PROCESS_JOB_QUEUE_V2,
DIRECT_INGEST_JAILS_PROCESS_JOB_QUEUE_V2,
DIRECT_INGEST_BQ_IMPORT_EXPORT_QUEUE_V2,
DIRECT_INGEST_SCHEDULER_QUEUE_V2
]:
queues.append(
_queue_config_with_name(client_wrapper,
DIRECT_INGEST_QUEUE_BASE_CONFIG,
queue_name))
queues.append(
ProtobufBuilder(queue_pb2.Queue).compose(
_queue_config_with_name(
client_wrapper, SCRAPER_PHASE_QUEUE_CONFIG,
SCRAPER_PHASE_QUEUE_V2)).update_args(
rate_limits=queue_pb2.RateLimits(
max_dispatches_per_second=1, ), ).build())
queues.append(
_queue_config_with_name(client_wrapper, BIGQUERY_QUEUE_CONFIG,
BIGQUERY_QUEUE_V2))
queues.append(
_queue_config_with_name(client_wrapper, JOB_MONITOR_QUEUE_CONFIG,
JOB_MONITOR_QUEUE_V2))
for vendor in vendors.get_vendors():
queue_params = vendors.get_vendor_queue_params(vendor)
if queue_params is None:
continue
vendor_queue_name = \
'vendor-{}-scraper-v2'.format(vendor.replace('_', '-'))
queue = ProtobufBuilder(queue_pb2.Queue).compose(
_queue_config_with_name(
client_wrapper,
BASE_SCRAPER_QUEUE_CONFIG, vendor_queue_name)).compose(
queue_pb2.Queue(**queue_params)).build()
queues.append(queue)
for region in regions.get_supported_regions():
if region.shared_queue or not region.is_ingest_launched_in_env():
continue
queue = _queue_config_with_name(client_wrapper,
BASE_SCRAPER_QUEUE_CONFIG,
region.get_queue_name())
if region.queue:
queue = ProtobufBuilder(queue_pb2.Queue).compose(queue).compose(
queue_pb2.Queue(**region.queue)).build()
queues.append(queue)
return queues
def test_simple_build(self) -> None:
duration = ProtobufBuilder(duration_pb2.Duration).build()
self.assertEqual(type(duration), duration_pb2.Duration)
def test_bad_arg_raises_error(self) -> None:
duration_builder = ProtobufBuilder(duration_pb2.Duration)
with self.assertRaises(ValueError):
duration_builder.update_args(name="this-field-does-not-exist")
def test_bad_type_raises_error(self) -> None:
duration_builder = ProtobufBuilder(duration_pb2.Duration)
with self.assertRaises(TypeError):
duration_builder.compose(queue_pb2.Queue())
def create_task(self,
*,
queue_name: str,
relative_uri: str,
task_id: Optional[str] = None,
body: Optional[Dict[str, Any]] = None,
schedule_delay_seconds: int = 0,
http_method: HttpMethod = HttpMethod.POST) -> None:
"""Creates a task with the given details.
Args:
task_id: (optional) ID of the task to include in the task name
Specifying task IDs enables task de-duplication for the queue. Subsequent requests to enqueue a task with the
same ID as a recently completed task will raise `409 Conflict (entity already exists)
If left unspecified, Cloud Tasks will automatically generate an ID for the task
https://cloud.google.com/tasks/docs/reference/rest/v2beta3/projects.locations.queues.tasks/create#body.request_body.FIELDS.task
queue_name: The queue on which to schedule the task
schedule_delay_seconds: The number of seconds by which to delay the
scheduling of the given task.
relative_uri: The relative uri to hit.
body: Dictionary of values that will be converted to JSON and
included in the request.
http_method: The method for this request (i.e. GET or POST)
"""
schedule_timestamp = None
if schedule_delay_seconds > 0:
schedule_timestamp = timestamp_pb2.Timestamp()
schedule_timestamp.FromDatetime(
datetime.now(tz=pytz.UTC) +
timedelta(seconds=schedule_delay_seconds))
task_builder = ProtobufBuilder(task_pb2.Task).update_args(
app_engine_http_request={
"relative_uri": relative_uri,
}, )
if task_id is not None:
task_name = self.format_task_path(queue_name, task_id)
task_builder.update_args(name=task_name, )
if schedule_timestamp:
task_builder.update_args(schedule_time=schedule_timestamp, )
if http_method is not None:
task_builder.update_args(app_engine_http_request={
"http_method": http_method.value,
}, )
if http_method in (HttpMethod.POST, HttpMethod.PUT, HttpMethod.PATCH):
if body is None:
body = {}
task_builder.update_args(
app_engine_http_request={"body": json.dumps(body).encode()}, )
task = task_builder.build()
logging.info("Queueing task to queue [%s]: [%s]", queue_name,
task.name)
retry_grpc(
self.NUM_GRPC_RETRIES,
self.client.create_task,
parent=self.format_queue_path(queue_name),
task=task,
)