diff --git a/batch/batch/driver/main.py b/batch/batch/driver/main.py index e65ddfa7953..66a466c1fb5 100644 --- a/batch/batch/driver/main.py +++ b/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): diff --git a/batch/batch/driver/templates/index.html b/batch/batch/driver/templates/index.html index 161f0292bc2..26016670cd5 100644 --- a/batch/batch/driver/templates/index.html +++ b/batch/batch/driver/templates/index.html @@ -25,6 +25,19 @@

Globals

{% endif %} +

Feature Flags

+
+
+ + + +
+
+

Instance Collections

Pools

diff --git a/batch/sql/add-feature-flags.sql b/batch/sql/add-feature-flags.sql new file mode 100644 index 00000000000..cb78ec627a5 --- /dev/null +++ b/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); diff --git a/batch/sql/delete-batch-tables.sql b/batch/sql/delete-batch-tables.sql index 754f42018c3..c59e928b6e8 100644 --- a/batch/sql/delete-batch-tables.sql +++ b/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`; diff --git a/batch/sql/estimated-current.sql b/batch/sql/estimated-current.sql index 5bdf4d430c8..2788c825321 100644 --- a/batch/sql/estimated-current.sql +++ b/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, diff --git a/build.yaml b/build.yaml index 4da45283ef9..a0c1dbcc3a5 100644 --- a/build.yaml +++ b/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