[v2] Add --apn-id global arg for PRM user-agent attribution

awscli/customizations/globalargs.py

@@ -42,6 +42,11 @@ def register_parse_global_args(cli):
         resolve_cli_connect_timeout,
         unique_id='resolve-cli-connect-timeout',
     )
+    cli.register(
+        'session-initialized',
+        resolve_apn_id,
+        unique_id='resolve-apn-id',
+    )
 
 
 def resolve_types(parsed_args, **kwargs):
@@ -133,3 +138,53 @@ def _update_default_client_config(session, arg_name, arg_value):
     if current_default_config is not None:
         new_default_config = current_default_config.merge(new_default_config)
     session.set_default_client_config(new_default_config)
+
+
+# Allowed characters for APN ID values: ASCII letters, digits,
+# underscore, and hyphen.  Simple set check avoids regex entirely.
+_APN_ID_ALLOWED_CHARS = frozenset(
+    'abcdefghijklmnopqrstuvwxyz' 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' '0123456789' '_-'
+)
+
+
+def _is_valid_apn_id(value):
+    """Validate an APN ID value: 1-255 chars from the allowed set."""
+    if not value or len(value) > 255:
+        return False
+    return all(c in _APN_ID_ALLOWED_CHARS for c in value)
+
+
+def resolve_apn_id(parsed_args, session, **kwargs):
+    """Resolve the APN ID from CLI flag, env var, or config file.
+
+    Registered on 'session-initialized' so the profile is already set
+    and config file values are accessible.
+    """
+    apn_id = getattr(parsed_args, 'apn_id', None)
+    if apn_id is None:
+        apn_id = os.environ.get('AWS_APN_ID')
+    if apn_id is None:
+        try:
+            scoped_config = session.get_scoped_config()
+            apn_id = scoped_config.get('apn_id')
+        except Exception:
+            pass
+    if not apn_id:
+        return
+    if not _is_valid_apn_id(apn_id):
+        raise ParamValidationError(
+            "Bad value for --apn-id %r: must be 1-255 characters and may "
+            "only contain letters, digits, '_', and '-'." % apn_id
+        )
+    # Format defined by the AWS PRM onboarding guide:
+    # https://docs.aws.amazon.com/PRM/latest/aws-prm-onboarding-guide/user-agent-string.html
+    # The trailing '$' is a required end delimiter, not a regex anchor.
+    # The user provides the full identifier including the type prefix
+    # (e.g. "pc_PRODUCTCODE" or "ra_ATTRIBUTIONID").
+    marker = 'APN_1.1/%s$' % apn_id
+    existing = session.user_agent_extra or ''
+    if marker in existing.split():
+        return
+    session.user_agent_extra = (
+        '%s %s' % (existing, marker) if existing else marker
+    )

awscli/data/cli.json

@@ -97,6 +97,10 @@
                 "enhanced"
             ],
             "help": "<p>The formatting style for error output. By default, errors are displayed in enhanced format.</p>"
+        },
+        "apn-id": {
+            "dest": "apn_id",
+            "help": "<p>An AWS Partner Network identifier used for Partner Revenue Measurement (PRM) attribution. The value should include the type prefix (e.g. <code>pc_PRODUCTCODE</code> for product codes or <code>ra_PARTNERID</code> for referral agents). The value is appended to the User-Agent header of all requests as <code>APN_1.1/&lt;value&gt;$</code>. Overrides the <code>AWS_APN_ID</code> environment variable and the <code>apn_id</code> setting in the shared config file. See <a href=\"https://docs.aws.amazon.com/PRM/latest/aws-prm-onboarding-guide/user-agent-string.html\">the PRM onboarding guide</a> for the canonical format.</p>"
         }
     }
 }

awscli/examples/global_options.rst

@@ -116,3 +116,7 @@
   *   enhanced
 
 
+``--apn-id`` (string)
+
+  An AWS Partner Network identifier used for Partner Revenue Measurement (PRM) attribution. The value should include the type prefix (e.g. ``pc_PRODUCTCODE`` for product codes or ``ra_PARTNERID`` for referral agents). The value is appended to the User-Agent header of all requests as ``APN_1.1/<value>$``. Overrides the ``AWS_APN_ID`` environment variable and the ``apn_id`` setting in the shared config file. See `the PRM onboarding guide <https://docs.aws.amazon.com/PRM/latest/aws-prm-onboarding-guide/user-agent-string.html>`__ for the canonical format.
+

awscli/examples/global_synopsis.rst

@@ -17,3 +17,4 @@
 [--cli-auto-prompt]
 [--no-cli-auto-prompt]
 [--cli-error-format <value>]
+[--apn-id <value>]

tests/unit/customizations/test_globalargs.py

@@ -194,3 +194,163 @@ def test_cli_connect_timeout_for_blocking(self):
         self.assertEqual(
             session.get_default_client_config().connect_timeout, None
         )
+
+
+class TestResolveApnId(unittest.TestCase):
+    def _make_session(self, config_value=None, user_agent_extra=''):
+        session = mock.Mock()
+        session.user_agent_extra = user_agent_extra
+        if config_value is not None:
+            session.get_scoped_config.return_value = {'apn_id': config_value}
+        else:
+            session.get_scoped_config.return_value = {}
+        return session
+
+    def test_no_apn_id_leaves_user_agent_unchanged(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs()
+        with mock.patch.dict(os.environ, {}, clear=True):
+            globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(session.user_agent_extra, 'botocore/1.0')
+
+    def test_cli_flag_appends_apn_marker(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='pc_abc123')
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_abc123$'
+        )
+
+    def test_env_var_used_when_flag_missing(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs()
+        with mock.patch.dict(os.environ, {'AWS_APN_ID': 'pc_env123'}):
+            globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_env123$'
+        )
+
+    def test_config_file_used_when_flag_and_env_missing(self):
+        session = self._make_session(
+            config_value='pc_cfg123', user_agent_extra='botocore/1.0'
+        )
+        parsed_args = FakeParsedArgs()
+        with mock.patch.dict(os.environ, {}, clear=True):
+            globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_cfg123$'
+        )
+
+    def test_cli_flag_overrides_env_and_config(self):
+        session = self._make_session(
+            config_value='pc_from-config', user_agent_extra='botocore/1.0'
+        )
+        parsed_args = FakeParsedArgs(apn_id='pc_from-flag')
+        with mock.patch.dict(os.environ, {'AWS_APN_ID': 'pc_from-env'}):
+            globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra,
+            'botocore/1.0 APN_1.1/pc_from-flag$',
+        )
+
+    def test_env_overrides_config(self):
+        session = self._make_session(
+            config_value='pc_from-config', user_agent_extra='botocore/1.0'
+        )
+        parsed_args = FakeParsedArgs()
+        with mock.patch.dict(os.environ, {'AWS_APN_ID': 'ra_from-env'}):
+            globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra,
+            'botocore/1.0 APN_1.1/ra_from-env$',
+        )
+
+    def test_appends_when_user_agent_extra_is_empty(self):
+        session = self._make_session(user_agent_extra='')
+        parsed_args = FakeParsedArgs(apn_id='pc_abc123')
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(session.user_agent_extra, 'APN_1.1/pc_abc123$')
+
+    def test_appends_when_user_agent_extra_is_none(self):
+        session = self._make_session(user_agent_extra=None)
+        session.user_agent_extra = None
+        parsed_args = FakeParsedArgs(apn_id='pc_abc123')
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(session.user_agent_extra, 'APN_1.1/pc_abc123$')
+
+    def test_idempotent_when_marker_already_present(self):
+        session = self._make_session(
+            user_agent_extra='botocore/1.0 APN_1.1/pc_abc123$'
+        )
+        parsed_args = FakeParsedArgs(apn_id='pc_abc123')
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_abc123$'
+        )
+
+    def test_empty_string_treated_as_unset(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='')
+        with mock.patch.dict(os.environ, {}, clear=True):
+            globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(session.user_agent_extra, 'botocore/1.0')
+
+    def test_rejects_value_with_whitespace(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='bad value')
+        with self.assertRaises(ParamValidationError):
+            globalargs.resolve_apn_id(parsed_args, session)
+
+    def test_rejects_value_with_special_chars(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='bad/value')
+        with self.assertRaises(ParamValidationError):
+            globalargs.resolve_apn_id(parsed_args, session)
+
+    def test_rejects_value_with_dot(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='1.0')
+        with self.assertRaises(ParamValidationError):
+            globalargs.resolve_apn_id(parsed_args, session)
+
+    def test_rejects_value_over_255_chars(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='a' * 256)
+        with self.assertRaises(ParamValidationError):
+            globalargs.resolve_apn_id(parsed_args, session)
+
+    def test_accepts_max_length_value(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        long_id = 'a' * 255
+        parsed_args = FakeParsedArgs(apn_id=long_id)
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra,
+            'botocore/1.0 APN_1.1/%s$' % long_id,
+        )
+
+    def test_accepts_dash_and_underscore(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='pc_Pc-1_test')
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra,
+            'botocore/1.0 APN_1.1/pc_Pc-1_test$',
+        )
+
+    def test_accepts_ra_prefix(self):
+        session = self._make_session(user_agent_extra='botocore/1.0')
+        parsed_args = FakeParsedArgs(apn_id='ra_MyPartner123')
+        globalargs.resolve_apn_id(parsed_args, session)
+        self.assertEqual(
+            session.user_agent_extra,
+            'botocore/1.0 APN_1.1/ra_MyPartner123$',
+        )
+
+    def test_handler_is_registered(self):
+        cli = mock.Mock()
+        globalargs.register_parse_global_args(cli)
+        registered = [
+            c.kwargs.get('unique_id') for c in cli.register.call_args_list
+        ]
+        self.assertIn('resolve-apn-id', registered)