CloudTrail customization fix to use built-in session

awscli/customizations/cloudtrail.py

@@ -15,7 +15,7 @@
 import sys
 
 from awscli.customizations.commands import BasicCommand
-from awscli.customizations.service import Service, OperationException
+from awscli.customizations.service import Service
 from botocore.vendored import requests
 
 
@@ -85,20 +85,18 @@ def _run_main(self, args, parsed_globals):
 
         # Initialize services
         LOG.debug('Initializing S3, SNS and CloudTrail...')
-        self.iam = Service('iam')
-        self.s3 = Service('s3', endpoint_args['region_name'])
-        self.sns = Service('sns', endpoint_args['region_name'])
-        self.cloudtrail = Service('cloudtrail', endpoint_args=endpoint_args)
+        self.iam = Service('iam', session=self._session)
+        self.s3 = Service('s3', endpoint_args['region_name'],
+                          session=self._session)
+        self.sns = Service('sns', endpoint_args['region_name'],
+                           session=self._session)
+        self.cloudtrail = Service('cloudtrail', endpoint_args=endpoint_args,
+                                  session=self._session)
 
-        # Run the command, watch for errors and report success/fail
-        success = True
-        try:
-            self._call(args, parsed_globals)
-        except OperationException as err:
-            success = False
-            self.handle_error(err)
+        # Run the command and report success
+        self._call(args, parsed_globals)
 
-        return int(success)
+        return 1
 
     def _call(self, options, parsed_globals):
         """
@@ -177,39 +175,6 @@ def _call(self, options, parsed_globals):
                 'Logs will be delivered to {bucket}:{prefix}\n'.format(
                     bucket=bucket, prefix=options.s3_prefix))
 
-    def handle_error(self, err):
-        """
-        Handle service errors by printing out exception information. Parses
-        error codes and messages and attempts to display them in a nice
-        format.
-        """
-        sys.stderr.write('Error(s) encountered:\n')
-
-        if err.data and 'Errors' in err.data:
-            try:
-                sys.stderr.write('\t{code}: {msg}\n'.format(
-                    code=err.data['ErrorCode'],
-                    msg=err.data['Message']
-                ))
-            except Exception:
-                pass
-
-            for error in err.data['Errors']:
-                msg = error.get('Message', 'No details')
-
-                code = 'Unknown Error'
-                for name in ['ErrorCode', 'Code', 'Type']:
-                    if name in error:
-                        code = error[name]
-                        break
-
-                sys.stderr.write('\t{code}: {msg}\n'.format(code=code, msg=msg))
-        elif err.data and 'Body' in err.data and \
-             hasattr(err.data['Body'], 'read'):
-            sys.stderr.write('\t{err}\n'.format(err=err.data['Body'].read()))
-        else:
-            sys.stderr.write('\t{err}\n'.format(err=err.data))
-
     def setup_new_bucket(self, bucket, prefix, policy_url=None):
         """
         Creates a new S3 bucket with an appropriate policy to let CloudTrail

awscli/customizations/service.py

@@ -12,20 +12,6 @@
 # language governing permissions and limitations under the License.
 import botocore.session
 
-class OperationException(Exception):
-    """
-    An exception to be thrown when service operations fail. The exception
-    contains a link to the operation, the arguments that were used to
-    invoke the operation, the HTTP response object and response data.
-    """
-    def __init__(self, operation, args=None, res=None, data=None):
-        super(OperationException, self).__init__()
-
-        self.operation = operation
-        self.args = args
-        self.res = res
-        self.data = data
-
 
 class OperationProxy(object):
     """
@@ -49,9 +35,6 @@ def __init__(self, service, name, endpoint):
     def __call__(self, **kwargs):
         res, data = self._operation.call(self.endpoint, **kwargs)
 
-        if res.status_code < 200 or res.status_code > 300:
-            raise OperationException(self, kwargs, res, data)
-
         return data
 
 
@@ -76,7 +59,10 @@ class Service(object):
 
        >>> s3 = Service('s3', {'endpoint_url': 'http://...'})
 
-    A custom session can be used:
+    Note that by default, a new session is used. When creating instances
+    from a custom command you should reuse the default CLI session, which
+    has extra error handling and profile handling enabled.A custom session
+    can be used by passing it to the constructor:
 
         >>> s3 = Service('s3', session=mysession)
 

tests/unit/customizations/test_cloudtrail.py

@@ -11,12 +11,15 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
+import botocore.session
 import io
+import os
 
+from awscli.clidriver import create_clidriver
 from awscli.customizations.cloudtrail import CloudTrailSubscribe
 from awscli.customizations.service import Service
-from mock import Mock
-from tests.unit.customizations.test_configure import FakeSession
+from mock import Mock, patch
+from tests.unit.test_clidriver import FakeSession
 from tests import unittest
 
 
@@ -104,3 +107,40 @@ def test_sns_create(self):
     def test_sns_create_already_exists(self):
         with self.assertRaises(Exception):
             self.subscribe.setup_new_topic('test2')
+
+
+class TestCloudTrailSessions(unittest.TestCase):
+    def test_sessions(self):
+        """
+        Make sure that the session passed to our custom command
+        is the same session used when making service calls.
+        """
+        # awscli/__init__.py injects AWS_DATA_PATH at import time
+        # so that we can find cli.json.  This might be fixed in the
+        # future, but for now we just grab that value out of the real
+        # os.environ so the patched os.environ has this data and
+        # the CLI works.
+        self.environ = {
+            'AWS_DATA_PATH': os.environ['AWS_DATA_PATH'],
+            'AWS_DEFAULT_REGION': 'us-east-1',
+            'AWS_ACCESS_KEY_ID': 'access_key',
+            'AWS_SECRET_ACCESS_KEY': 'secret_key',
+        }
+        self.environ_patch = patch('os.environ', self.environ)
+        self.environ_patch.start()
+
+        # Get a new session we will use to test
+        driver = create_clidriver()
+
+        def _mock_call(self, *args, **kwargs):
+            # Make sure every service uses the same session
+            assert driver.session == self.iam.session
+            assert driver.session == self.s3.session
+            assert driver.session == self.sns.session
+            assert driver.session == self.cloudtrail.session
+
+        with patch.object(CloudTrailSubscribe, '_call', _mock_call):
+            rc = driver.main('cloudtrail create-subscription --name test --s3-use-bucket test'.split())
+            # If any assert above fails, then rc will be set to a
+            # non-zero value and this will cause the test to fail
+            self.assertEqual(rc, None)