Skip to content

Import botocore lazily for AWS_MSK_IAM #37

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dogukancagatay wants to merge 1 commit into from
Closed

Import botocore lazily for AWS_MSK_IAM #37

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 6 additions & 8 deletions kafka/sasl/msk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,6 @@
import logging
import string

# needed for AWS_MSK_IAM authentication:
try:
from botocore.session import Session as BotoSession
except ImportError:
# no botocore available, will disable AWS_MSK_IAM mechanism
BotoSession = None

from kafka.errors import KafkaConfigurationError
from kafka.sasl.abc import SaslMechanism
from kafka.vendor.six.moves import urllib
Expand All @@ -24,7 +17,12 @@

class SaslMechanismAwsMskIam(SaslMechanism):
def __init__(self, **config):
try:
from botocore.session import Session as BotoSession
except ImportError:
BotoSession = None
assert BotoSession is not None, 'AWS_MSK_IAM requires the "botocore" package'
self._boto_session_cls = BotoSession
assert config.get('security_protocol', '') == 'SASL_SSL', 'AWS_MSK_IAM requires SASL_SSL'
assert 'host' in config, 'AWS_MSK_IAM requires host configuration'
self.host = config['host']
Expand All @@ -33,7 +31,7 @@ def __init__(self, **config):
self._is_authenticated = False

def _build_client(self):
session = BotoSession()
session = self._boto_session_cls()
credentials = session.get_credentials().get_frozen_credentials()
if not session.get_config_variable('region'):
raise KafkaConfigurationError('Unable to determine region for AWS MSK cluster. Is AWS_DEFAULT_REGION set?')
Expand Down
2 changes: 1 addition & 1 deletion test/sasl/test_msk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ def test_aws_msk_iam_client_temporary_credentials():


def test_aws_msk_iam_sasl_mechanism():
with mock.patch('kafka.sasl.msk.BotoSession'):
with mock.patch.dict(sys.modules, {'botocore': mock.MagicMock(), 'botocore.session': mock.MagicMock()}):
sasl = SaslMechanismAwsMskIam(security_protocol='SASL_SSL', host='localhost')
with mock.patch.object(sasl, '_build_client', return_value=client_factory(token=None)):
assert sasl.auth_bytes() != b''
Expand Down