From 26304e94019ff757a5c53e27198505c46bb4dde6 Mon Sep 17 00:00:00 2001
From: Gaetano Guerriero <gaetano@athenian.co>
Date: Thu, 20 Apr 2023 15:03:10 +0200
Subject: [PATCH] [DEV-6438] Define models for account monitor

---
 .../api/models/persistentdata/models.py       | 23 +++++++-
 .../versions/765542234720_account_monitor.py  | 53 +++++++++++++++++++
 server/tests/conftest.py                      |  4 +-
 3 files changed, 77 insertions(+), 3 deletions(-)
 create mode 100644 server/athenian/api/models/persistentdata/versions/765542234720_account_monitor.py

diff --git a/server/athenian/api/models/persistentdata/models.py b/server/athenian/api/models/persistentdata/models.py
index 7e8f30a47b..92810327e6 100644
--- a/server/athenian/api/models/persistentdata/models.py
+++ b/server/athenian/api/models/persistentdata/models.py
@@ -18,6 +18,8 @@
 
 from athenian.api.models import create_base
 
+JSONType = JSONB().with_variant(JSON(), sqlite.dialect.name)
+
 
 class ShardedByAccount:
     """All the tables contain `account_id` primary key."""
@@ -106,7 +108,7 @@ class DeployedLabel(Base):
 
     deployment_name = Column(Text(), primary_key=True)
     key = Column(Text(), primary_key=True)
-    value = Column(JSONB().with_variant(JSON(), sqlite.dialect.name))
+    value = Column(JSONType)
 
 
 class DeployedComponent(create_time_mixin(created_at=True, updated_at=False), Base):
@@ -143,7 +145,7 @@ class HealthMetric(Base):
         default=lambda: datetime.now(timezone.utc),
         server_default=func.now(),
     )
-    value = Column(JSONB().with_variant(JSON(), sqlite.dialect.name), nullable=False)
+    value = Column(JSONType, nullable=False)
 
 
 class VitallyAccount(Base, create_time_mixin(created_at=False, updated_at=True)):
@@ -154,3 +156,20 @@ class VitallyAccount(Base, create_time_mixin(created_at=False, updated_at=True))
     name = Column(Text())
     mrr = Column(DECIMAL())
     health_score = Column(Float())
+
+
+class AccMonitorCheckLog(Base):
+    """The log of the execution of a check in the account monitor."""
+
+    __table_args__ = {"schema": "acc_monitor"}
+    __tablename__ = "check_logs"
+
+    created_at = Column(
+        TIMESTAMP(timezone=True),
+        primary_key=True,
+        default=lambda: datetime.now(timezone.utc),
+        server_default=func.now(),
+    )
+    check_name = Column(Text, primary_key=True)
+    passed = Column(Boolean, nullable=False)
+    result = Column(JSONType)
diff --git a/server/athenian/api/models/persistentdata/versions/765542234720_account_monitor.py b/server/athenian/api/models/persistentdata/versions/765542234720_account_monitor.py
new file mode 100644
index 0000000000..9057596040
--- /dev/null
+++ b/server/athenian/api/models/persistentdata/versions/765542234720_account_monitor.py
@@ -0,0 +1,53 @@
+"""Account monitor
+
+Revision ID: 765542234720
+Revises: e06c9947fbdb
+Create Date: 2023-04-19 08:50:36.182075+00:00
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "765542234720"
+down_revision = "e06c9947fbdb"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    name, schema_arg = _name_schema_arg("acc_monitor", "check_logs")
+
+    op.create_table(
+        name,
+        sa.Column("account_id", sa.Integer(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.func.now(),
+            nullable=False,
+        ),
+        sa.Column("check_name", sa.Text(), nullable=False),
+        sa.Column("passed", sa.Boolean(), nullable=False),
+        sa.Column(
+            "result",
+            postgresql.JSONB(astext_type=sa.Text()).with_variant(sa.JSON(), "sqlite"),
+            nullable=True,
+        ),
+        sa.PrimaryKeyConstraint("account_id", "created_at", "check_name"),
+        **schema_arg,
+    )
+
+
+def downgrade():
+    name, schema_arg = _name_schema_arg("acc_monitor", "check_logs")
+    op.drop_table(name, **schema_arg)
+
+
+def _name_schema_arg(schema, name):
+    if op.get_bind().dialect.name == "postgresql":
+        op.execute(f"CREATE SCHEMA IF NOT EXISTS {schema}")
+        return name, {"schema": schema}
+    else:
+        return f"schema.{name}", {}
diff --git a/server/tests/conftest.py b/server/tests/conftest.py
index 36b85f8c04..228309c314 100644
--- a/server/tests/conftest.py
+++ b/server/tests/conftest.py
@@ -636,7 +636,9 @@ def persistentdata_db(worker_id) -> str:
         PersistentdataBase,
         worker_id,
         {
-            "postgresql": "create schema if not exists athenian;",
+            "postgresql": (
+                "CREATE SCHEMA IF NOT EXISTS athenian; CREATE SCHEMA IF NOT EXISTS acc_monitor "
+            ),
         },
     )