Skip to main content

spin_factor_outbound_pg/
lib.rs

1mod allowed_hosts;
2pub mod client;
3mod host;
4mod types;
5
6use std::{collections::HashMap, sync::Arc};
7
8use allowed_hosts::AllowedHostChecker;
9use client::ClientFactory;
10use spin_factor_otel::OtelFactorState;
11use spin_factor_outbound_networking::OutboundNetworkingFactor;
12use spin_factors::{
13    ConfigureAppContext, Factor, PrepareContext, RuntimeFactors, SelfInstanceBuilder, anyhow,
14};
15
16pub struct OutboundPgFactor<CF = crate::client::PooledTokioClientFactory> {
17    _phantom: std::marker::PhantomData<CF>,
18}
19
20impl<CF: ClientFactory> Factor for OutboundPgFactor<CF> {
21    type RuntimeConfig = ();
22    type AppState = HashMap<String, Arc<CF>>;
23    type InstanceBuilder = InstanceState<CF>;
24
25    fn init(&mut self, ctx: &mut impl spin_factors::InitContext<Self>) -> anyhow::Result<()> {
26        ctx.link_bindings(spin_world::v1::postgres::add_to_linker::<_, PgFactorData<CF>>)?;
27        ctx.link_bindings(spin_world::v2::postgres::add_to_linker::<_, PgFactorData<CF>>)?;
28        ctx.link_bindings(
29            spin_world::spin::postgres3_0_0::postgres::add_to_linker::<_, PgFactorData<CF>>,
30        )?;
31        ctx.link_bindings(
32            spin_world::spin::postgres4_2_0::postgres::add_to_linker::<_, PgFactorData<CF>>,
33        )?;
34        Ok(())
35    }
36
37    fn configure_app<T: RuntimeFactors>(
38        &self,
39        ctx: ConfigureAppContext<T, Self>,
40    ) -> anyhow::Result<Self::AppState> {
41        let mut client_factories = HashMap::new();
42        for comp in ctx.app().components() {
43            client_factories.insert(comp.id().to_string(), Arc::new(CF::default()));
44        }
45        Ok(client_factories)
46    }
47
48    fn prepare<T: RuntimeFactors>(
49        &self,
50        mut ctx: PrepareContext<T, Self>,
51    ) -> anyhow::Result<Self::InstanceBuilder> {
52        let allowed_hosts = ctx
53            .instance_builder::<OutboundNetworkingFactor>()?
54            .allowed_hosts();
55        let otel = OtelFactorState::from_prepare_context(&mut ctx)?;
56        let cf = ctx.app_state().get(ctx.app_component().id()).unwrap();
57
58        Ok(InstanceState {
59            allowed_host_checker: AllowedHostChecker::new(allowed_hosts),
60            client_factory: cf.clone(),
61            connections: Default::default(),
62            otel,
63            builders: Default::default(),
64        })
65    }
66}
67
68impl<C> Default for OutboundPgFactor<C> {
69    fn default() -> Self {
70        Self {
71            _phantom: Default::default(),
72        }
73    }
74}
75
76impl<C> OutboundPgFactor<C> {
77    pub fn new() -> Self {
78        Self::default()
79    }
80}
81
82pub struct InstanceState<CF: ClientFactory> {
83    allowed_host_checker: AllowedHostChecker,
84    client_factory: Arc<CF>,
85    connections: spin_resource_table::Table<CF::Client>,
86    otel: OtelFactorState,
87    builders: spin_resource_table::Table<host::ConnectionBuilder>,
88}
89
90impl<CF: ClientFactory> SelfInstanceBuilder for InstanceState<CF> {}
91
92pub struct PgFactorData<CF: ClientFactory>(OutboundPgFactor<CF>);
93
94impl<CF: ClientFactory> spin_core::wasmtime::component::HasData for PgFactorData<CF> {
95    type Data<'a> = &'a mut InstanceState<CF>;
96}
97
98impl<CF: ClientFactory> spin_core::wasmtime::component::HasData for InstanceState<CF> {
99    type Data<'a> = &'a mut InstanceState<CF>;
100}