new APIs to avoid exceptions

Author: JoshuaFox

.gcloudignore

@@ -1,6 +1,6 @@
 *.bak
 .DS_Store
-
+*-test.txt
 test_scripts/
 test_*
 sample_data/

.gitignore

@@ -3,7 +3,7 @@
 TEMP*
 
 test_scripts/
-
+*-test.text
 test_*
 # Environments
 .env

config.yaml

@@ -3,12 +3,15 @@
 
 # Keys:
 
-# Change this before first deployment for added security in communication between PubSub and the Iris App on App Engine.
-# You could even re-generate a new token per deployment.
-# Note that this approach is not very secure, though it was was once recommended by Google.
-# However, so long as  your GCP project is secure, this token protects against unwanted invocations.
-# Also, the worst an invocation can do is re-trigger idempotent behavior, with possible DOS.
-pubsub_verification_token: 2a343f4c1b76512039fe763412756c4fbb30c
+# projects: Only resources in these projects will get labeled.
+# But if the value is empty, *all* projects in the organization are included.
+projects:
+  - joshua-playground
+
+# plugins: Only these plugins are enabled.
+# For example, list some of these: bigquery, bigtable, disks, buckets, cloudsql, instances, snaphots, subscriptions, topics
+# But if the value is empty, *all* plugins are enabled
+plugins: {}
 
 # iris_prefix is prefixed to the key of each label that is added.
 iris_prefix: iris
@@ -22,6 +25,9 @@ from_project: True
 # get labeled on cron.
 label_all_on_cron: True
 
-# projects: Only resources in these projects will get labeled.
-# But if the value is empty, *all* projects in the organization are included.
-projects: {}
+# Change this before first deployment for added security in communication between PubSub and the Iris App on App Engine.
+# You could even re-generate a new token per deployment.
+# Note that this approach is not very secure, though it was was once recommended by Google.
+# However, so long as  your GCP project is secure, this token protects against unwanted invocations.
+# Also, the worst an invocation can do is re-trigger idempotent behavior, with possible DOS.
+pubsub_verification_token: 2a343f4c1b76512039fe763412756c4fbb30c

config.yaml.original

@@ -3,12 +3,14 @@
 
 # Keys:
 
-# Change this before first deployment for added security in communication between PubSub and the Iris App on App Engine.
-# You could even re-generate a new token per deployment.
-# Note that this approach is not very secure, though it was was once recommended by Google.
-# However, so long as  your GCP project is secure, this token protects against unwanted invocations.
-# Also, the worst an invocation can do is re-trigger idempotent behavior, with possible DOS.
-pubsub_verification_token: 2a343f4c1b76512039fe763412756c4fbb30c
+# projects: Only resources in these projects will get labeled.
+# But if the value is empty, *all* projects in the organization are included.
+projects: {}
+
+# plugins: Only these plugins are enabled.
+# For example, list some of these: bigquery, buckets, bigtable, disks,  cloudsql, instances, snaphots, subscriptions, topics
+# But if the value is empty, *all* plugins are enabled
+plugins: {}
 
 # iris_prefix is prefixed to the key of each label that is added.
 iris_prefix: iris
@@ -22,6 +24,9 @@ from_project: True
 # get labeled on cron.
 label_all_on_cron: True
 
-# projects: Only resources in these projects will get labeled.
-# But if the value is empty, *all* projects in the organization are included.
-projects: {}
+# Change this before first deployment for added security in communication between PubSub and the Iris App on App Engine.
+# You could even re-generate a new token per deployment.
+# Note that this approach is not very secure, though it was was once recommended by Google.
+# However, so long as  your GCP project is secure, this token protects against unwanted invocations.
+# Also, the worst an invocation can do is re-trigger idempotent behavior, with possible DOS.
+pubsub_verification_token: 2a343f4c1b76512039fe763412756c4fbb30c

config.yaml.test.template

@@ -3,10 +3,11 @@
 # In integration_test.sh, config.yaml.test.template is used to make a temporary config.yaml, then config.yaml is reverted at the end.
 #
 # Generate random token per-test-run as an additional isolation measure
-pubsub_verification_token: ${PUBSUB_TEST_TOKEN}
-iris_prefix: ${RUN_ID}
-from_project: True
-label_all_on_cron: False
 projects:
   - ${DEPLOYMENT_PROJECT}
   - ${TEST_PROJECT}
+plugins: {}
+iris_prefix: ${RUN_ID}
+from_project: True
+label_all_on_cron: False
+pubsub_verification_token: ${PUBSUB_TEST_TOKEN}

gce_base/gce_zonal_base.py

@@ -4,6 +4,8 @@
 import util.gcp_utils
 from gce_base.gce_base import GceBase
 
+from google.cloud import compute_v1
+
 
 class GceZonalBase(GceBase, metaclass=ABCMeta):
     def _gcp_zone(self, gcp_object):
@@ -27,9 +29,12 @@ def _all_zones(self, project_id):
         """
         Get all available zones.
         """
-        response = self._google_client.zones().list(project=project_id).execute()
-        zones = [zone["description"] for zone in response["items"]]
-        return zones
+        zones_client = compute_v1.ZonesClient()
+
+        request = compute_v1.ListZonesRequest(project=project_id)
+        zones = zones_client.list(request)
+        ret = [z.name for z in zones]
+        return ret
 
     def block_labeling(self, gcp_object, original_labels):
         # goog-gke-node appears in Nodes and Disks; and goog-gke-volume appears in Disks

main.py

@@ -14,7 +14,7 @@
 from util import pubsub_utils, gcp_utils, utils, config_utils
 from util.gcp_utils import detect_gae
 
-from util.config_utils import iris_prefix, configured_project, get_config, pubsub_token
+from util.config_utils import iris_prefix, is_project_enabled, get_config, pubsub_token
 from util.utils import init_logging, log_time, timing
 
 import googlecloudprofiler
@@ -88,7 +88,7 @@ def schedule():
         nonappscript_projects = [p for p in all_projects if p not in appscript_projects]
 
         configured_projects = [
-            p for p in nonappscript_projects if config_utils.configured_project(p)
+            p for p in nonappscript_projects if config_utils.is_project_enabled(p)
         ]
 
         skipped_nonappscript_projects = [
@@ -126,7 +126,7 @@ def schedule():
                 raise Exception(msg)
 
         for project_id in configured_projects:
-            for _, plugin_ in Plugin.instances.items():
+            for _, plugin_ in Plugin.plugins.items():
                 if (
                     not plugin_.is_labeled_on_creation()
                     or plugin_.relabel_on_cron()
@@ -175,7 +175,7 @@ def label_one():
         method_from_log = data["protoPayload"]["methodName"]
 
         plugins_found = []
-        for plugin_name, plugin in Plugin.instances.items():
+        for plugin_name, plugin in Plugin.plugins.items():
             for supported_method in plugin.method_names():
                 if supported_method.lower() in method_from_log.lower():
                     if plugin.is_labeled_on_creation():
@@ -185,8 +185,19 @@ def label_one():
                         plugin_name
                     )  # Append it even if not used due to is_labeled_on_creation False
 
-        if len(plugins_found) != 1:
-            logging.error(f"Error: plugins found {plugins_found} for {method_from_log}")
+        if not plugins_found:
+            logging.info(
+                "(OK if plugin is disabled.) No plugins found for %s. Enabled plugins are %s",
+                method_from_log,
+                config_utils.enabled_plugins(),
+            )
+
+        if len(plugins_found) > 1:
+            raise Exception(
+                f"Error: Multiple plugins found %s for %s"
+                % (plugins_found, method_from_log)
+            )
+
         return "OK", 200
     except Exception as e:
         project_id = data.get("resource", {}).get("labels", {}).get("project_id")
@@ -200,7 +211,7 @@ def __label_one_0(data, plugin):
     gcp_object = plugin.get_gcp_object(data)
     if gcp_object is not None:
         project_id = data["resource"]["labels"]["project_id"]
-        if configured_project(project_id):
+        if is_project_enabled(project_id):
             logging.info("Will label_one() in %s, %s", project_id, gcp_object)
             plugin.label_resource(gcp_object, project_id)
             plugin.do_batch()
@@ -254,13 +265,20 @@ def do_label():
         plugin_class_name = data["plugin"]
 
         plugin = Plugin.get_plugin(plugin_class_name)
-        project_id = data["project_id"]
-        with timing(f"do_label {plugin_class_name} {project_id}"):
+        if not plugin:
             logging.info(
-                "do_label() for %s in %s", plugin.__class__.__name__, project_id
+                "(OK if plugin is disabled.) No plugins found for %s. Enabled plugins are %s",
+                plugin_class_name,
+                config_utils.enabled_plugins(),
             )
-            plugin.label_all(project_id)
-        logging.info("OK on do_label %s %s", plugin_class_name, project_id)
+        else:
+            project_id = data["project_id"]
+            with timing(f"do_label {plugin_class_name} {project_id}"):
+                logging.info(
+                    "do_label() for %s in %s", plugin.__class__.__name__, project_id
+                )
+                plugin.label_all(project_id)
+            logging.info("OK on do_label %s %s", plugin_class_name, project_id)
         return "OK", 200
     except Exception as e:
         logging.exception(

plugin.py

@@ -7,7 +7,7 @@
 from googleapiclient import discovery
 from googleapiclient import errors
 
-from util import gcp_utils
+from util import gcp_utils, config_utils
 from util.config_utils import is_copying_labels_from_project, iris_prefix
 from util.utils import methods, cls_by_name, log_time, timed_lru_cache
 
@@ -22,8 +22,8 @@ class Plugin(object, metaclass=ABCMeta):
 
     # For a class to know its subclasses and their instances is generally bad.
     # We could create a separate PluginManager but let's not get too Java-ish.
-    instances: typing.Dict[str, "Plugin"]
-    instances = {}
+    plugins: typing.Dict[str, "Plugin"]
+    plugins = {}
 
     def __init__(self):
         self._google_client = discovery.build(*self.discovery_api())
@@ -139,15 +139,16 @@ def load_plugin_class(name) -> typing.Type:
             return plugin_cls
 
         for _, module, _ in pkgutil.iter_modules([PLUGINS_MODULE]):
-            plugin_class = load_plugin_class(module)
-            instance = plugin_class()
-            Plugin.instances[plugin_class.__name__] = instance
+            if config_utils.is_plugin_enabled(module):
+                plugin_class = load_plugin_class(module)
+                instance = plugin_class()
+                Plugin.plugins[plugin_class.__name__] = instance
 
-        assert Plugin.instances, "No plugins defined"
+        assert Plugin.plugins, "No plugins defined"
 
     @staticmethod
     def get_plugin(plugin_name: str) -> "Plugin":
-        return Plugin.instances[plugin_name]
+        return Plugin.plugins.get(plugin_name)
 
     def _build_labels(self, gcp_object, project_id):
         """

plugins/instances.py

@@ -1,5 +1,6 @@
 import logging
 
+from google.cloud import compute_v1
 from googleapiclient import errors
 
 from gce_base.gce_zonal_base import GceZonalBase
@@ -24,27 +25,42 @@ def method_names(self):
         return ["compute.instances.insert", "compute.instances.start"]
 
     def __list_instances(self, project_id, zone):
-        instances = []
-        page_token = None
-        more_results = True
-        while more_results:
-            result = (
-                self._google_client.instances()
-                .list(
-                    project=project_id,
-                    zone=zone,
-                    pageToken=page_token,
+        def list1():
+            instances1 = []
+            page_token = None
+            more_results = True
+            while more_results:
+                result = (
+                    self._google_client.instances()
+                    .list(
+                        project=project_id,
+                        zone=zone,
+                        pageToken=page_token,
+                    )
+                    .execute()
                 )
-                .execute()
+                if "items" in result:
+                    instances = instances1 + result["items"]
+                if "nextPageToken" in result:
+                    page_token = result["nextPageToken"]
+                else:
+                    more_results = False
+            return instances1
+
+        ret1 = list1()
+
+        def list2():
+            instances_client = compute_v1.InstancesClient()
+
+            request = compute_v1.ListInstancesRequest(
+                project=project_id, zone="us-central1-a"
             )
-            if "items" in result:
-                instances = instances + result["items"]
-            if "nextPageToken" in result:
-                page_token = result["nextPageToken"]
-            else:
-                more_results = False
-
-        return instances
+            instances2 = instances_client.list(request)
+
+            return list(instances2)
+
+        ret2 = list2()
+        return ret1
 
     def __get_instance(self, project_id, zone, name):
         try:

requirements.txt

@@ -4,7 +4,8 @@ google-cloud-pubsub==2.8.0
 oauth2client==4.1.3
 
 ratelimit==2.2.1
-#google-cloud-resource-manager==0.30.3
+google-cloud-compute==0.6.0
+google-cloud-resource-manager==1.3.0
 black
 google-cloud-profiler
 # gunicorn is only needed if you add entrypoint: gunicorn -b :$PORT --worker-class=sync --workers=1 main:app

test_scripts/test_do_label.py

@@ -3,6 +3,7 @@
 import sys
 
 from test_scripts.utils_for_tests import do_local_http
+from util import config_utils
 
 """
 This is a debugging tool useful in development.
@@ -61,7 +62,8 @@ def main():
         chosen_plugins = [s.strip() for s in chosen_plugins]
         unsupported = [p for p in chosen_plugins if p not in PLUGINS]
         if unsupported:
-            raise Exception(f"Unsupported: {', '.join(unsupported)}")
+            raise Exception(f"Error: \"{', '.join(unsupported)}\" not a legal value. "
+                            f"For this test, you can use these (comma-separated) in the env variable: {PLUGINS}")
     print(f"Will do_label on{msg}: {', '.join(chosen_plugins)} ")
     test_do_label(chosen_plugins)