[batch] Add feature switches (#13262)

Let me know if you think this is good and whether I need to test the UI
with dev deploy.

---------

Co-authored-by: Dan King <daniel.zidan.king@gmail.com>

batch/batch/driver/main.py

@@ -38,7 +38,14 @@
 from hailtop import aiotools, httpx
 from hailtop.config import get_deploy_config
 from hailtop.hail_logging import AccessLogger
-from hailtop.utils import AsyncWorkerPool, Notice, dump_all_stacktraces, flatten, periodically_call, time_msecs
+from hailtop.utils import (
+    AsyncWorkerPool,
+    Notice,
+    dump_all_stacktraces,
+    flatten,
+    periodically_call,
+    time_msecs,
+)
 from web_common import render_template, set_message, setup_aiohttp_jinja2, setup_common_static_routes
 
 from ..batch import cancel_batch_in_db
@@ -459,6 +466,7 @@ async def get_index(request, userdata):
         'total_provisioned_cores_mcpu': inst_coll_manager.global_total_provisioned_cores_mcpu,
         'live_free_cores_mcpu': inst_coll_manager.global_current_version_live_free_cores_mcpu,
         'frozen': app['frozen'],
+        'feature_flags': app['feature_flags'],
     }
     return await render_template('batch-driver', request, userdata, 'index.html', page_context)
 
@@ -552,6 +560,27 @@ def validate_int(session, name, value, predicate, description):
     return validate(session, name, i, predicate, description)
 
 
+@routes.post('/configure-feature-flags')
+@check_csrf_token
+@auth.web_authenticated_developers_only()
+async def configure_feature_flags(request, userdata):  # pylint: disable=unused-argument
+    app = request.app
+    db: Database = app['db']
+    post = await request.post()
+
+    compact_billing_tables = 'compact_billing_tables' in post
+
+    await db.execute_update(
+        '''
+UPDATE feature_flags SET compact_billing_tables = %s;
+''',
+        (compact_billing_tables,),
+    )
+
+    row = await db.select_and_fetchone('SELECT * FROM feature_flags')
+    app['feature_flags'] = row
+
+
 @routes.post('/config-update/pool/{pool}')
 @check_csrf_token
 @auth.web_authenticated_developers_only()
@@ -1335,6 +1364,16 @@ async def monitor_system(app):
     monitor_instances(app)
 
 
+async def compact_agg_billing_project_users_table(app):
+    if not app['feature_flags']['compact_billing_tables']:
+        return
+
+
+async def compact_agg_billing_project_users_by_date_table(app):
+    if not app['feature_flags']['compact_billing_tables']:
+        return
+
+
 async def scheduling_cancelling_bump(app):
     log.info('scheduling cancelling bump loop')
     app['scheduler_state_changed'].notify()
@@ -1412,6 +1451,9 @@ async def on_startup(app):
     app['batch_headers'] = {'Authorization': f'Bearer {row["internal_token"]}'}
     app['frozen'] = row['frozen']
 
+    row = await db.select_and_fetchone('SELECT * FROM feature_flags')
+    app['feature_flags'] = row
+
     await refresh_globals_from_db(app, db)
 
     app['scheduler_state_changed'] = Notice()
@@ -1437,6 +1479,8 @@ async def on_startup(app):
     task_manager.ensure_future(periodically_call(60, scheduling_cancelling_bump, app))
     task_manager.ensure_future(periodically_call(15, monitor_system, app))
     task_manager.ensure_future(periodically_call(5, refresh_globals_from_db, app, db))
+    task_manager.ensure_future(periodically_call(60, compact_agg_billing_project_users_table, app))
+    task_manager.ensure_future(periodically_call(60, compact_agg_billing_project_users_by_date_table, app))
 
 
 async def on_cleanup(app):

batch/batch/driver/templates/index.html

@@ -25,6 +25,19 @@ <h1>Globals</h1>
     {% endif %}
   </div>
 
+<h1>Feature Flags</h1>
+  <div>
+    <form action="{{ base_path }}/configure-feature-flags" method="POST">
+      <input type="hidden" name="_csrf" value="{{ csrf_token }}">
+      <input type="checkbox"
+             id="compact_billing_tables"
+             name="compact_billing_tables"
+             {% if feature_flags['compact_billing_tables'] %}checked{% endif %}
+      value="true" />
+      <button>Update</button>
+    </form>
+  </div>
+
   <h1>Instance Collections</h1>
 
   <h2>Pools</h2>

batch/sql/add-feature-flags.sql

@@ -0,0 +1,5 @@
+CREATE TABLE IF NOT EXISTS `feature_flags` (
+  `compact_billing_tables` BOOLEAN NOT NULL
+) ENGINE = InnoDB;
+
+INSERT INTO `feature_flags` (compact_billing_tables) VALUES (0);

batch/sql/delete-batch-tables.sql

@@ -79,4 +79,5 @@ DROP TABLE IF EXISTS `pools`;
 DROP TABLE IF EXISTS `inst_colls`;
 DROP TABLE IF EXISTS `latest_product_versions`;
 DROP TABLE IF EXISTS `resources`;
-DROP TABLE IF EXISTS `regions`
+DROP TABLE IF EXISTS `regions`;
+DROP TABLE IF EXISTS `feature_flags`;

batch/sql/estimated-current.sql

@@ -5,6 +5,10 @@ CREATE TABLE IF NOT EXISTS `globals` (
   `frozen` BOOLEAN NOT NULL DEFAULT FALSE
 ) ENGINE = InnoDB;
 
+CREATE TABLE IF NOT EXISTS `feature_flags` (
+  `compact_billing_tables` BOOLEAN NOT NULL
+) ENGINE = InnoDB;
+
 CREATE TABLE IF NOT EXISTS `resources` (
   `resource` VARCHAR(100) NOT NULL,
   `rate` DOUBLE NOT NULL,

build.yaml

@@ -2301,6 +2301,9 @@ steps:
       - name: add-list-batches-index
         script: /io/sql/add-list-batches-index.sql
         online: true
+      - name: add-feature-flags
+        script: /io/sql/add-feature-flags.sql
+        online: true
     inputs:
       - from: /repo/batch/sql
         to: /io/sql