From 5405a275ab5a665d5b57696bb44ca06c9f6d6cff Mon Sep 17 00:00:00 2001
From: Timothy Messier <tim.messier@gmail.com>
Date: Fri, 16 Feb 2024 16:19:25 +0000
Subject: [PATCH] feat(census): Update census tables to support additional
 metrics

---
 .../oss/postgres/72/01_census.up.sql          |  1 +
 .../postgres/82/09_cenus_add_metric.up.sql    | 47 +++++++++++++++++++
 .../db/sqltest/tests/census/last_uploaded.sql | 47 +++++++++++++++----
 3 files changed, 85 insertions(+), 10 deletions(-)
 create mode 100644 internal/db/schema/migrations/oss/postgres/82/09_cenus_add_metric.up.sql

diff --git a/internal/db/schema/migrations/oss/postgres/72/01_census.up.sql b/internal/db/schema/migrations/oss/postgres/72/01_census.up.sql
index 71b14d6611..5ca5710c87 100644
--- a/internal/db/schema/migrations/oss/postgres/72/01_census.up.sql
+++ b/internal/db/schema/migrations/oss/postgres/72/01_census.up.sql
@@ -12,6 +12,7 @@ begin;
 
   -- This index ensures that there will only ever be one row in the table.
   -- See: https://www.postgresql.org/docs/current/indexes-expressional.html
+  -- Dropped in 82/09_census_add_metric.up.sql
   create unique index census_last_uploaded_one_row
     on census_last_uploaded((last_uploaded_at is not null));
 
diff --git a/internal/db/schema/migrations/oss/postgres/82/09_cenus_add_metric.up.sql b/internal/db/schema/migrations/oss/postgres/82/09_cenus_add_metric.up.sql
new file mode 100644
index 0000000000..5b1ce7b777
--- /dev/null
+++ b/internal/db/schema/migrations/oss/postgres/82/09_cenus_add_metric.up.sql
@@ -0,0 +1,47 @@
+-- Copyright (c) HashiCorp, Inc.
+-- SPDX-License-Identifier: BUSL-1.1
+
+begin;
+  -- Create enum table for valid cenus metrics.
+  create table census_metric_enm (
+    name text primary key
+    constraint only_predefined_census_metrics_allowed
+    check (
+      name in (
+        'sessions',
+        'active_users'
+      )
+    )
+  );
+  comment on table census_metric_enm is
+    'census_metric_enm is an enumeration table for census metric types.';
+
+  insert into census_metric_enm
+              (name)
+       values ('sessions'),
+              ('active_users');
+
+  -- Drop constraint and index from 72/01_census.up.sql
+  drop index census_last_uploaded_one_row;
+  alter table census_last_uploaded drop constraint census_last_uploaded_pkey;
+
+  -- Add new column and populate the default value.
+  alter table census_last_uploaded
+   add column metric text
+              default 'sessions'
+              not null
+              primary key
+              constraint census_metric_enm_fkey
+                references census_metric_enm (name)
+                on delete restrict
+                on update restrict;
+  alter table census_last_uploaded alter column metric drop default;
+  comment on table census_last_uploaded is
+    'census_last_uploaded is a table which contains the timestamp '
+    'of the last time census data was uploaded, per census metric.';
+
+  -- Add new row for active_users.
+  insert into census_last_uploaded
+              (last_uploaded_at,         metric)
+       values ('-infinity'::timestamptz, 'active_users');
+commit;
diff --git a/internal/db/sqltest/tests/census/last_uploaded.sql b/internal/db/sqltest/tests/census/last_uploaded.sql
index 79a8c4a02b..7f7eaff3bd 100644
--- a/internal/db/sqltest/tests/census/last_uploaded.sql
+++ b/internal/db/sqltest/tests/census/last_uploaded.sql
@@ -3,28 +3,55 @@
 
 begin;
 
-  select plan(6);
+  select plan(7);
 
   select has_table('census_last_uploaded');
-  select is(count(*), 1::bigint, 'census_last_uploaded should have only 1 row') from census_last_uploaded;
-  select ok(not isfinite(last_uploaded_at)) from census_last_uploaded;
+  select is(count(*), 2::bigint, 'census_last_uploaded should have only 2 row') from census_last_uploaded;
+  select results_eq(
+           $$
+           select last_uploaded_at::timestamptz, metric
+             from census_last_uploaded
+           $$,
+           $$
+           values ('-infinity'::timestamptz, 'sessions'),
+                  ('-infinity'::timestamptz, 'active_users')
+           $$);
+
 
   prepare insert_row as
    insert into census_last_uploaded
-    (last_uploaded_at)
-   values
-    (now());
+          (last_uploaded_at, metric)
+   values (now(), 'sessions');
+
+  select throws_ok('insert_row',
+                   '23505',
+                   'duplicate key value violates unique constraint "census_last_uploaded_pkey"',
+                   'insert into census_last_uploaded should fail');
+
+  prepare insert_row_invalid_metric as
+   insert into census_last_uploaded
+          (last_uploaded_at, metric)
+   values (now(), 'foo');
 
-  select throws_ok('insert_row', '23505',
-    'duplicate key value violates unique constraint "census_last_uploaded_one_row"',
-    'insert into census_last_uploaded should fail');
+  select throws_ok('insert_row_invalid_metric',
+                   '23503',
+                   'insert or update on table "census_last_uploaded" violates foreign key constraint "census_metric_enm_fkey"',
+                   'insert into census_last_uploaded should fail');
 
   prepare update_census_last_uploaded as
    update census_last_uploaded
       set last_uploaded_at = now();
 
   select lives_ok('update_census_last_uploaded');
-  select ok(isfinite(last_uploaded_at)) from census_last_uploaded;
+  select results_eq(
+           $$
+           select last_uploaded_at::timestamptz, metric
+             from census_last_uploaded
+           $$,
+           $$
+           values (now(), 'sessions'),
+                  (now(), 'active_users')
+           $$);
 
   select * from finish();