Merge pull request jantman#240 from jantman/issues/239

Fixes jantman#239 - add capability to refresh TA checks

CHANGES.rst

@@ -7,6 +7,11 @@ Changelog
 This release deprecates support for Python 3.2. It will be removed in the
 next release.
 
+This release introduces support for automatically refreshing Trusted Advisor
+checks on accounts that support this. If you use this new feature,
+awslimitchecker will require a new permission, ``trustedadvisor:RefreshCheck``.
+See `Getting Started - Trusted Advisor <http://awslimitchecker.readthedocs.io/en/latest/getting_started.html#trusted-advisor>`_ for further information.
+
 * `#231 <https://github.com/jantman/awslimitchecker/issues/231>`_ - add support
   for new f1, r4 and t2.(xlarge|2xlarge) instance types, introduced in November
   2016.
@@ -19,6 +24,12 @@ next release.
   Python 3.2 support. There don't appear to have been any downloads on py32
   in the last 6 months, and the effort to support it is too high.
 * A bunch of Sphinx work to use README.rst in the generated documentation.
+* Changed DEBUG-level logging format to include timestamp.
+* `#239 <https://github.com/jantman/awslimitchecker/issues/239>`_ - Support
+  refreshing Trusted Advisor check results during the run, and optionally waiting
+  for refresh to finish. See
+  `Getting Started - Trusted Advisor <http://awslimitchecker.readthedocs.io/en/latest/getting_started.html#trusted-advisor>`_
+  for further information.
 
 0.6.0 (2016-11-12)
 ------------------

README.rst

@@ -93,6 +93,11 @@ What It Does
 - where possible, pull current limits from each service's API (for services that provide this information)
 - Supports explicitly setting the AWS region
 - Supports using `STS <http://docs.aws.amazon.com/STS/latest/APIReference/Welcome.html>`_ to assume roles in other accounts, including using ``external_id``.
+- Optionally refresh Trusted Advisor "Service Limits" check before polling
+  Trusted Advisor data, and optionally wait for the refresh to complete (up to
+  an optional maximum time limit). See
+  `Getting Started - Trusted Advisor <http://awslimitchecker.readthedocs.io/en/latest/getting_started.html#trusted-advisor>`_
+  for more information.
 
 Requirements
 ------------

awslimitchecker/checker.py

@@ -52,7 +52,7 @@ class AwsLimitChecker(object):
     def __init__(self, warning_threshold=80, critical_threshold=99,
                  profile_name=None, account_id=None, account_role=None,
                  region=None, external_id=None, mfa_serial_number=None,
-                 mfa_token=None):
+                 mfa_token=None, ta_refresh_mode=None, ta_refresh_timeout=None):
         """
         Main AwsLimitChecker class - this should be the only externally-used
         portion of awslimitchecker.
@@ -96,6 +96,23 @@ def __init__(self, warning_threshold=80, critical_threshold=99,
         :param mfa_token: (optional) the `MFA Token` string to use when
           assuming a role via STS.
         :type mfa_token: str
+        :param ta_refresh_mode: How to handle refreshing Trusted Advisor checks;
+          this is either None (do not refresh at all), the string "wait"
+          (trigger refresh of all limit-related checks and wait for the refresh
+          to complete), the string "trigger" (trigger refresh of all
+          limit-related checks but do not wait for the refresh to complete), or
+          an integer, which causes any limit-related checks more than this
+          number of seconds old to be refreshed, waiting for the refresh to
+          complete. Note that "trigger" will likely result in the current run
+          getting stale data, but the check being refreshed in time for the
+          next run.
+        :type ta_refresh_mode: :py:class:`str` or :py:class:`int` or
+          :py:data:`None`
+        :param ta_refresh_timeout: If ``ta_refresh_mode`` is "wait" or an
+          integer (any mode that will wait for the refresh to complete), if this
+          parameter is not None, only wait up to this number of seconds for the
+          refresh to finish before continuing on anyway.
+        :type ta_refresh_timeout: :py:class:`int` or :py:data:`None`
         """
         # ###### IMPORTANT license notice ##########
         # Pursuant to Sections 5(b) and 13 of the GNU Affero General Public
@@ -151,6 +168,8 @@ def __init__(self, warning_threshold=80, critical_threshold=99,
             external_id=external_id,
             mfa_serial_number=mfa_serial_number,
             mfa_token=mfa_token,
+            ta_refresh_mode=ta_refresh_mode,
+            ta_refresh_timeout=ta_refresh_timeout
         )
 
     def get_version(self):
@@ -450,6 +469,7 @@ def get_required_iam_policy(self):
         required_actions = [
             'support:*',
             'trustedadvisor:Describe*',
+            'trustedadvisor:RefreshCheck'
         ]
         for cls in self.services.values():
             required_actions.extend(cls.required_iam_permissions())

awslimitchecker/runner.py

@@ -162,6 +162,30 @@ def parse_args(self, argv):
         p.add_argument('--skip-ta', action='store_true', default=False,
                        help='do not attempt to pull *any* information on limits'
                        ' from Trusted Advisor')
+        g = p.add_mutually_exclusive_group()
+        g.add_argument('--ta-refresh-wait', dest='ta_refresh_wait',
+                       action='store_true', default=False,
+                       help='If applicable, refresh all Trusted Advisor '
+                            'limit-related checks, and wait for the refresh to'
+                            ' complete before continuing.')
+        g.add_argument('--ta-refresh-trigger', dest='ta_refresh_trigger',
+                       action='store_true', default=False,
+                       help='If applicable, trigger refreshes for all Trusted '
+                            'Advisor limit-related checks, but do not wait for '
+                            'them to finish refreshing; trigger the refresh '
+                            'and continue on (useful to ensure checks are '
+                            'refreshed before the next scheduled run).')
+        g.add_argument('--ta-refresh-older', dest='ta_refresh_older',
+                       action='store', type=int, default=None,
+                       help='If applicable, trigger refreshes for all Trusted '
+                            'Advisor limit-related checks with results more '
+                            'than this number of seconds old. Wait for the '
+                            'refresh to complete before continuing.')
+        p.add_argument('--ta-refresh-timeout', dest='ta_refresh_timeout',
+                       type=int, action='store', default=None,
+                       help='If waiting for TA checks to refresh, wait up to '
+                            'this number of seconds before continuing on '
+                            'anyway.')
         p.add_argument('--no-color', action='store_true', default=False,
                        help='do not colorize output')
         p.add_argument('-v', '--verbose', dest='verbose', action='count',
@@ -172,6 +196,13 @@ def parse_args(self, argv):
                        default=False,
                        help='print version number and exit.')
         args = p.parse_args(argv)
+        args.ta_refresh_mode = None
+        if args.ta_refresh_wait:
+            args.ta_refresh_mode = 'wait'
+        elif args.ta_refresh_trigger:
+            args.ta_refresh_mode = 'trigger'
+        elif args.ta_refresh_older is not None:
+            args.ta_refresh_mode = args.ta_refresh_older
         return args
 
     def list_services(self):
@@ -301,7 +332,7 @@ def console_entry_point(self):
             logger.setLevel(logging.INFO)
         elif args.verbose > 1:
             # debug-level logging hacks
-            FORMAT = "[%(levelname)s %(filename)s:%(lineno)s - " \
+            FORMAT = "%(asctime)s [%(levelname)s %(filename)s:%(lineno)s - " \
                      "%(name)s.%(funcName)s() ] %(message)s"
             debug_formatter = logging.Formatter(fmt=FORMAT)
             logger.handlers[0].setFormatter(debug_formatter)
@@ -323,7 +354,9 @@ def console_entry_point(self):
             region=args.region,
             external_id=args.external_id,
             mfa_serial_number=args.mfa_serial_number,
-            mfa_token=args.mfa_token
+            mfa_token=args.mfa_token,
+            ta_refresh_mode=args.ta_refresh_mode,
+            ta_refresh_timeout=args.ta_refresh_timeout
         )
 
         if args.version:

awslimitchecker/tests/test_checker.py

@@ -120,7 +120,8 @@ def test_init(self):
         assert self.mock_ta_constr.mock_calls == [
             call(services, account_id=None, account_role=None, region=None,
                  external_id=None, mfa_serial_number=None, mfa_token=None,
-                 profile_name=None)
+                 profile_name=None, ta_refresh_mode=None,
+                 ta_refresh_timeout=None)
         ]
         assert self.mock_svc1.mock_calls == []
         assert self.mock_svc2.mock_calls == []
@@ -187,7 +188,8 @@ def test_init_thresholds(self):
         assert mock_ta_constr.mock_calls == [
             call(services, account_id=None, account_role=None, region=None,
                  external_id=None, mfa_serial_number=None, mfa_token=None,
-                 profile_name=None)
+                 profile_name=None, ta_refresh_mode=None,
+                 ta_refresh_timeout=None)
         ]
         assert mock_svc1.mock_calls == []
         assert mock_svc2.mock_calls == []
@@ -234,7 +236,8 @@ def test_init_region_profile(self):
         assert mock_ta_constr.mock_calls == [
             call(services, profile_name='foo', account_id=None,
                  account_role=None, region='myregion', external_id=None,
-                 mfa_serial_number=None, mfa_token=None)
+                 mfa_serial_number=None, mfa_token=None, ta_refresh_mode=None,
+                 ta_refresh_timeout=None)
         ]
         assert mock_svc1.mock_calls == []
         assert mock_svc2.mock_calls == []
@@ -292,15 +295,17 @@ def test_init_sts(self):
                 external_id=None,
                 mfa_serial_number=None,
                 mfa_token=None,
-                profile_name=None
+                profile_name=None,
+                ta_refresh_mode=None,
+                ta_refresh_timeout=None
             )
         ]
         assert mock_svc1.mock_calls == []
         assert mock_svc2.mock_calls == []
         assert self.mock_version.mock_calls == [call()]
         assert self.cls.vinfo == self.mock_ver_info
 
-    def test_init_sts_external_id(self):
+    def test_init_sts_external_id_ta_refresh(self):
         mock_svc1 = Mock(spec_set=_AwsService)
         mock_svc2 = Mock(spec_set=_AwsService)
         mock_foo = Mock(spec_set=_AwsService)
@@ -328,7 +333,9 @@ def test_init_sts_external_id(self):
                     region='myregion',
                     external_id='myextid',
                     mfa_serial_number=None,
-                    mfa_token=None
+                    mfa_token=None,
+                    ta_refresh_mode=123,
+                    ta_refresh_timeout=456
                 )
         # dict should be of _AwsService instances
         services = {
@@ -354,7 +361,9 @@ def test_init_sts_external_id(self):
                 external_id='myextid',
                 mfa_serial_number=None,
                 mfa_token=None,
-                profile_name=None
+                profile_name=None,
+                ta_refresh_mode=123,
+                ta_refresh_timeout=456
             )
         ]
         assert mock_svc1.mock_calls == []
@@ -657,7 +666,8 @@ def test_get_required_iam_policy(self):
                     'foo:perm1',
                     'foo:perm2',
                     'support:*',
-                    'trustedadvisor:Describe*'
+                    'trustedadvisor:Describe*',
+                    'trustedadvisor:RefreshCheck'
                 ],
             }],
         }