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            0 : pub fn create_availability_check_data(client: &mut Client) -> Result<()> {
      11            0 :     let query = "
      12            0 :         DO $$
      13            0 :         BEGIN
      14            0 :             IF NOT EXISTS(
      15            0 :                 SELECT 1
      16            0 :                 FROM pg_catalog.pg_tables
      17            0 :                 WHERE tablename = 'health_check'
      18            0 :             )
      19            0 :             THEN
      20            0 :             CREATE TABLE health_check (
      21            0 :                 id serial primary key,
      22            0 :                 updated_at timestamptz default now()
      23            0 :             );
      24            0 :             INSERT INTO health_check VALUES (1, now())
      25            0 :                 ON CONFLICT (id) DO UPDATE
      26            0 :                  SET updated_at = now();
      27            0 :             END IF;
      28            0 :         END
      29            0 :         $$;";
      30            0 :     client.execute(query, &[])?;
      31              : 
      32            0 :     Ok(())
      33            0 : }
      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            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              :                     warn!("Tenant disk is full");
      71              :                     return Ok(());
      72              :                 }
      73              :             }
      74              :             return Err(err.into());
      75              :         }
      76              :     }
      77              : 
      78              :     Ok(())
      79              : }