store: Remove PermittedConnection.into_inner

lutter · lutter · commit 0ff94bdf2cc4 · 2025-12-09T10:25:48.000-08:00
diff --git a/core/graphman/src/commands/deployment/pause.rs b/core/graphman/src/commands/deployment/pause.rs
@@ -37,7 +37,10 @@ pub async fn load_active_deployment(
     primary_pool: ConnectionPool,
     deployment: &DeploymentSelector,
 ) -> Result<ActiveDeployment, PauseDeploymentError> {
-    let mut primary_conn = primary_pool.get().await.map_err(GraphmanError::from)?;
+    let mut primary_conn = primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
 
     let locator = crate::deployment::load_deployment_locator(
         &mut primary_conn,
@@ -76,7 +79,7 @@ pub async fn pause_active_deployment(
     notification_sender: Arc<NotificationSender>,
     active_deployment: ActiveDeployment,
 ) -> Result<(), GraphmanError> {
-    let primary_conn = primary_pool.get().await?;
+    let primary_conn = primary_pool.get_permitted().await?;
     let mut catalog_conn = catalog::Connection::new(primary_conn);
 
     let changes = catalog_conn.pause_subgraph(&active_deployment.site).await?;
diff --git a/core/graphman/src/commands/deployment/reassign.rs b/core/graphman/src/commands/deployment/reassign.rs
@@ -29,7 +29,10 @@ impl Deployment {
         &self,
         primary_pool: ConnectionPool,
     ) -> Result<Option<NodeId>, GraphmanError> {
-        let primary_conn = primary_pool.get().await.map_err(GraphmanError::from)?;
+        let primary_conn = primary_pool
+            .get_permitted()
+            .await
+            .map_err(GraphmanError::from)?;
         let mut catalog_conn = catalog::Connection::new(primary_conn);
         let node = catalog_conn
             .assigned_node(&self.site)
@@ -58,7 +61,10 @@ pub async fn load_deployment(
     primary_pool: ConnectionPool,
     deployment: &DeploymentSelector,
 ) -> Result<Deployment, ReassignDeploymentError> {
-    let mut primary_conn = primary_pool.get().await.map_err(GraphmanError::from)?;
+    let mut primary_conn = primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
 
     let locator = crate::deployment::load_deployment_locator(
         &mut primary_conn,
@@ -87,7 +93,10 @@ pub async fn reassign_deployment(
     node: &NodeId,
     curr_node: Option<NodeId>,
 ) -> Result<ReassignResult, ReassignDeploymentError> {
-    let primary_conn = primary_pool.get().await.map_err(GraphmanError::from)?;
+    let primary_conn = primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
     let mut catalog_conn = catalog::Connection::new(primary_conn);
     let changes: Vec<AssignmentChange> = match &curr_node {
         Some(curr) => {
diff --git a/core/graphman/src/commands/deployment/resume.rs b/core/graphman/src/commands/deployment/resume.rs
@@ -37,7 +37,10 @@ pub async fn load_paused_deployment(
     primary_pool: ConnectionPool,
     deployment: &DeploymentSelector,
 ) -> Result<PausedDeployment, ResumeDeploymentError> {
-    let mut primary_conn = primary_pool.get().await.map_err(GraphmanError::from)?;
+    let mut primary_conn = primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
 
     let locator = crate::deployment::load_deployment_locator(
         &mut primary_conn,
@@ -76,7 +79,7 @@ pub async fn resume_paused_deployment(
     notification_sender: Arc<NotificationSender>,
     paused_deployment: PausedDeployment,
 ) -> Result<(), GraphmanError> {
-    let primary_conn = primary_pool.get().await?;
+    let primary_conn = primary_pool.get_permitted().await?;
     let mut catalog_conn = catalog::Connection::new(primary_conn);
 
     let changes = catalog_conn
diff --git a/core/graphman/src/commands/deployment/unassign.rs b/core/graphman/src/commands/deployment/unassign.rs
@@ -37,7 +37,10 @@ pub async fn load_assigned_deployment(
     primary_pool: ConnectionPool,
     deployment: &DeploymentSelector,
 ) -> Result<AssignedDeployment, UnassignDeploymentError> {
-    let mut primary_conn = primary_pool.get().await.map_err(GraphmanError::from)?;
+    let mut primary_conn = primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
 
     let locator = crate::deployment::load_deployment_locator(
         &mut primary_conn,
@@ -73,7 +76,7 @@ pub async fn unassign_deployment(
     notification_sender: Arc<NotificationSender>,
     deployment: AssignedDeployment,
 ) -> Result<(), GraphmanError> {
-    let primary_conn = primary_pool.get().await?;
+    let primary_conn = primary_pool.get_permitted().await?;
     let mut catalog_conn = catalog::Connection::new(primary_conn);
 
     let changes = catalog_conn.unassign_subgraph(&deployment.site).await?;
diff --git a/node/src/manager/commands/assign.rs b/node/src/manager/commands/assign.rs
@@ -13,7 +13,7 @@ pub async fn unassign(
 ) -> Result<(), Error> {
     let locator = search.locate_unique(&primary).await?;
 
-    let pconn = primary.get().await?;
+    let pconn = primary.get_permitted().await?;
     let mut conn = catalog::Connection::new(pconn);
 
     let site = conn
@@ -38,7 +38,7 @@ pub async fn reassign(
     let node = NodeId::new(node.clone()).map_err(|()| anyhow!("illegal node id `{}`", node))?;
     let locator = search.locate_unique(&primary).await?;
 
-    let pconn = primary.get().await?;
+    let pconn = primary.get_permitted().await?;
     let mut conn = catalog::Connection::new(pconn);
 
     let site = conn
@@ -81,7 +81,7 @@ pub async fn pause_or_resume(
     locator: &DeploymentLocator,
     should_pause: bool,
 ) -> Result<(), Error> {
-    let pconn = primary.get().await?;
+    let pconn = primary.get_permitted().await?;
     let mut conn = catalog::Connection::new(pconn);
 
     let site = conn
diff --git a/node/src/manager/commands/chain.rs b/node/src/manager/commands/chain.rs
@@ -157,8 +157,9 @@ pub async fn info(
 
 pub async fn remove(primary: ConnectionPool, store: BlockStore, name: String) -> Result<(), Error> {
     let sites = {
-        let mut conn =
-            graph_store_postgres::command_support::catalog::Connection::new(primary.get().await?);
+        let mut conn = graph_store_postgres::command_support::catalog::Connection::new(
+            primary.get_permitted().await?,
+        );
         conn.find_sites_for_network(&name).await?
     };
 
diff --git a/node/src/manager/commands/rewind.rs b/node/src/manager/commands/rewind.rs
@@ -78,7 +78,7 @@ pub async fn run(
     if !start_block && (block_hash.is_none() || block_number.is_none()) {
         bail!("--block-hash and --block-number must be specified when --start-block is not set");
     }
-    let pconn = primary.get().await?;
+    let pconn = primary.get_permitted().await?;
     let mut conn = store_catalog::Connection::new(pconn);
 
     let subgraph_store = store.subgraph_store();
diff --git a/node/src/manager/commands/stats.rs b/node/src/manager/commands/stats.rs
@@ -23,7 +23,7 @@ async fn site_and_conn(
     let primary_pool = pools.get(&*PRIMARY_SHARD).unwrap();
     let locator = search.locate_unique(primary_pool).await?;
 
-    let pconn = primary_pool.get().await?;
+    let pconn = primary_pool.get_permitted().await?;
     let mut conn = store_catalog::Connection::new(pconn);
 
     let site = conn
diff --git a/server/graphman/src/resolvers/deployment_mutation/create.rs b/server/graphman/src/resolvers/deployment_mutation/create.rs
@@ -7,7 +7,11 @@ use crate::resolvers::context::GraphmanContext;
 use graphman::GraphmanError;
 
 pub async fn run(ctx: &GraphmanContext, name: &String) -> Result<()> {
-    let primary_pool = ctx.primary_pool.get().await.map_err(GraphmanError::from)?;
+    let primary_pool = ctx
+        .primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
     let mut catalog_conn = catalog::Connection::new(primary_pool);
 
     let name = match SubgraphName::new(name) {
diff --git a/server/graphman/src/resolvers/deployment_mutation/remove.rs b/server/graphman/src/resolvers/deployment_mutation/remove.rs
@@ -7,7 +7,11 @@ use crate::resolvers::context::GraphmanContext;
 use graphman::GraphmanError;
 
 pub async fn run(ctx: &GraphmanContext, name: &String) -> Result<()> {
-    let primary_pool = ctx.primary_pool.get().await.map_err(GraphmanError::from)?;
+    let primary_pool = ctx
+        .primary_pool
+        .get_permitted()
+        .await
+        .map_err(GraphmanError::from)?;
     let mut catalog_conn = catalog::Connection::new(primary_pool);
 
     let name = match SubgraphName::new(name) {
diff --git a/store/postgres/src/pool/mod.rs b/store/postgres/src/pool/mod.rs
@@ -51,14 +51,6 @@ pub struct PermittedConnection {
     _permit: tokio::sync::OwnedSemaphorePermit,
 }
 
-impl PermittedConnection {
-    /// Returns the underlying connection, consuming the permit.
-    /// The permit is released when this method is called.
-    pub fn into_inner(self) -> AsyncPgConnection {
-        self.conn
-    }
-}
-
 impl Deref for PermittedConnection {
     type Target = AsyncPgConnection;
     fn deref(&self) -> &Self::Target {
diff --git a/store/postgres/src/primary.rs b/store/postgres/src/primary.rs
@@ -4,7 +4,7 @@
 use crate::{
     block_range::UNVERSIONED_RANGE,
     detail::DeploymentDetail,
-    pool::PRIMARY_PUBLIC,
+    pool::{PermittedConnection, PRIMARY_PUBLIC},
     subgraph_store::{unused, Shard, PRIMARY_SHARD},
     AsyncPgConnection, ConnectionPool, ForeignServer, NotificationSender,
 };
@@ -779,11 +779,11 @@ mod queries {
 /// A wrapper for a database connection that provides access to functionality
 /// that works only on the primary database
 pub struct Connection {
-    conn: AsyncPgConnection,
+    conn: PermittedConnection,
 }
 
 impl Connection {
-    pub fn new(conn: AsyncPgConnection) -> Self {
+    pub fn new(conn: PermittedConnection) -> Self {
         Self { conn }
     }
 
@@ -807,13 +807,13 @@ impl Connection {
         type TM = <AsyncPgConnection as diesel_async::AsyncConnection>::TransactionManager;
 
         async move {
-            TM::begin_transaction(&mut self.conn).await?;
+            TM::begin_transaction(&mut *self.conn).await?;
             match callback(self).await {
                 Ok(value) => {
-                    TM::commit_transaction(&mut self.conn).await?;
+                    TM::commit_transaction(&mut *self.conn).await?;
                     Ok(value)
                 }
-                Err(user_error) => match TM::rollback_transaction(&mut self.conn).await {
+                Err(user_error) => match TM::rollback_transaction(&mut *self.conn).await {
                     Ok(()) => Err(user_error),
                     Err(diesel::result::Error::BrokenTransactionManager) => {
                         // In this case we are probably more interested by the
diff --git a/store/postgres/src/subgraph_store.rs b/store/postgres/src/subgraph_store.rs
@@ -819,7 +819,7 @@ impl Inner {
     /// of connections in between getting the first one and trying to get the
     /// second one.
     pub(crate) async fn primary_conn(&self) -> Result<primary::Connection, StoreError> {
-        let conn = self.mirror.primary().get().await?;
+        let conn = self.mirror.primary().get_permitted().await?;
         Ok(primary::Connection::new(conn))
     }
 
@@ -1381,15 +1381,15 @@ impl EnsLookup {
     }
 
     async fn is_table_empty(pool: &ConnectionPool) -> Result<bool, StoreError> {
-        let conn = pool.get_permitted().await?.into_inner();
+        let conn = pool.get_permitted().await?;
         primary::Connection::new(conn).is_ens_table_empty().await
     }
 }
 
 #[async_trait]
 impl EnsLookupTrait for EnsLookup {
     async fn find_name(&self, hash: &str) -> Result<Option<String>, StoreError> {
-        let conn = self.primary.get_permitted().await?.into_inner();
+        let conn = self.primary.get_permitted().await?;
         primary::Connection::new(conn).find_ens_name(hash).await
     }
 
diff --git a/store/test-store/src/store.rs b/store/test-store/src/store.rs
@@ -672,7 +672,7 @@ fn build_store() -> (Arc<Store>, ConnectionPool, Config, Arc<SubscriptionManager
 }
 
 pub async fn primary_connection() -> graph_store_postgres::layout_for_tests::Connection {
-    let conn = PRIMARY_POOL.get().await.unwrap();
+    let conn = PRIMARY_POOL.get_permitted().await.unwrap();
     graph_store_postgres::layout_for_tests::Connection::new(conn)
 }
 

Original file line number	Diff line number	Diff line change
`@@ -78,7 +78,7 @@ pub async fn run(`
`78`	`78`	`if !start_block && (block_hash.is_none() \|\| block_number.is_none()) {`
`79`	`79`	`bail!("--block-hash and --block-number must be specified when --start-block is not set");`
`80`	`80`	`}`
`81`		`- let pconn = primary.get().await?;`
	`81`	`+ let pconn = primary.get_permitted().await?;`
`82`	`82`	`let mut conn = store_catalog::Connection::new(pconn);`
`83`	`83`
`84`	`84`	`let subgraph_store = store.subgraph_store();`