TLA  Line data    Source code

       1                 : use anyhow::{anyhow, Ok, Result};
       2                 : use postgres::Client;
       3                 : use tokio_postgres::NoTls;
       4                 : use tracing::{error, instrument, warn};
       5                 : 
       6                 : use crate::compute::ComputeNode;
       7                 : 
       8                 : /// Create a special service table for availability checks
       9                 : /// only if it does not exist already.
      10 CBC           1 : pub fn create_availability_check_data(client: &mut Client) -> Result<()> {
      11               1 :     let query = "
      12               1 :         DO $$
      13               1 :         BEGIN
      14               1 :             IF NOT EXISTS(
      15               1 :                 SELECT 1
      16               1 :                 FROM pg_catalog.pg_tables
      17               1 :                 WHERE tablename = 'health_check'
      18               1 :             )
      19               1 :             THEN
      20               1 :             CREATE TABLE health_check (
      21               1 :                 id serial primary key,
      22               1 :                 updated_at timestamptz default now()
      23               1 :             );
      24               1 :             INSERT INTO health_check VALUES (1, now())
      25               1 :                 ON CONFLICT (id) DO UPDATE
      26               1 :                  SET updated_at = now();
      27               1 :             END IF;
      28               1 :         END
      29               1 :         $$;";
      30               1 :     client.execute(query, &[])?;
      31                 : 
      32               1 :     Ok(())
      33               1 : }
      34                 : 
      35                 : /// Update timestamp in a row in a special service table to check
      36                 : /// that we can actually write some data in this particular timeline.
      37 UBC           0 : #[instrument(skip_all)]
      38                 : pub async fn check_writability(compute: &ComputeNode) -> Result<()> {
      39                 :     // Connect to the database.
      40                 :     let (client, connection) = tokio_postgres::connect(compute.connstr.as_str(), NoTls).await?;
      41                 :     if client.is_closed() {
      42                 :         return Err(anyhow!("connection to postgres closed"));
      43                 :     }
      44                 : 
      45                 :     // The connection object performs the actual communication with the database,
      46                 :     // so spawn it off to run on its own.
      47               0 :     tokio::spawn(async move {
      48               0 :         if let Err(e) = connection.await {
      49               0 :             error!("connection error: {}", e);
      50               0 :         }
      51               0 :     });
      52                 : 
      53                 :     let query = "
      54                 :     INSERT INTO health_check VALUES (1, now())
      55                 :         ON CONFLICT (id) DO UPDATE
      56                 :          SET updated_at = now();";
      57                 : 
      58                 :     match client.simple_query(query).await {
      59                 :         Result::Ok(result) => {
      60                 :             if result.len() != 1 {
      61                 :                 return Err(anyhow::anyhow!(
      62                 :                     "expected 1 query results, but got {}",
      63                 :                     result.len()
      64                 :                 ));
      65                 :             }
      66                 :         }
      67                 :         Err(err) => {
      68                 :             if let Some(state) = err.code() {
      69                 :                 if state == &tokio_postgres::error::SqlState::DISK_FULL {
      70               0 :                     warn!("Tenant disk is full");
      71                 :                     return Ok(());
      72                 :                 }
      73                 :             }
      74                 :             return Err(err.into());
      75                 :         }
      76                 :     }
      77                 : 
      78                 :     Ok(())
      79                 : }