Repository: MystenLabs/sui
File: crates/sui-pg-db/src/lib.rs (Rust)

Code before:
    }

    /// Construct a new DB connection pool talking to the database at `database_url` that defaults
    /// to read-only transactions. Instances of [Db] can be cloned to share access to the same
    /// pool.
    pub async fn for_read(database_url: Url, config: DbArgs) -> anyhow::Result<Self> {
        Self::from_pool(database_url, config, true).await
    }

    async fn from_pool(database_url: Url, config: DbArgs, read_only: bool) -> anyhow::Result<Self> {
        Ok(Db {
            pool: pool(database_url, config, read_only).await?,
            pool_stats: Default::default(),
        })
    }

    /// Retrieves a connection from the pool. Can fail with a timeout if a connection cannot be
    /// established before the [DbArgs::connection_timeout] has elapsed.
    pub async fn connect(&self) -> anyhow::Result<Connection<'_>> {
        self.pool_stats.requested.fetch_add(1, Ordering::Relaxed);
        let future = PoolFuture {
            completed: Default::default(),
            pool_stats: self.pool_stats.clone(),
            inner: self.pool.get(),
        };
        let completed = future.completed.clone();
        match future.await {
            Ok(c) => {
                inc_if_false(&completed, &self.pool_stats.acquired);
                Ok(Connection(c))
            }
            Err(e) => {
                inc_if_false(&completed, &self.pool_stats.unacquired_error);
                Err(e.into())
            }
        }
    }

    /// Statistics about the connection pool
    pub fn state(&self) -> bb8::State {
        self.pool.state()
    }

    pub fn pool_stats(&self) -> &PoolStats {
        self.pool_stats.as_ref()
    }

    async fn clear_database(&self) -> anyhow::Result<()> {
        use diesel_async::RunQueryDsl;

        info!("Clearing the database...");

Review comment (bug):
I think we can avoid the `Future impl`, `AtomicBool`, `pin_project` with some kind of drop guard? May be incorrect but I think when a future is canceled mid-await Rust will still call drop

Could use scopeguard or something homemade

```
  struct CancelGuard<'a> {                                                                                                                                                                                                                                                                                                                                                   
      stats: &'a PoolStats,                                                                                                                                                                                                                                                                                                                                                  
      disarmed: bool,                                                                                                                                                                                                                                                                                                                                                        
  }                                                                                                                                                                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                                                                                                                             
  impl Drop for CancelGuard<'_> {                                                                                                                                                                                                                                                                                                                                            
      fn drop(&mut self) {                                                                                                                                                                                                                                                                                                                                                   
          if !self.disarmed {                                                                                                                                                                                                                                                                                                                                                
              self.stats.unacquired_canceled.fetch_add(1, Ordering::Relaxed);                                                                                                                                                                                                                                                                                                
          }                                                                                                                                                                                                                                                                                                                                                                  
      }                                                                                                                                                                                                                                                                                                                                                                      
  }                                                                                                                                                                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                                                                                                                             
  pub async fn connect(&self) -> anyhow::Result<Connection<'_>> {                                                                                                                                                                                                                                                                                                            
      self.pool_stats.requested.fetch_add(1, Ordering::Relaxed);                                                                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                                                                                                                             
      let mut guard = CancelGuard { stats: &self.pool_stats, disarmed: false };                                                                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                                                                                                                             
      match self.pool.get().await {                                                                                                                                                                                                                                                                                                                                          
          Ok(c) => {                                                                                                                                                                                                                                                                                                                                                         
              self.pool_stats.acquired.fetch_add(1, Ordering::Relaxed);                                                                                                                                                                                                                                                                                                      
              guard.disarmed = true;                                                                                                                                                                                                                                                                                                                                         
              Ok(Connection(c))                                                                                                                                                                                                                                                                                                                                              
          }                                                                                                                                                                                                                                                                                                                                                                  
          Err(e) => {                                                                                                                                                                                                                                                                                                                                                        
              self.pool_stats.unacquired_error.fetch_add(1, Ordering::Relaxed);                                                                                                                                                                                                                                                                                              
              guard.disarmed = true;                                                                                                                                                                                                                                                                                                                                         
              Err(e.into())                                                                                                                                                                                                                                                                                                                                                  
          }                                                                                                                                                                                                                                                                                                                                                                  
      }                                                                                                                                                                                                                                                                                                                                                                      
  }
```




Code after:
    }

    async fn clear_database(&self) -> anyhow::Result<()> {
        info!("Clearing the database...");
        let mut conn = self.connect().await?;
        let drop_all_tables = "
        DO $$ DECLARE
            r RECORD;
        BEGIN
        FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public')
            LOOP
                EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
            END LOOP;
        END $$;";
        diesel::sql_query(drop_all_tables)
            .execute(&mut conn)
            .await?;
        info!("Dropped all tables.");

        let drop_all_procedures = "
        DO $$ DECLARE
            r RECORD;
        BEGIN
            FOR r IN (SELECT proname, oidvectortypes(proargtypes) as argtypes
                      FROM pg_proc INNER JOIN pg_namespace ns ON (pg_proc.pronamespace = ns.oid)
                      WHERE ns.nspname = 'public' AND prokind = 'p')
            LOOP
                EXECUTE 'DROP PROCEDURE IF EXISTS ' || quote_ident(r.proname) || '(' || r.argtypes || ') CASCADE';
            END LOOP;
        END $$;";
        diesel::sql_query(drop_all_procedures)
            .execute(&mut conn)
            .await?;
        info!("Dropped all procedures.");

        let drop_all_functions = "
        DO $$ DECLARE
            r RECORD;
        BEGIN
            FOR r IN (SELECT proname, oidvectortypes(proargtypes) as argtypes
                      FROM pg_proc INNER JOIN pg_namespace ON (pg_proc.pronamespace = pg_namespace.oid)
                      WHERE pg_namespace.nspname = 'public' AND prokind = 'f')
            LOOP
                EXECUTE 'DROP FUNCTION IF EXISTS ' || quote_ident(r.proname) || '(' || r.argtypes || ') CASCADE';
            END LOOP;
        END $$;";
        diesel::sql_query(drop_all_functions)
            .execute(&mut conn)
            .await?;
        info!("Database cleared.");
        Ok(())