diff --git a/backend/jobs/workers/bigquery/bq_worker.py b/backend/jobs/workers/bigquery/bq_worker.py index 463a75c8a..499120443 100644 --- a/backend/jobs/workers/bigquery/bq_worker.py +++ b/backend/jobs/workers/bigquery/bq_worker.py @@ -1,4 +1,4 @@ -# Copyright 2020 Google Inc. All rights reserved. +# Copyright 2024 Google Inc. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,11 +14,11 @@ """CRMint's abstract worker dealing with BigQuery.""" - +import os import time +from google.api_core.client_info import ClientInfo from google.cloud import bigquery - from jobs.workers import worker @@ -42,7 +42,16 @@ class BQWorker(worker.Worker): ] def _get_client(self): - return bigquery.Client(client_options={'scopes': self._SCOPES}) + client_info = None + if 'REPORT_USAGE_ID' in os.environ: + client_id = os.getenv('REPORT_USAGE_ID') + opt_out = not bool(client_id) + if not opt_out: + client_info = ClientInfo(user_agent='cloud-solutions/crmint-usage-v3') + return bigquery.Client( + client_options={'scopes': self._SCOPES}, + client_info=client_info, + ) def _get_prefix(self): return f'{self._pipeline_id}_{self._job_id}_{self.__class__.__name__}' diff --git a/backend/tests/jobs/unit/workers/bq_worker_tests.py b/backend/tests/jobs/unit/workers/bq_worker_tests.py index 5846378e6..bb1bb9019 100644 --- a/backend/tests/jobs/unit/workers/bq_worker_tests.py +++ b/backend/tests/jobs/unit/workers/bq_worker_tests.py @@ -1,11 +1,14 @@ """Tests for bq_worker.""" from unittest import mock +import os from absl.testing import absltest from absl.testing import parameterized + from google.auth import credentials from google.cloud import bigquery +from google.api_core.client_info import ClientInfo from jobs.workers import worker from jobs.workers.bigquery import bq_worker @@ -82,6 +85,49 @@ def test_generates_proper_bq_table_name_from_params(self): self.assertEqual('a_project.a_dataset_id.a_table_id', worker._generate_qualified_bq_table_name()) +class BQWorkerGetClientTest(parameterized.TestCase): + + @parameterized.parameters( + { + 'report_usage_id_present': True, + 'client_info_user_agent': 'cloud-solutions/crmint-usage-v3', + }, + { + 'report_usage_id_present': False, + 'client_info_user_agent': None + }, + ) + def test_get_client_handles_report_usage_id( + self, report_usage_id_present, client_info_user_agent): + report_usage_id = 'some-usage-id' if report_usage_id_present else '' + with ( + mock.patch.dict( + os.environ, + {'REPORT_USAGE_ID': report_usage_id} + if report_usage_id_present + else {}, + ), + mock.patch('os.getenv', return_value=report_usage_id) as getenv_mock, + mock.patch('google.cloud.bigquery.Client') as client_mock, + ): + worker_inst = bq_worker.BQWorker({}, 0, 0) + worker_inst._get_client() + + if report_usage_id_present: + getenv_mock.assert_called_with('REPORT_USAGE_ID') + else: + getenv_mock.assert_not_called() + + client_mock.assert_called_once() + _, kwargs = client_mock.call_args + if report_usage_id_present: + self.assertIsInstance(kwargs['client_info'], ClientInfo) + self.assertEqual( + kwargs['client_info'].user_agent, client_info_user_agent + ) + else: + self.assertIsNone(kwargs.get('client_info')) + if __name__ == '__main__': absltest.main() diff --git a/cli/appcli.py b/cli/appcli.py index f92911d26..59bf14691 100644 --- a/cli/appcli.py +++ b/cli/appcli.py @@ -54,7 +54,7 @@ def _ask_permission(self): fg='yellow') msg += click.style(pkg_name, fg='red', bold=True) msg += click.style( - ' better! \nMay we anonymously report usage statistics to improve the' + ' better! \nMay we anonymously report usage statistics to improve the ' 'tool over time? \nMore info: https://github.com/google/crmint & ' 'https://google.github.io/crmint', fg='yellow') diff --git a/terraform/services.tf b/terraform/services.tf index a7b8fcf1b..df5959f1e 100644 --- a/terraform/services.tf +++ b/terraform/services.tf @@ -209,6 +209,10 @@ resource "google_cloud_run_service" "jobs_run" { name = "PUBSUB_VERIFICATION_TOKEN" value = random_id.pubsub_verification_token.b64_url } + env { + name = "REPORT_USAGE_ID" + value = var.report_usage_id + } } timeout_seconds = 900 # 15min