spin_factor_outbound_pg/
lib.rs

1mod allowed_hosts;
2pub mod client;
3mod host;
4pub mod runtime_config;
5mod types;
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use allowed_hosts::AllowedHostChecker;
11use client::ClientFactory;
12use runtime_config::RuntimeConfig;
13use spin_factor_otel::OtelFactorState;
14use spin_factor_outbound_networking::{
15    ConnectionSemaphore, OutboundNetworkingFactor, build_connection_semaphore,
16};
17use spin_factors::{
18    ConfigureAppContext, Factor, PrepareContext, RuntimeFactors, SelfInstanceBuilder, anyhow,
19};
20
21pub struct OutboundPgFactor<CF = crate::client::PooledTokioClientFactory> {
22    _phantom: std::marker::PhantomData<CF>,
23}
24
25pub struct AppState<CF> {
26    pub client_factories: HashMap<String, Arc<CF>>,
27    /// Semaphore to limit concurrent outbound PostgreSQL connections.
28    pub semaphore: ConnectionSemaphore,
29}
30
31impl<CF: ClientFactory> Factor for OutboundPgFactor<CF> {
32    type RuntimeConfig = RuntimeConfig;
33    type AppState = AppState<CF>;
34    type InstanceBuilder = InstanceState<CF>;
35
36    fn init<T: spin_factors::InitContext<Self>>(&mut self, ctx: &mut T) -> anyhow::Result<()> {
37        ctx.link_bindings(spin_world::v1::postgres::add_to_linker::<_, PgFactorData<CF>>)?;
38        ctx.link_bindings(spin_world::v2::postgres::add_to_linker::<_, PgFactorData<CF>>)?;
39        ctx.link_bindings(
40            spin_world::spin::postgres3_0_0::postgres::add_to_linker::<_, PgFactorData<CF>>,
41        )?;
42        ctx.link_bindings(
43            spin_world::spin::postgres4_2_0::postgres::add_to_linker::<_, PgFactorData<CF>>,
44        )?;
45        Ok(())
46    }
47
48    fn configure_app<T: RuntimeFactors>(
49        &self,
50        mut ctx: ConfigureAppContext<T, Self>,
51    ) -> anyhow::Result<Self::AppState> {
52        let config = ctx.take_runtime_config().unwrap_or_default();
53        let mut client_factories = HashMap::new();
54        for comp in ctx.app().components() {
55            client_factories.insert(comp.id().to_string(), Arc::new(CF::default()));
56        }
57        let networking = ctx.app_state::<OutboundNetworkingFactor>().ok();
58
59        Ok(AppState {
60            client_factories,
61            semaphore: build_connection_semaphore(
62                networking,
63                "pg",
64                config.max_connections,
65                config.wait_timeout,
66            ),
67        })
68    }
69
70    fn prepare<T: RuntimeFactors>(
71        &self,
72        mut ctx: PrepareContext<T, Self>,
73    ) -> anyhow::Result<Self::InstanceBuilder> {
74        let allowed_hosts = ctx
75            .instance_builder::<OutboundNetworkingFactor>()?
76            .allowed_hosts();
77        let otel = OtelFactorState::from_prepare_context(&mut ctx)?;
78        let cf = ctx
79            .app_state()
80            .client_factories
81            .get(ctx.app_component().id())
82            .unwrap();
83
84        Ok(InstanceState {
85            allowed_host_checker: AllowedHostChecker::new(allowed_hosts),
86            client_factory: cf.clone(),
87            connections: Default::default(),
88            otel,
89            builders: Default::default(),
90            semaphore: ctx.app_state().semaphore.clone(),
91        })
92    }
93}
94
95impl<C> Default for OutboundPgFactor<C> {
96    fn default() -> Self {
97        Self {
98            _phantom: Default::default(),
99        }
100    }
101}
102
103impl<C> OutboundPgFactor<C> {
104    pub fn new() -> Self {
105        Self::default()
106    }
107}
108
109pub struct InstanceState<CF: ClientFactory> {
110    allowed_host_checker: AllowedHostChecker,
111    client_factory: Arc<CF>,
112    connections: spin_resource_table::Table<(
113        CF::Client,
114        spin_factor_outbound_networking::ConnectionPermit,
115    )>,
116    otel: OtelFactorState,
117    builders: spin_resource_table::Table<host::ConnectionBuilder>,
118    pub semaphore: ConnectionSemaphore,
119}
120
121impl<CF: ClientFactory> SelfInstanceBuilder for InstanceState<CF> {}
122
123pub struct PgFactorData<CF: ClientFactory>(OutboundPgFactor<CF>);
124
125impl<CF: ClientFactory> spin_core::wasmtime::component::HasData for PgFactorData<CF> {
126    type Data<'a> = &'a mut InstanceState<CF>;
127}
128
129impl<CF: ClientFactory> spin_core::wasmtime::component::HasData for InstanceState<CF> {
130    type Data<'a> = &'a mut InstanceState<CF>;
131}
spin_factor_outbound_pg/lib.rs

spin_factor_outbound_pg/
lib.rs
