feat(flags): Fetch flags and team from database

.github/workflows/ci-plugin-server.yml

@@ -44,7 +44,7 @@ jobs:
                         - 'plugin-server/**'
                         - 'posthog/clickhouse/migrations/**'
                         - 'ee/migrations/**'
-                        - 'ee/management/commands/setup_test_environment.py'
+                        - 'posthog/management/commands/setup_test_environment.py'
                         - 'posthog/migrations/**'
                         - 'posthog/plugins/**'
                         - 'docker*.yml'

.github/workflows/rust.yml

@@ -34,6 +34,9 @@ jobs:
                         - '.github/workflows/rust.yml'
                         - '.github/workflows/rust-docker-build.yml'
                         - '.github/workflows/rust-hook-migrator-docker.yml'
+                        - 'posthog/management/commands/setup_test_environment.py'
+                        - 'posthog/migrations/**'
+                        - 'ee/migrations/**'
 
     build:
         name: Build rust services
@@ -73,6 +76,11 @@ jobs:
 
     test:
         name: Test rust services
+        strategy:
+            matrix:
+                package:
+                    - feature-flags
+                    - others
         needs: changes
         runs-on: depot-ubuntu-22.04-4
         timeout-minutes: 10
@@ -86,11 +94,15 @@ jobs:
             # Use sparse checkout to only select files in rust directory
             # Turning off cone mode ensures that files in the project root are not included during checkout
             - uses: actions/checkout@v3
-              if: needs.changes.outputs.rust == 'true'
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'others'
               with:
                   sparse-checkout: 'rust/'
                   sparse-checkout-cone-mode: false
 
+            # For flags checkout entire repository
+            - uses: actions/checkout@v3
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+
             - name: Login to DockerHub
               if: needs.changes.outputs.rust == 'true'
               uses: docker/login-action@v2
@@ -99,8 +111,15 @@ jobs:
                   username: posthog
                   password: ${{ secrets.DOCKERHUB_TOKEN }}
 
+            - name: Setup main repo dependencies for flags
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+              run: |
+                  docker compose -f ../docker-compose.dev.yml down
+                  docker compose -f ../docker-compose.dev.yml up -d
+                  echo "127.0.0.1 kafka" | sudo tee -a /etc/hosts
+
             - name: Setup dependencies
-              if: needs.changes.outputs.rust == 'true'
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'others'
               run: |
                   docker compose up kafka redis db echo_server -d --wait
                   docker compose up setup_test_db
@@ -119,9 +138,46 @@ jobs:
                       rust/target
                   key: ${ runner.os }-cargo-debug-${{ hashFiles('**/Cargo.lock') }}
 
+            - name: Set up Python
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+              uses: actions/setup-python@v5
+              with:
+                  python-version: 3.11.9
+                  cache: 'pip'
+                  cache-dependency-path: '**/requirements*.txt'
+                  token: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
+
+            # uv is a fast pip alternative: https://github.com/astral-sh/uv/
+            - run: pip install uv
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+
+            - name: Install SAML (python3-saml) dependencies
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+              run: |
+                  sudo apt-get update
+                  sudo apt-get install libxml2-dev libxmlsec1-dev libxmlsec1-openssl
+
+            - name: Install python dependencies
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+              run: |
+                  uv pip install --system -r ../requirements-dev.txt
+                  uv pip install --system -r ../requirements.txt
+
+            - name: Set up databases
+              if: needs.changes.outputs.rust == 'true' && matrix.package == 'feature-flags'
+              env:
+                  DEBUG: 'true'
+                  TEST: 'true'
+                  SECRET_KEY: 'abcdef' # unsafe - for testing only
+                  DATABASE_URL: 'postgres://posthog:posthog@localhost:5432/posthog'
+              run: cd ../ && python manage.py setup_test_environment --only-postgres
+
             - name: Run cargo test
               if: needs.changes.outputs.rust == 'true'
-              run: cargo test --all-features
+              run: |
+                  echo "Starting cargo test"
+                  cargo test --all-features ${{ matrix.package == 'feature-flags' && '--package feature-flags' || '--workspace --exclude feature-flags' }}
+                  echo "Cargo test completed"
 
     linting:
         name: Lint rust services

posthog/management/commands/setup_test_environment.py

@@ -26,6 +26,12 @@
 class Command(BaseCommand):
     help = "Set up databases for non-Python tests that depend on the Django server"
 
+    # has optional arg to only run postgres setup
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--only-postgres", action="store_true", help="Only set up the Postgres database", default=False
+        )
+
     def handle(self, *args, **options):
         if not TEST:
             raise ValueError("TEST environment variable needs to be set for this command to function")
@@ -36,6 +42,10 @@ def handle(self, *args, **options):
         test_runner.setup_databases()
         test_runner.setup_test_environment()
 
+        if options["only_postgres"]:
+            print("Only setting up Postgres database")  # noqa: T201
+            return
+
         print("\nCreating test ClickHouse database...")  # noqa: T201
         database = Database(
             CLICKHOUSE_DATABASE,

rust/feature-flags/Cargo.toml

@@ -27,6 +27,8 @@ thiserror = { workspace = true }
 serde-pickle = { version = "1.1.1"}
 sha1 = "0.10.6"
 regex = "1.10.4"
+sqlx = { workspace = true }
+uuid = { workspace = true }
 
 [lints]
 workspace = true

rust/feature-flags/README.md

@@ -1,6 +1,23 @@
 
 # Testing
 
+First, make sure docker compose is running (from main posthog repo), and test database exists:
+
+```
+docker compose -f ../docker-compose.dev.yml up -d
+```
+
+```
+TEST=1 python manage.py setup_test_environment --only-postgres
+```
+
+We only need to run the above once, when the test database is created.
+
+TODO: Would be nice to make the above automatic.
+
+
+Then, run the tests:
+
 ```
 cargo test --package feature-flags
 ```

rust/feature-flags/src/api.rs

@@ -5,6 +5,9 @@ use axum::response::{IntoResponse, Response};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
+use crate::database::CustomDatabaseError;
+use crate::redis::CustomRedisError;
+
 #[derive(Debug, PartialEq, Eq, Deserialize, Serialize)]
 pub enum FlagsResponseCode {
     Ok = 1,
@@ -42,6 +45,14 @@ pub enum FlagError {
     DataParsingError,
     #[error("redis unavailable")]
     RedisUnavailable,
+    #[error("database unavailable")]
+    DatabaseUnavailable,
+    #[error("Timed out while fetching data")]
+    TimeoutError,
+    // TODO: Consider splitting top-level errors (that are returned to the client)
+    // and FlagMatchingError, like timeouterror which we can gracefully handle.
+    // This will make the `into_response` a lot clearer as well, since it wouldn't
+    // have arbitrary errors that actually never make it to the client.
 }
 
 impl IntoResponse for FlagError {
@@ -58,10 +69,53 @@ impl IntoResponse for FlagError {
 
             FlagError::RateLimited => (StatusCode::TOO_MANY_REQUESTS, self.to_string()),
 
-            FlagError::DataParsingError | FlagError::RedisUnavailable => {
-                (StatusCode::SERVICE_UNAVAILABLE, self.to_string())
-            }
+            FlagError::DataParsingError
+            | FlagError::RedisUnavailable
+            | FlagError::DatabaseUnavailable
+            | FlagError::TimeoutError => (StatusCode::SERVICE_UNAVAILABLE, self.to_string()),
         }
         .into_response()
     }
 }
+
+impl From<CustomRedisError> for FlagError {
+    fn from(e: CustomRedisError) -> Self {
+        match e {
+            CustomRedisError::NotFound => FlagError::TokenValidationError,
+            CustomRedisError::PickleError(e) => {
+                tracing::error!("failed to fetch data: {}", e);
+                FlagError::DataParsingError
+            }
+            CustomRedisError::Timeout(_) => FlagError::TimeoutError,
+            CustomRedisError::Other(e) => {
+                tracing::error!("Unknown redis error: {}", e);
+                FlagError::RedisUnavailable
+            }
+        }
+    }
+}
+
+impl From<CustomDatabaseError> for FlagError {
+    fn from(e: CustomDatabaseError) -> Self {
+        match e {
+            CustomDatabaseError::NotFound => FlagError::TokenValidationError,
+            CustomDatabaseError::Other(_) => {
+                tracing::error!("failed to get connection: {}", e);
+                FlagError::DatabaseUnavailable
+            }
+            CustomDatabaseError::Timeout(_) => FlagError::TimeoutError,
+        }
+    }
+}
+
+impl From<sqlx::Error> for FlagError {
+    fn from(e: sqlx::Error) -> Self {
+        // TODO: Be more precise with error handling here
+        tracing::error!("sqlx error: {}", e);
+        println!("sqlx error: {}", e);
+        match e {
+            sqlx::Error::RowNotFound => FlagError::TokenValidationError,
+            _ => FlagError::DatabaseUnavailable,
+        }
+    }
+}

rust/feature-flags/src/config.rs

@@ -7,10 +7,10 @@ pub struct Config {
     #[envconfig(default = "127.0.0.1:3001")]
     pub address: SocketAddr,
 
-    #[envconfig(default = "postgres://posthog:posthog@localhost:15432/test_database")]
+    #[envconfig(default = "postgres://posthog:posthog@localhost:5432/test_posthog")]
     pub write_database_url: String,
 
-    #[envconfig(default = "postgres://posthog:posthog@localhost:15432/test_database")]
+    #[envconfig(default = "postgres://posthog:posthog@localhost:5432/test_posthog")]
     pub read_database_url: String,
 
     #[envconfig(default = "1024")]

rust/feature-flags/src/database.rs

@@ -0,0 +1,86 @@
+use std::time::Duration;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use sqlx::{
+    pool::PoolConnection,
+    postgres::{PgPoolOptions, PgRow},
+    Postgres,
+};
+use thiserror::Error;
+use tokio::time::timeout;
+
+const DATABASE_TIMEOUT_MILLISECS: u64 = 1000;
+
+#[derive(Error, Debug)]
+pub enum CustomDatabaseError {
+    #[error("Not found in database")]
+    NotFound,
+
+    #[error("Pg error: {0}")]
+    Other(#[from] sqlx::Error),
+
+    #[error("Timeout error")]
+    Timeout(#[from] tokio::time::error::Elapsed),
+}
+
+/// A simple db wrapper
+/// Supports running any arbitrary query with a timeout.
+/// TODO: Make sqlx prepared statements work with pgbouncer, potentially by setting pooling mode to session.
+#[async_trait]
+pub trait Client {
+    async fn get_connection(&self) -> Result<PoolConnection<Postgres>, CustomDatabaseError>;
+    async fn run_query(
+        &self,
+        query: String,
+        parameters: Vec<String>,
+        timeout_ms: Option<u64>,
+    ) -> Result<Vec<PgRow>, CustomDatabaseError>;
+}
+
+pub struct PgClient {
+    pool: sqlx::PgPool,
+}
+
+impl PgClient {
+    pub async fn new(addr: String) -> Result<PgClient, CustomDatabaseError> {
+        // TODO: Get these vals from config
+        let pool = PgPoolOptions::new()
+            .max_connections(5)
+            .acquire_timeout(Duration::from_secs(1))
+            .test_before_acquire(true)
+            .connect(&addr)
+            .await?;
+
+        Ok(PgClient { pool })
+    }
+}
+
+#[async_trait]
+impl Client for PgClient {
+    async fn run_query(
+        &self,
+        query: String,
+        parameters: Vec<String>,
+        timeout_ms: Option<u64>,
+    ) -> Result<Vec<PgRow>, CustomDatabaseError> {
+        let built_query = sqlx::query(&query);
+        let built_query = parameters
+            .iter()
+            .fold(built_query, |acc, param| acc.bind(param));
+        let query_results = built_query.fetch_all(&self.pool);
+
+        let timeout_ms = match timeout_ms {
+            Some(ms) => ms,
+            None => DATABASE_TIMEOUT_MILLISECS,
+        };
+
+        let fut = timeout(Duration::from_secs(timeout_ms), query_results).await?;
+
+        Ok(fut?)
+    }
+
+    async fn get_connection(&self) -> Result<PoolConnection<Postgres>, CustomDatabaseError> {
+        Ok(self.pool.acquire().await?)
+    }
+}