diff --git a/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py b/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py index aee24ac2c052..c67ea3202465 100644 --- a/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py +++ b/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py @@ -17,6 +17,7 @@ # pytype: skip-file +from apache_beam import version as beam_version from apache_beam.io.aws.clients.s3 import messages from apache_beam.options import pipeline_options from apache_beam.utils import retry @@ -25,6 +26,7 @@ # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=ungrouped-imports import boto3 + from botocore.config import Config except ImportError: boto3 = None @@ -72,7 +74,9 @@ def __init__(self, options): endpoint_url=endpoint_url, aws_access_key_id=access_key_id, aws_secret_access_key=secret_access_key, - aws_session_token=session_token) + aws_session_token=session_token, + config=Config( + user_agent_extra='apache-beam/%s' % beam_version.__version__)) self._download_request = None self._download_stream = None diff --git a/sdks/python/apache_beam/io/aws/clients/s3/client_test.py b/sdks/python/apache_beam/io/aws/clients/s3/client_test.py index 67797b1736b9..3be216e66300 100644 --- a/sdks/python/apache_beam/io/aws/clients/s3/client_test.py +++ b/sdks/python/apache_beam/io/aws/clients/s3/client_test.py @@ -20,11 +20,20 @@ import os import unittest +from apache_beam import version as beam_version from apache_beam.io.aws import s3io from apache_beam.io.aws.clients.s3 import fake_client from apache_beam.io.aws.clients.s3 import messages from apache_beam.options import pipeline_options +# Protect against environments where boto3 library is not available. +# pylint: disable=wrong-import-order, wrong-import-position +try: + from apache_beam.io.aws.clients.s3 import boto3_client +except ImportError: + boto3_client = None # type: ignore[assignment] +# pylint: enable=wrong-import-order, wrong-import-position + class ClientErrorTest(unittest.TestCase): def setUp(self): @@ -222,6 +231,15 @@ def test_complete_multipart_upload_too_many(self): self.assertEqual(e.code, 400) +@unittest.skipIf(boto3_client is None, 'AWS dependencies are not installed') +class ClientUserAgentTest(unittest.TestCase): + def test_user_agent_contains_beam_version(self): + client = boto3_client.Client(pipeline_options.S3Options()) + self.assertIn( + 'apache-beam/%s' % beam_version.__version__, + client.client.meta.config.user_agent) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main()