From 95a39f759abe8d7c26535d8a2da70d456e882443 Mon Sep 17 00:00:00 2001
From: Iain Learmonth <irl@sr2.uk>
Date: Mon, 25 Apr 2022 15:05:28 +0100
Subject: [PATCH] azure_cdn: use subgroups, attempt 2

---
 app/terraform/proxy/__init__.py  | 16 ++++++++++++++++
 app/terraform/proxy/azure_cdn.py | 19 ++-----------------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/app/terraform/proxy/__init__.py b/app/terraform/proxy/__init__.py
index a26ab6c..3deed73 100644
--- a/app/terraform/proxy/__init__.py
+++ b/app/terraform/proxy/__init__.py
@@ -1,4 +1,5 @@
 import datetime
+from collections import defaultdict
 
 from app import app
 from app.extensions import db
@@ -8,6 +9,20 @@ from app.terraform import BaseAutomation
 
 
 class ProxyAutomation(BaseAutomation):
+    def get_subgroups(self):
+        conn = db.engine.connect()
+        result = conn.execute("""
+        SELECT origin.group_id, proxy.psg, COUNT(proxy.id) FROM proxy, origin
+        WHERE proxy.origin_id = origin.id
+            AND proxy.destroyed IS NULL
+            AND proxy.provider = :provider
+        GROUP BY origin.group_id, proxy.psg;
+        """, provider=self.provider)
+        subgroups = defaultdict(lambda: defaultdict(lambda: 0))
+        for row in result:
+            subgroups[row[0]][row[1]] = row[2]
+        return subgroups
+
     def create_missing_proxies(self):
         origins = Origin.query.all()
         for origin in origins:
@@ -44,6 +59,7 @@ class ProxyAutomation(BaseAutomation):
                 Proxy.provider == self.provider,
                 Proxy.destroyed == None
             ).all(),
+            subgroups = self.get_subgroups(),
             global_namespace=app.config['GLOBAL_NAMESPACE'],
             bypass_token=app.config['BYPASS_TOKEN'],
             **{
diff --git a/app/terraform/proxy/azure_cdn.py b/app/terraform/proxy/azure_cdn.py
index ed2d139..68ccdaf 100644
--- a/app/terraform/proxy/azure_cdn.py
+++ b/app/terraform/proxy/azure_cdn.py
@@ -1,7 +1,6 @@
 import datetime
 import string
 import random
-from collections import defaultdict
 
 from azure.identity import ClientSecretCredential
 from azure.mgmt.alertsmanagement import AlertsManagementClient
@@ -11,7 +10,7 @@ from app import app
 from app.alarms import get_proxy_alarm
 from app.extensions import db
 from app.models.base import Group
-from app.models.mirrors import Proxy, Origin
+from app.models.mirrors import Proxy
 from app.models.alarms import AlarmState
 from app.terraform.proxy import ProxyAutomation
 
@@ -168,23 +167,9 @@ class ProxyAzureCdnAutomation(ProxyAutomation):
     {% endfor %}
     """
 
-    def get_subgroup_counts(self):
-        conn = db.engine.connect()
-        result = conn.execute("""
-        SELECT origin.group_id, proxy.psg, COUNT(proxy.id) FROM proxy, origin
-        WHERE proxy.origin_id = origin.id
-            AND proxy.destroyed IS NULL
-            AND proxy.provider = 'azure_cdn'
-        GROUP BY origin.group_id, proxy.psg;
-        """)
-        subgroups = defaultdict(lambda: defaultdict(lambda: 0))
-        for row in result:
-            subgroups[row[0]][row[1]] = row[2]
-        return subgroups
-
     def create_missing_proxies(self):
         groups = Group.query.all()
-        subgroups = self.get_subgroup_counts()
+        subgroups = self.get_subgroups()
         for group in groups:
             subgroup = 0
             for origin in group.origins: