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, )