Line data    Source code

       1              : use std::{collections::HashMap, sync::Arc, time::Duration};
       2              : 
       3              : use posthog_client_lite::{
       4              :     FeatureResolverBackgroundLoop, PostHogClientConfig, PostHogEvaluationError,
       5              : };
       6              : use tokio_util::sync::CancellationToken;
       7              : use utils::id::TenantId;
       8              : 
       9              : use crate::config::PageServerConf;
      10              : 
      11              : #[derive(Clone)]
      12              : pub struct FeatureResolver {
      13              :     inner: Option<Arc<FeatureResolverBackgroundLoop>>,
      14              : }
      15              : 
      16              : impl FeatureResolver {
      17          117 :     pub fn new_disabled() -> Self {
      18          117 :         Self { inner: None }
      19          117 :     }
      20              : 
      21            0 :     pub fn spawn(
      22            0 :         conf: &PageServerConf,
      23            0 :         shutdown_pageserver: CancellationToken,
      24            0 :         handle: &tokio::runtime::Handle,
      25            0 :     ) -> anyhow::Result<Self> {
      26              :         // DO NOT block in this function: make it return as fast as possible to avoid startup delays.
      27            0 :         if let Some(posthog_config) = &conf.posthog_config {
      28            0 :             let inner = FeatureResolverBackgroundLoop::new(
      29            0 :                 PostHogClientConfig {
      30            0 :                     server_api_key: posthog_config.server_api_key.clone(),
      31            0 :                     client_api_key: posthog_config.client_api_key.clone(),
      32            0 :                     project_id: posthog_config.project_id.clone(),
      33            0 :                     private_api_url: posthog_config.private_api_url.clone(),
      34            0 :                     public_api_url: posthog_config.public_api_url.clone(),
      35            0 :                 },
      36            0 :                 shutdown_pageserver,
      37            0 :             );
      38            0 :             let inner = Arc::new(inner);
      39            0 :             // TODO: make this configurable
      40            0 :             inner.clone().spawn(handle, Duration::from_secs(60));
      41            0 :             Ok(FeatureResolver { inner: Some(inner) })
      42              :         } else {
      43            0 :             Ok(FeatureResolver { inner: None })
      44              :         }
      45            0 :     }
      46              : 
      47              :     /// Evaluate a multivariate feature flag. Currently, we do not support any properties.
      48            0 :     pub fn evaluate_multivariate(
      49            0 :         &self,
      50            0 :         flag_key: &str,
      51            0 :         tenant_id: TenantId,
      52            0 :     ) -> Result<String, PostHogEvaluationError> {
      53            0 :         if let Some(inner) = &self.inner {
      54            0 :             inner.feature_store().evaluate_multivariate(
      55            0 :                 flag_key,
      56            0 :                 &tenant_id.to_string(),
      57            0 :                 &HashMap::new(),
      58            0 :             )
      59              :         } else {
      60            0 :             Err(PostHogEvaluationError::NotAvailable(
      61            0 :                 "PostHog integration is not enabled".to_string(),
      62            0 :             ))
      63              :         }
      64            0 :     }
      65              : }