From 60eff3d026b0feaf39b83a65b1afd17f9a67b341 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 12 Oct 2026 04:18:27 -0700 Subject: [PATCH] [Python] Pass custom OSS User-Agent to PyArrow S3 clients --- .../pypaimon/filesystem/oss_user_agent.py | 12 +++++ .../pypaimon/filesystem/pyarrow_file_io.py | 18 ++++++- .../pypaimon/tests/oss_user_agent_test.py | 53 ++++++++++++++++++- 3 files changed, 80 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/filesystem/oss_user_agent.py b/paimon-python/pypaimon/filesystem/oss_user_agent.py index 844499c3a53d..36a1e2fd6609 100644 --- a/paimon-python/pypaimon/filesystem/oss_user_agent.py +++ b/paimon-python/pypaimon/filesystem/oss_user_agent.py @@ -22,6 +22,7 @@ from pypaimon import build_info from pypaimon.common.options import Options +from pypaimon.common.user_agent import format_user_agent USER_AGENT_MODULE = "fs.oss.user.agent.module" USER_AGENT_FEATURES = "fs.oss.user.agent.features" @@ -63,6 +64,17 @@ def extended(options: Options) -> Optional[str]: return " ".join(parts) if parts else None +def legacy_user_agent(options: Options, transport: str) -> Optional[str]: + """Build a per-client User-Agent only when the caller requests customization.""" + if not any(_effective(options, key) for key in ( + USER_AGENT_MODULE, USER_AGENT_FEATURES, USER_AGENT_EXTENDED, + COMMON_USER_AGENT_MODULE, COMMON_USER_AGENT_FEATURES, COMMON_USER_AGENT_EXTENDED)): + return None + name = module(options) or identity() + tags = [feature for feature in features(options).split() if feature != name] + return format_user_agent(name, transport, tags, extended(options)) + + def _effective(options: Options, *keys: str) -> Optional[str]: """The first non-blank value among ``keys``, stripped.""" data = options.to_map() diff --git a/paimon-python/pypaimon/filesystem/pyarrow_file_io.py b/paimon-python/pypaimon/filesystem/pyarrow_file_io.py index dcb1c5d5c349..9314e06587d7 100644 --- a/paimon-python/pypaimon/filesystem/pyarrow_file_io.py +++ b/paimon-python/pypaimon/filesystem/pyarrow_file_io.py @@ -258,7 +258,7 @@ def _initialize_oss_fs(self, path) -> FileSystem: # Uses setdefault so that an explicit user setting is never overridden. # Note: this is process-wide and affects all AWS SDK clients. os.environ.setdefault("AWS_EC2_METADATA_DISABLED", "true") - # S3FileSystem takes no User-Agent; this adds app/pypaimon/ to it, process-wide. + # Keep the default SDK application marker for PyArrow builds without a User-Agent option. os.environ.setdefault("AWS_SDK_UA_APP_ID", oss_user_agent.identity()) client_kwargs = { @@ -278,7 +278,21 @@ def _initialize_oss_fs(self, path) -> FileSystem: retry_config = self._create_s3_retry_config() client_kwargs.update(retry_config) - return pafs.S3FileSystem(**client_kwargs) + user_agent = oss_user_agent.legacy_user_agent( + self.properties, "pyarrow/{}".format(pyarrow.__version__)) + if user_agent is not None: + client_kwargs["user_agent"] = user_agent + try: + return pafs.S3FileSystem(**client_kwargs) + except TypeError as error: + if user_agent is not None and "unexpected keyword argument 'user_agent'" in str(error): + self.logger.warning( + "This PyArrow build does not support S3FileSystem(user_agent=...); " + "custom OSS User-Agent options are ignored. Use a PyArrow build " + "with this option or fs.oss.impl=jindo to apply them.") + del client_kwargs["user_agent"] + return pafs.S3FileSystem(**client_kwargs) + raise def _initialize_s3_fs(self) -> FileSystem: self._configure_s3_checksums() diff --git a/paimon-python/pypaimon/tests/oss_user_agent_test.py b/paimon-python/pypaimon/tests/oss_user_agent_test.py index dee7fd822ce3..8f0d00628392 100644 --- a/paimon-python/pypaimon/tests/oss_user_agent_test.py +++ b/paimon-python/pypaimon/tests/oss_user_agent_test.py @@ -121,7 +121,7 @@ def test_blank_extended_is_ignored(self): class LegacyOssUserAgentTest(unittest.TestCase): - def _new_legacy_file_io(self): + def _new_legacy_file_io(self, extra=None): options = Options({ OssOptions.OSS_ACCESS_KEY_ID.key(): "ak", OssOptions.OSS_ACCESS_KEY_SECRET.key(): "sk", @@ -129,10 +129,61 @@ def _new_legacy_file_io(self): OssOptions.OSS_REGION.key(): "cn-test", OssOptions.OSS_IMPL.key(): "legacy", }) + options.to_map().update(extra or {}) with mock.patch.object(pyarrow_file_io.pafs, "S3FileSystem") as s3_filesystem, \ mock.patch.object(oss_user_agent, "identity", return_value=IDENTITY): PyArrowFileIO("oss://test-bucket/", options) s3_filesystem.assert_called_once() + return s3_filesystem.call_args[1] + + def test_custom_user_agent_is_per_filesystem(self): + with mock.patch.dict(os.environ, {"AWS_SDK_UA_APP_ID": "existing-app"}): + first = self._new_legacy_file_io({ + "user-agent.module": "first/1.0", "user-agent.features": "training batch", + "user-agent.extended": "env=dev"}) + second = self._new_legacy_file_io({ + "user-agent.module": "second/2.0", "user-agent.features": "serving"}) + default = self._new_legacy_file_io() + transport = "pyarrow/" + pyarrow_file_io.pyarrow.__version__ + self.assertEqual("first/1.0({};{};training;batch) env=dev".format(transport, IDENTITY), + first["user_agent"]) + self.assertEqual("second/2.0({};{};serving)".format(transport, IDENTITY), + second["user_agent"]) + self.assertNotIn("user_agent", default) + self.assertEqual("existing-app", os.environ["AWS_SDK_UA_APP_ID"]) + + def test_oss_overrides_and_tracking_reach_pyarrow(self): + kwargs = self._new_legacy_file_io({ + "user-agent.module": "common", "user-agent.features": "ignored", + "user-agent.extended": "env=dev", "fs.oss.user.agent.module": "oss-app", + "fs.oss.user.agent.features": "read", "dlf.access-tracking.extended-info": "uid/123"}) + self.assertEqual("oss-app(pyarrow/{};{};read) env=dev uid/123".format( + pyarrow_file_io.pyarrow.__version__, IDENTITY), kwargs["user_agent"]) + + def test_default_module_is_not_duplicated(self): + kwargs = self._new_legacy_file_io({"user-agent.features": "training"}) + self.assertEqual("{}(pyarrow/{};training)".format( + IDENTITY, pyarrow_file_io.pyarrow.__version__), kwargs["user_agent"]) + + def test_unsupported_pyarrow_warns_and_preserves_existing_behavior(self): + options = Options({"fs.oss.impl": "legacy", "user-agent.module": "app"}) + error = TypeError("__init__() got an unexpected keyword argument 'user_agent'") + filesystem = mock.sentinel.filesystem + with mock.patch.object(pyarrow_file_io.pafs, "S3FileSystem", + side_effect=[error, filesystem]) as create: + with self.assertLogs(level="WARNING") as logs: + file_io = PyArrowFileIO("oss://test-bucket/", options) + self.assertIs(file_io.filesystem, filesystem) + self.assertIn("custom OSS User-Agent options are ignored", " ".join(logs.output)) + self.assertIn("user_agent", create.call_args_list[0][1]) + self.assertNotIn("user_agent", create.call_args_list[1][1]) + + def test_other_constructor_errors_are_not_hidden(self): + options = Options({"fs.oss.impl": "legacy", "user-agent.module": "app"}) + error = TypeError("invalid endpoint") + with mock.patch.object(pyarrow_file_io.pafs, "S3FileSystem", side_effect=error): + with self.assertRaisesRegex(TypeError, "invalid endpoint"): + PyArrowFileIO("oss://test-bucket/", options) def test_sets_aws_app_id_when_absent(self): with mock.patch.dict(os.environ):