def __init__(self):
# Be sure to set blocking to false for timeout_ms to work properly
self.pipeline = TestPipeline(is_integration_test=True, blocking=False)
assert not self.pipeline.blocking
options = self.pipeline.get_pipeline_options().view_as(LoadTestOptions)
self.timeout_ms = options.timeout_ms
self.input_options = options.input_options
self.metrics_namespace = options.metrics_table or 'default'
publish_to_bq = options.publish_to_big_query
if publish_to_bq is None:
logging.info(
'Missing --publish_to_big_query option. Metrics will not '
'be published to BigQuery.')
if options.input_options is None:
logging.error('--input_options argument is required.')
sys.exit(1)
gcloud_options = self.pipeline.get_pipeline_options().view_as(
GoogleCloudOptions)
self.project_id = gcloud_options.project
self._metrics_monitor = MetricsReader(
publish_to_bq=publish_to_bq,
project_name=self.project_id,
bq_table=options.metrics_table,
bq_dataset=options.metrics_dataset,
influxdb_options=InfluxDBMetricsPublisherOptions(
options.influx_measurement,
options.influx_db_name,
options.influx_hostname,
os.getenv('INFLUXDB_USER'),
os.getenv('INFLUXDB_USER_PASSWORD')),
# Apply filter to prevent system metrics from being published
filters=MetricsFilter().with_namespace(self.metrics_namespace))
def _publish_metrics(self, pipeline, metric_value):
influx_options = InfluxDBMetricsPublisherOptions(
pipeline.get_option('influx_measurement'),
pipeline.get_option('influx_db_name'),
pipeline.get_option('influx_hostname'),
os.getenv('INFLUXDB_USER'),
os.getenv('INFLUXDB_USER_PASSWORD'),
)
metric_reader = MetricsReader(
project_name=pipeline.get_option('project'),
bq_table=pipeline.get_option('metrics_table'),
bq_dataset=pipeline.get_option('metrics_dataset'),
publish_to_bq=True,
influxdb_options=influx_options,
)
metric_reader.publish_values([(
'runtime',
metric_value,
)])