compliance: compliance-server: db: Define migrations and query db

compliance/compliance-api/src/lib.rs

@@ -1,3 +1,11 @@
+//! The API for the compliance server
+
+#![deny(missing_docs)]
+#![deny(clippy::missing_docs_in_private_items)]
+#![deny(unsafe_code)]
+#![deny(clippy::needless_pass_by_value)]
+#![deny(clippy::needless_pass_by_ref_mut)]
+
 use serde::{Deserialize, Serialize};
 
 /// The response type for a compliance check
@@ -13,5 +21,6 @@ pub enum ComplianceStatus {
     /// The wallet is compliant
     Compliant,
     /// The wallet is not compliant
-    NotCompliant,
+    #[allow(missing_docs)]
+    NotCompliant { reason: String },
 }

compliance/compliance-server/Cargo.toml

@@ -9,6 +9,9 @@ http-body-util = "0.1.0"
 warp = "0.3"
 compliance-api = { path = "../compliance-api" }
 
+# === Database === #
+diesel = { version = "2.2", features = ["postgres", "r2d2"] }
+
 # === Renegade Dependencies === #
 renegade-util = { workspace = true }
 

compliance/compliance-server/diesel.toml

@@ -0,0 +1,9 @@
+# For documentation on how to configure this file,
+# see https://diesel.rs/guides/configuring-diesel-cli
+
+[print_schema]
+file = "src/schema.rs"
+custom_type_derives = ["diesel::query_builder::QueryId", "Clone"]
+
+[migrations_directory]
+dir = "./migrations"

compliance/compliance-server/migrations/00000000000000_diesel_initial_setup/down.sql

@@ -0,0 +1,6 @@
+-- This file was automatically created by Diesel to setup helper functions
+-- and other internal bookkeeping. This file is safe to edit, any future
+-- changes will be added to existing projects as new migrations.
+
+DROP FUNCTION IF EXISTS diesel_manage_updated_at(_tbl regclass);
+DROP FUNCTION IF EXISTS diesel_set_updated_at();

compliance/compliance-server/migrations/00000000000000_diesel_initial_setup/up.sql

@@ -0,0 +1,36 @@
+-- This file was automatically created by Diesel to setup helper functions
+-- and other internal bookkeeping. This file is safe to edit, any future
+-- changes will be added to existing projects as new migrations.
+
+
+
+
+-- Sets up a trigger for the given table to automatically set a column called
+-- `updated_at` whenever the row is modified (unless `updated_at` was included
+-- in the modified columns)
+--
+-- # Example
+--
+-- ```sql
+-- CREATE TABLE users (id SERIAL PRIMARY KEY, updated_at TIMESTAMP NOT NULL DEFAULT NOW());
+--
+-- SELECT diesel_manage_updated_at('users');
+-- ```
+CREATE OR REPLACE FUNCTION diesel_manage_updated_at(_tbl regclass) RETURNS VOID AS $$
+BEGIN
+    EXECUTE format('CREATE TRIGGER set_updated_at BEFORE UPDATE ON %s
+                    FOR EACH ROW EXECUTE PROCEDURE diesel_set_updated_at()', _tbl);
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION diesel_set_updated_at() RETURNS trigger AS $$
+BEGIN
+    IF (
+        NEW IS DISTINCT FROM OLD AND
+        NEW.updated_at IS NOT DISTINCT FROM OLD.updated_at
+    ) THEN
+        NEW.updated_at := current_timestamp;
+    END IF;
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;

compliance/compliance-server/migrations/2024-06-28-231756_create_compliance_table/down.sql

@@ -0,0 +1,2 @@
+-- Drop the wallet compliance table
+DROP TABLE IF EXISTS wallet_compliance;

compliance/compliance-server/migrations/2024-06-28-231756_create_compliance_table/up.sql

@@ -0,0 +1,8 @@
+-- Create a table for caching wallet compliance information
+CREATE TABLE IF NOT EXISTS wallet_compliance (
+    address TEXT PRIMARY KEY,
+    is_compliant BOOLEAN NOT NULL,
+    reason TEXT NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
+    expires_at TIMESTAMP NOT NULL DEFAULT NOW() + INTERVAL '1 year'
+);

compliance/compliance-server/src/db.rs

@@ -0,0 +1,59 @@
+//! Database helpers for the server
+
+use std::time::SystemTime;
+
+use compliance_api::ComplianceStatus;
+use diesel::{ExpressionMethods, Insertable, PgConnection, QueryDsl, Queryable, RunQueryDsl};
+use renegade_util::err_str;
+
+use crate::{
+    error::ComplianceServerError,
+    schema::{
+        wallet_compliance,
+        wallet_compliance::dsl::{address as address_col, wallet_compliance as compliance_table},
+    },
+};
+
+// ----------
+// | Models |
+// ----------
+
+/// A compliance entry for a wallet
+#[derive(Debug, Clone, Queryable, Insertable)]
+#[table_name = "wallet_compliance"]
+#[allow(missing_docs)]
+pub struct ComplianceEntry {
+    pub address: String,
+    pub is_compliant: bool,
+    pub reason: String,
+    pub created_at: SystemTime,
+    pub expires_at: SystemTime,
+}
+
+impl ComplianceEntry {
+    /// Get the compliance status for an entry
+    pub fn compliance_status(&self) -> ComplianceStatus {
+        if self.is_compliant {
+            ComplianceStatus::Compliant
+        } else {
+            ComplianceStatus::NotCompliant { reason: self.reason.clone() }
+        }
+    }
+}
+
+// -----------
+// | Queries |
+// -----------
+
+/// Get a compliance entry by address
+pub fn get_compliance_entry(
+    address: &str,
+    conn: &mut PgConnection,
+) -> Result<Option<ComplianceEntry>, ComplianceServerError> {
+    let query = compliance_table
+        .filter(address_col.eq(address))
+        .load::<ComplianceEntry>(conn)
+        .map_err(err_str!(ComplianceServerError::Db))?;
+
+    Ok(query.first().cloned())
+}

compliance/compliance-server/src/error.rs

@@ -6,11 +6,16 @@ use warp::reject::Reject;
 
 /// The error type emitted by the compliance server
 #[derive(Debug, Clone)]
-pub enum ComplianceServerError {}
+pub enum ComplianceServerError {
+    /// An error with a database query
+    Db(String),
+}
 
 impl Display for ComplianceServerError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "ComplianceServerError")
+        match self {
+            ComplianceServerError::Db(e) => write!(f, "Database error: {}", e),
+        }
     }
 }
 impl Error for ComplianceServerError {}

compliance/compliance-server/src/main.rs

@@ -1,10 +1,31 @@
+//! The server that serves wallet compliance information
+
+#![deny(missing_docs)]
+#![deny(clippy::missing_docs_in_private_items)]
+#![deny(unsafe_code)]
+#![deny(clippy::needless_pass_by_value)]
+#![deny(clippy::needless_pass_by_ref_mut)]
+
+use std::sync::Arc;
+
 use clap::Parser;
 use compliance_api::{ComplianceCheckResponse, ComplianceStatus};
+use diesel::pg::PgConnection;
+use diesel::r2d2::{ConnectionManager, Pool};
 use error::ComplianceServerError;
+use renegade_util::err_str;
 use renegade_util::telemetry::{setup_system_logger, LevelFilter};
 use warp::{reply::Json, Filter};
 
+use crate::db::get_compliance_entry;
+
+pub mod db;
 pub mod error;
+#[allow(missing_docs, clippy::missing_docs_in_private_items)]
+pub mod schema;
+
+/// The type of the connection pool
+type ConnectionPool = Arc<Pool<ConnectionManager<PgConnection>>>;
 
 /// The CLI for the compliance server
 #[derive(Debug, Clone, Parser)]
@@ -16,13 +37,21 @@ struct Cli {
     /// The Chainalysis API key
     #[arg(long)]
     chainalysis_api_key: String,
+    /// The url of the compliance database
+    #[arg(long)]
+    db_url: String,
 }
 
 #[tokio::main]
 async fn main() {
     setup_system_logger(LevelFilter::INFO);
     let cli = Cli::parse();
 
+    // Create the connection pool
+    let manager = ConnectionManager::<PgConnection>::new(cli.db_url.clone());
+    let pool = Pool::builder().build(manager).expect("Failed to create pool");
+    let pool = Arc::new(pool);
+
     // Get compliance information for a wallet
     let chainalysis_key = cli.chainalysis_api_key.clone();
     let compliance_check = warp::get()
@@ -31,8 +60,10 @@ async fn main() {
         .and(warp::path::param::<String>()) // wallet_address
         .and_then(move |wallet_address| {
             let key = chainalysis_key.clone();
+            let pool = pool.clone();
+
             async move {
-                handle_req(wallet_address, &key).await
+                handle_req(wallet_address, &key, pool).await
             }
         });
 
@@ -49,8 +80,10 @@ async fn main() {
 async fn handle_req(
     wallet_address: String,
     chainalysis_api_key: &str,
+    pool: ConnectionPool,
 ) -> Result<Json, warp::Rejection> {
-    let compliance_status = check_wallet_compliance(wallet_address, chainalysis_api_key).await?;
+    let compliance_status =
+        check_wallet_compliance(wallet_address, chainalysis_api_key, pool).await?;
     let resp = ComplianceCheckResponse { compliance_status };
     Ok(warp::reply::json(&resp))
 }
@@ -59,10 +92,15 @@ async fn handle_req(
 async fn check_wallet_compliance(
     wallet_address: String,
     chainalysis_api_key: &str,
+    pool: ConnectionPool,
 ) -> Result<ComplianceStatus, ComplianceServerError> {
     // 1. Check the DB first
+    let mut conn = pool.get().map_err(err_str!(ComplianceServerError::Db))?;
+    let compliance_entry = get_compliance_entry(&wallet_address, &mut conn)?;
+    if let Some(compliance_entry) = compliance_entry {
+        return Ok(compliance_entry.compliance_status());
+    }
 
     // 2. If not present, check the chainalysis API
-
-    todo!()
+    Ok(ComplianceStatus::Compliant)
 }

compliance/compliance-server/src/schema.rs

@@ -0,0 +1,11 @@
+// @generated automatically by Diesel CLI.
+
+diesel::table! {
+    wallet_compliance (address) {
+        address -> Text,
+        is_compliant -> Bool,
+        reason -> Text,
+        created_at -> Timestamp,
+        expires_at -> Timestamp,
+    }
+}