rls

backend/src/db.rs

@@ -1,12 +1,40 @@
 use std::time::Duration;
 
 use diesel::pg::PgConnection;
-use diesel::r2d2::{ConnectionManager, Pool};
+use diesel::r2d2::{ConnectionManager, CustomizeConnection, Pool};
+use diesel::RunQueryDsl;
 
 pub type PgPool = Pool<ConnectionManager<PgConnection>>;
 
 pub const DEFAULT_MAX_POOL_SIZE: u32 = 2;
 
+#[derive(Debug)]
+struct SchemaCustomizer;
+
+impl CustomizeConnection<PgConnection, diesel::r2d2::Error> for SchemaCustomizer {
+    fn on_acquire(&self, conn: &mut PgConnection) -> Result<(), diesel::r2d2::Error> {
+        diesel::sql_query(
+            "SELECT set_config('search_path', (
+                SELECT string_agg(schema_name, ', ')
+                FROM (
+                    SELECT 'tenant' AS schema_name WHERE EXISTS (
+                        SELECT 1 FROM pg_namespace WHERE nspname = 'tenant'
+                    )
+                    UNION ALL
+                    SELECT 'shared' AS schema_name WHERE EXISTS (
+                        SELECT 1 FROM pg_namespace WHERE nspname = 'shared'
+                    )
+                    UNION ALL
+                    SELECT 'public' AS schema_name
+                ) AS schemas
+            ), false)",
+        )
+        .execute(conn)
+        .map(|_| ())
+        .map_err(diesel::r2d2::Error::QueryError)
+    }
+}
+
 pub fn init_pool(database_url: &str) -> anyhow::Result<PgPool> {
     init_pool_with_size(database_url, DEFAULT_MAX_POOL_SIZE)
 }
@@ -17,6 +45,7 @@ pub fn init_pool_with_size(database_url: &str, max_size: u32) -> anyhow::Result<
     let pool = Pool::builder()
         .max_size(pool_size)
         .connection_timeout(Duration::from_secs(10))
+        .connection_customizer(Box::new(SchemaCustomizer))
         .build(manager)?;
     Ok(pool)
 }