store: Use AsyncPgConnection in remaining places

Author: lutter

store/postgres/src/advisory_lock.rs

@@ -20,7 +20,6 @@ use diesel_async::RunQueryDsl;
 use graph::prelude::StoreError;
 
 use crate::command_support::catalog::Site;
-use crate::pool::PgConnection;
 use crate::primary::DeploymentId;
 use crate::AsyncPgConnection;
 
@@ -57,7 +56,7 @@ impl Scope {
 
     /// Lock the deployment in this scope with the given id. Blocks until we
     /// can get the lock
-    async fn lock(&self, conn: &mut PgConnection, id: DeploymentId) -> Result<(), StoreError> {
+    async fn lock(&self, conn: &mut AsyncPgConnection, id: DeploymentId) -> Result<(), StoreError> {
         sql_query(format!("select pg_advisory_lock({}, {id})", self.id))
             .execute(conn)
             .await
@@ -120,12 +119,18 @@ where
 
 /// Take the lock used to keep two copy operations to run simultaneously on
 /// the same deployment. Block until we can get the lock
-pub(crate) async fn lock_copying(conn: &mut PgConnection, dst: &Site) -> Result<(), StoreError> {
+pub(crate) async fn lock_copying(
+    conn: &mut AsyncPgConnection,
+    dst: &Site,
+) -> Result<(), StoreError> {
     COPY.lock(conn, dst.id).await
 }
 
 /// Release the lock acquired with `lock_copying`.
-pub(crate) async fn unlock_copying(conn: &mut PgConnection, dst: &Site) -> Result<(), StoreError> {
+pub(crate) async fn unlock_copying(
+    conn: &mut AsyncPgConnection,
+    dst: &Site,
+) -> Result<(), StoreError> {
     COPY.unlock(conn, dst.id).await
 }
 

store/postgres/src/block_store.rs

@@ -25,9 +25,9 @@ use graph::{prelude::StoreError, util::timed_cache::TimedCache};
 use crate::{
     chain_head_listener::ChainHeadUpdateSender,
     chain_store::{ChainStoreMetrics, Storage},
-    pool::{ConnectionPool, PgConnection},
+    pool::ConnectionPool,
     primary::Mirror as PrimaryMirror,
-    ChainStore, NotificationSender, Shard, PRIMARY_SHARD,
+    AsyncPgConnection, ChainStore, NotificationSender, Shard, PRIMARY_SHARD,
 };
 
 use self::primary::Chain;
@@ -58,7 +58,7 @@ pub mod primary {
         prelude::StoreError,
     };
 
-    use crate::{chain_store::Storage, pool::PgConnection, AsyncPgConnection};
+    use crate::{chain_store::Storage, AsyncPgConnection};
     use crate::{ConnectionPool, Shard};
 
     table! {
@@ -119,7 +119,7 @@ pub mod primary {
     }
 
     pub async fn add_chain(
-        conn: &mut PgConnection,
+        conn: &mut AsyncPgConnection,
         name: &str,
         shard: &Shard,
         ident: ChainIdentifier,
@@ -174,7 +174,7 @@ pub mod primary {
 
     // update chain name where chain name is 'name'
     pub async fn update_chain_name(
-        conn: &mut PgConnection,
+        conn: &mut AsyncPgConnection,
         name: &str,
         new_name: &str,
     ) -> Result<(), StoreError> {
@@ -344,7 +344,7 @@ impl BlockStore {
     }
 
     pub async fn allocate_chain(
-        conn: &mut PgConnection,
+        conn: &mut AsyncPgConnection,
         name: &String,
         shard: &Shard,
         ident: &ChainIdentifier,

store/postgres/src/catalog.rs

@@ -23,7 +23,6 @@ use graph::{
     prelude::{lazy_static, StoreError, BLOCK_NUMBER_MAX},
 };
 
-use crate::pool::PgConnection;
 use crate::AsyncPgConnection;
 use crate::{
     block_range::BLOCK_RANGE_COLUMN,
@@ -578,7 +577,7 @@ pub async fn drop_schema(conn: &mut AsyncPgConnection, nsp: &str) -> Result<(),
     Ok(conn.batch_execute(&query).await?)
 }
 
-pub async fn migration_count(conn: &mut PgConnection) -> Result<usize, StoreError> {
+pub async fn migration_count(conn: &mut AsyncPgConnection) -> Result<usize, StoreError> {
     use __diesel_schema_migrations as m;
 
     if !table_exists(conn, NAMESPACE_PUBLIC, &MIGRATIONS_TABLE).await? {

store/postgres/src/deployment.rs

@@ -1,10 +1,7 @@
 //! Utilities for dealing with deployment metadata. Any connection passed
 //! into these methods must be for the shard that holds the actual
 //! deployment data and metadata
-use crate::{
-    advisory_lock, detail::GraphNodeVersion, pool::PgConnection, primary::DeploymentId,
-    AsyncPgConnection,
-};
+use crate::{advisory_lock, detail::GraphNodeVersion, primary::DeploymentId, AsyncPgConnection};
 use diesel::{
     dsl::{count, delete, insert_into, now, select, sql, update},
     sql_types::{Bool, Integer},
@@ -989,7 +986,7 @@ pub(crate) async fn insert_subgraph_errors(
 
 #[cfg(debug_assertions)]
 pub(crate) async fn error_count(
-    conn: &mut PgConnection,
+    conn: &mut AsyncPgConnection,
     id: &DeploymentHash,
 ) -> Result<usize, StoreError> {
     use subgraph_error as e;
@@ -1116,7 +1113,7 @@ pub(crate) async fn revert_subgraph_errors(
 }
 
 pub(crate) async fn delete_error(
-    conn: &mut PgConnection,
+    conn: &mut AsyncPgConnection,
     error_id: &str,
 ) -> Result<(), StoreError> {
     use subgraph_error as e;
@@ -1130,7 +1127,7 @@ pub(crate) async fn delete_error(
 /// Copy the dynamic data sources for `src` to `dst`. All data sources that
 /// were created up to and including `target_block` will be copied.
 pub(crate) async fn copy_errors(
-    conn: &mut PgConnection,
+    conn: &mut AsyncPgConnection,
     src: &Site,
     dst: &Site,
     target_block: &BlockPtr,
@@ -1390,7 +1387,7 @@ pub async fn set_earliest_block(
 /// go across shards and use the metadata tables mapped into the shard for
 /// `conn` which must be the shard for `dst`
 pub async fn copy_earliest_block(
-    conn: &mut PgConnection,
+    conn: &mut AsyncPgConnection,
     src: &Site,
     dst: &Site,
 ) -> Result<(), StoreError> {

store/postgres/src/primary.rs

@@ -4,7 +4,7 @@
 use crate::{
     block_range::UNVERSIONED_RANGE,
     detail::DeploymentDetail,
-    pool::{PgConnection, PRIMARY_PUBLIC},
+    pool::PRIMARY_PUBLIC,
     subgraph_store::{unused, Shard, PRIMARY_SHARD},
     AsyncPgConnection, ConnectionPool, ForeignServer, NotificationSender,
 };
@@ -2019,7 +2019,7 @@ impl Primary {
 
 /// Return `true` if we deem this installation to be empty, defined as
 /// having no deployments and no subgraph names in the database
-pub async fn is_empty(conn: &mut PgConnection) -> Result<bool, StoreError> {
+pub async fn is_empty(conn: &mut AsyncPgConnection) -> Result<bool, StoreError> {
     use deployment_schemas as ds;
     use subgraph as s;
 

store/postgres/src/relational/index.rs

@@ -18,7 +18,6 @@ use graph::prelude::{
 use crate::block_range::{BLOCK_COLUMN, BLOCK_RANGE_COLUMN};
 use crate::command_support::catalog::Site;
 use crate::deployment_store::DeploymentStore;
-use crate::pool::PgConnection;
 use crate::primary::Namespace;
 use crate::relational::{BYTE_ARRAY_PREFIX_SIZE, STRING_PREFIX_SIZE};
 use crate::{catalog, AsyncPgConnection};
@@ -822,7 +821,7 @@ impl IndexList {
 
     pub async fn recreate_invalid_indexes(
         &self,
-        conn: &mut PgConnection,
+        conn: &mut AsyncPgConnection,
         layout: &Layout,
     ) -> Result<(), StoreError> {
         #[derive(QueryableByName, Debug)]