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

@@ -15,6 +15,7 @@ tokio = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }
 bytes = { workspace = true }
+once_cell = "1.18.0"
 rand = { workspace = true }
 redis = { version = "0.23.3", features = [
     "tokio-comp",
@@ -27,12 +28,13 @@ 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
 
 [dev-dependencies]
 assert-json-diff = { workspace = true }
-once_cell = "1.18.0"
 reqwest = { 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

@@ -1,16 +1,17 @@
-use std::net::SocketAddr;
-
 use envconfig::Envconfig;
+use once_cell::sync::Lazy;
+use std::net::SocketAddr;
+use std::str::FromStr;
 
-#[derive(Envconfig, Clone)]
+#[derive(Envconfig, Clone, Debug)]
 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")]
@@ -21,4 +22,83 @@ pub struct Config {
 
     #[envconfig(default = "redis://localhost:6379/")]
     pub redis_url: String,
+
+    #[envconfig(default = "1")]
+    pub acquire_timeout_secs: u64,
+}
+
+impl Config {
+    pub fn default_test_config() -> Self {
+        Self {
+            address: SocketAddr::from_str("127.0.0.1:0").unwrap(),
+            redis_url: "redis://localhost:6379/".to_string(),
+            write_database_url: "postgres://posthog:posthog@localhost:5432/test_posthog"
+                .to_string(),
+            read_database_url: "postgres://posthog:posthog@localhost:5432/test_posthog".to_string(),
+            max_concurrent_jobs: 1024,
+            max_pg_connections: 100,
+            acquire_timeout_secs: 1,
+        }
+    }
+}
+
+pub static DEFAULT_TEST_CONFIG: Lazy<Config> = Lazy::new(Config::default_test_config);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_default_config() {
+        let config = Config::init_from_env().unwrap();
+        assert_eq!(
+            config.address,
+            SocketAddr::from_str("127.0.0.1:3001").unwrap()
+        );
+        assert_eq!(
+            config.write_database_url,
+            "postgres://posthog:posthog@localhost:5432/test_posthog"
+        );
+        assert_eq!(
+            config.read_database_url,
+            "postgres://posthog:posthog@localhost:5432/test_posthog"
+        );
+        assert_eq!(config.max_concurrent_jobs, 1024);
+        assert_eq!(config.max_pg_connections, 100);
+        assert_eq!(config.redis_url, "redis://localhost:6379/");
+    }
+
+    #[test]
+    fn test_default_test_config() {
+        let config = Config::default_test_config();
+        assert_eq!(config.address, SocketAddr::from_str("127.0.0.1:0").unwrap());
+        assert_eq!(
+            config.write_database_url,
+            "postgres://posthog:posthog@localhost:5432/test_posthog"
+        );
+        assert_eq!(
+            config.read_database_url,
+            "postgres://posthog:posthog@localhost:5432/test_posthog"
+        );
+        assert_eq!(config.max_concurrent_jobs, 1024);
+        assert_eq!(config.max_pg_connections, 100);
+        assert_eq!(config.redis_url, "redis://localhost:6379/");
+    }
+
+    #[test]
+    fn test_default_test_config_static() {
+        let config = &*DEFAULT_TEST_CONFIG;
+        assert_eq!(config.address, SocketAddr::from_str("127.0.0.1:0").unwrap());
+        assert_eq!(
+            config.write_database_url,
+            "postgres://posthog:posthog@localhost:5432/test_posthog"
+        );
+        assert_eq!(
+            config.read_database_url,
+            "postgres://posthog:posthog@localhost:5432/test_posthog"
+        );
+        assert_eq!(config.max_concurrent_jobs, 1024);
+        assert_eq!(config.max_pg_connections, 100);
+        assert_eq!(config.redis_url, "redis://localhost:6379/");
+    }
 }