diff --git a/infra_lib/telemetry/__init__.py b/infra_lib/telemetry/__init__.py index 7aef80ddc..841c66357 100644 --- a/infra_lib/telemetry/__init__.py +++ b/infra_lib/telemetry/__init__.py @@ -72,13 +72,15 @@ def initialize(service_name, if not sys.platform.startswith('linux'): return - if not is_google_host(): - return - cfg = config.Config(cfg_file) if cfg.trace_config.disabled(): return + bot_enabled = (cfg.trace_config.has_enabled() + and cfg.trace_config.enabled_reason == 'BOT_USER') + if not is_google_host() and not bot_enabled: + return + if not cfg.trace_config.has_enabled(): if cfg.root_config.notice_countdown > -1: print(notice.format(run_count=cfg.root_config.notice_countdown, diff --git a/infra_lib/telemetry/__main__.py b/infra_lib/telemetry/__main__.py index ca44b655f..408e8c813 100755 --- a/infra_lib/telemetry/__main__.py +++ b/infra_lib/telemetry/__main__.py @@ -26,14 +26,26 @@ def main(): default=None, help='Enable telemetry collection.') + parser.add_argument('--bot-enable', + '-b', + dest='bot_enable', + action='store_true', + default=False, + help='Enable for bots. Ignores googler check. ' + 'Not for human users.') + args = parser.parse_args() if args.enable is not None: cfg = config.Config(config.DEFAULT_CONFIG_FILE) cfg.trace_config.update(args.enable, 'USER') cfg.flush() + elif args.bot_enable: + cfg = config.Config(config.DEFAULT_CONFIG_FILE) + cfg.trace_config.update(args.bot_enable, 'BOT_USER') + cfg.flush() else: - print('Error: --enable or --disable flag is required.') + print('Error: --enable --disable or --bot-enable flag is required.') if __name__ == '__main__': diff --git a/infra_lib/telemetry/config.py b/infra_lib/telemetry/config.py index 3a9198379..8d66147f6 100644 --- a/infra_lib/telemetry/config.py +++ b/infra_lib/telemetry/config.py @@ -45,7 +45,8 @@ class TraceConfig: if not self.has_enabled() or self.enabled: self.gen_id() - def update(self, enabled: bool, reason: Literal["AUTO", "USER"]) -> None: + def update(self, enabled: bool, reason: Literal["AUTO", "USER", + "BOT_USER"]) -> None: """Update the config.""" self._config[TRACE_SECTION_KEY][ENABLED_KEY] = str(enabled) self._config[TRACE_SECTION_KEY][ENABLED_REASON_KEY] = reason @@ -100,7 +101,7 @@ class TraceConfig: return self._config[TRACE_SECTION_KEY].getboolean(ENABLED_KEY, False) @property - def enabled_reason(self) -> Literal["AUTO", "USER"]: + def enabled_reason(self) -> Literal["AUTO", "USER", "BOT_USER"]: """Value of trace.enabled_reason property in telemetry.cfg.""" return self._config[TRACE_SECTION_KEY].get(ENABLED_REASON_KEY, "AUTO")