Skip to main content

micromegas_analytics/lakehouse/
view_definition_store.rs

1//! Postgres-backed persistence for `lakehouse_view_set_definitions`.
2//!
3//! [`ViewDefinitionStore`] is reduced to the single pool-backed [`ViewDefinitionStore::list`]
4//! method [`super::view_registry::ViewRegistry::reload`] needs -- its only implementors are
5//! [`PgViewDefinitionStore`] and a test fake. The DDL executor
6//! (`rust/public/src/servers/view_ddl.rs`) instead calls the free functions below directly,
7//! against its own open `sqlx::Transaction`, so a `CREATE`/`DROP`'s existence check, upsert/delete
8//! and dependent-protection re-read all run inside one transaction, alongside `retire_partitions`
9//! on a `DROP` -- with no trait indirection for methods a test fake could never implement.
10
11use super::view_definition::{ViewDefinition, ViewOptions};
12use anyhow::{Context, Result};
13use async_trait::async_trait;
14use chrono::{DateTime, Utc};
15use sqlx::Row;
16use std::fmt::Debug;
17
18/// One row of `lakehouse_view_set_definitions`, exactly as read back. `view_options` is kept as
19/// its raw JSON text here (parsed into a [`ViewOptions`] only by [`ViewDefinitionRow::into_definition`])
20/// so a row that fails to parse can still be reported by name rather than losing its identity.
21#[derive(Debug, Clone)]
22pub struct ViewDefinitionRow {
23    pub view_set_name: String,
24    pub definition_sql: String,
25    pub extract_query: String,
26    pub count_src_query: String,
27    pub merge_partitions_query: String,
28    pub update_group: i32,
29    pub view_options: String,
30    pub created_at: DateTime<Utc>,
31    pub updated_at: DateTime<Utc>,
32    pub updated_by: Option<String>,
33}
34
35impl ViewDefinitionRow {
36    /// Parses `view_options` and assembles the normalized [`ViewDefinition`] the DDL executor and
37    /// registry loader both validate and build a `SqlBatchView` from.
38    pub fn into_definition(self) -> Result<ViewDefinition> {
39        let options: ViewOptions = serde_json::from_str(&self.view_options).with_context(|| {
40            format!(
41                "'{}': parsing view_options {:?}",
42                self.view_set_name, self.view_options
43            )
44        })?;
45        Ok(ViewDefinition {
46            view_set_name: self.view_set_name,
47            extract_query: self.extract_query,
48            count_src_query: self.count_src_query,
49            merge_partitions_query: self.merge_partitions_query,
50            update_group: self.update_group,
51            options,
52        })
53    }
54}
55
56fn row_from_sqlx(row: sqlx::postgres::PgRow) -> Result<ViewDefinitionRow> {
57    Ok(ViewDefinitionRow {
58        view_set_name: row.try_get("view_set_name")?,
59        definition_sql: row.try_get("definition_sql")?,
60        extract_query: row.try_get("extract_query")?,
61        count_src_query: row.try_get("count_src_query")?,
62        merge_partitions_query: row.try_get("merge_partitions_query")?,
63        update_group: row.try_get("update_group")?,
64        view_options: row.try_get("view_options")?,
65        created_at: row.try_get("created_at")?,
66        updated_at: row.try_get("updated_at")?,
67        updated_by: row.try_get("updated_by")?,
68    })
69}
70
71const SELECT_COLUMNS: &str = "view_set_name, definition_sql, extract_query, count_src_query, \
72     merge_partitions_query, update_group, view_options, created_at, updated_at, updated_by";
73
74/// The rows-in seam [`super::view_registry::ViewRegistry::reload`] needs: every definition
75/// currently in force, ordered by `(update_group, view_set_name)` -- the order
76/// `super::view_registry::build_factory` relies on regardless of which caller fed it the rows.
77#[async_trait]
78pub trait ViewDefinitionStore: Send + Sync + Debug {
79    async fn list(&self) -> Result<Vec<ViewDefinitionRow>>;
80}
81
82/// The production, pool-backed implementation.
83#[derive(Debug)]
84pub struct PgViewDefinitionStore {
85    pool: sqlx::Pool<sqlx::Postgres>,
86}
87
88impl PgViewDefinitionStore {
89    pub fn new(pool: sqlx::Pool<sqlx::Postgres>) -> Self {
90        Self { pool }
91    }
92}
93
94#[async_trait]
95impl ViewDefinitionStore for PgViewDefinitionStore {
96    async fn list(&self) -> Result<Vec<ViewDefinitionRow>> {
97        let sql = format!(
98            "SELECT {SELECT_COLUMNS} FROM lakehouse_view_set_definitions \
99             ORDER BY update_group, view_set_name"
100        );
101        let rows = sqlx::query(&sql)
102            .fetch_all(&self.pool)
103            .await
104            .with_context(|| "listing lakehouse_view_set_definitions")?;
105        rows.into_iter().map(row_from_sqlx).collect()
106    }
107}
108
109/// The transaction-scoped counterpart of [`ViewDefinitionStore::list`], for a caller (the DDL
110/// executor) already holding an open `sqlx::Transaction` -- so a pre-mutation read cannot miss a
111/// concurrent, uncommitted write from a sibling transaction, and a post-mutation read sees this
112/// transaction's own write.
113pub async fn list_tx(
114    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
115) -> Result<Vec<ViewDefinitionRow>> {
116    let sql = format!(
117        "SELECT {SELECT_COLUMNS} FROM lakehouse_view_set_definitions \
118         ORDER BY update_group, view_set_name"
119    );
120    let rows = sqlx::query(&sql)
121        .fetch_all(&mut **tx)
122        .await
123        .with_context(|| "listing lakehouse_view_set_definitions (tx)")?;
124    rows.into_iter().map(row_from_sqlx).collect()
125}
126
127/// Inserts a new row, or -- on a name collision -- updates every column including `updated_at`/
128/// `updated_by` explicitly: the column defaults only fire on `INSERT`, and `reload()`'s digest is
129/// over `(view_set_name, updated_at)`, so a `REPLACE` that left it unset would go live only on the
130/// replacing node while every other replica kept serving/materializing the old definition.
131pub async fn upsert_tx(
132    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
133    def: &ViewDefinition,
134    definition_sql: &str,
135    updated_by: &str,
136) -> Result<()> {
137    let view_options = serde_json::to_string(&def.options)
138        .with_context(|| format!("serializing view_options for '{}'", def.view_set_name))?;
139    sqlx::query(
140        "INSERT INTO lakehouse_view_set_definitions
141             (view_set_name, definition_sql, extract_query, count_src_query,
142              merge_partitions_query, update_group, view_options, updated_at, updated_by)
143         VALUES ($1, $2, $3, $4, $5, $6, $7, now(), $8)
144         ON CONFLICT (view_set_name) DO UPDATE SET
145             definition_sql = EXCLUDED.definition_sql,
146             extract_query = EXCLUDED.extract_query,
147             count_src_query = EXCLUDED.count_src_query,
148             merge_partitions_query = EXCLUDED.merge_partitions_query,
149             update_group = EXCLUDED.update_group,
150             view_options = EXCLUDED.view_options,
151             updated_at = now(),
152             updated_by = EXCLUDED.updated_by;",
153    )
154    .bind(&def.view_set_name)
155    .bind(definition_sql)
156    .bind(&def.extract_query)
157    .bind(&def.count_src_query)
158    .bind(&def.merge_partitions_query)
159    .bind(def.update_group)
160    .bind(view_options)
161    .bind(updated_by)
162    .execute(&mut **tx)
163    .await
164    .with_context(|| format!("upserting '{}'", def.view_set_name))?;
165    Ok(())
166}
167
168/// Deletes a row by name; `true` if a row was actually deleted.
169pub async fn delete_tx(
170    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
171    view_set_name: &str,
172) -> Result<bool> {
173    let result = sqlx::query("DELETE FROM lakehouse_view_set_definitions WHERE view_set_name = $1")
174        .bind(view_set_name)
175        .execute(&mut **tx)
176        .await
177        .with_context(|| format!("deleting '{view_set_name}'"))?;
178    Ok(result.rows_affected() > 0)
179}
180
181/// The exact insert-time range currently materialized for `view_set_name`'s `'global'` instance --
182/// `None` when it has no partitions -- for `DROP`'s call to `retire_partitions`. Deliberately
183/// exact rather than an unbounded range: `retire_partitions` is hash-agnostic but still needs an
184/// explicit range to key its deletion on.
185pub async fn partition_insert_range(
186    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
187    view_set_name: &str,
188) -> Result<Option<(DateTime<Utc>, DateTime<Utc>)>> {
189    let row = sqlx::query(
190        "SELECT min(begin_insert_time) AS begin, max(end_insert_time) AS end
191         FROM lakehouse_partitions
192         WHERE view_set_name = $1 AND view_instance_id = 'global';",
193    )
194    .bind(view_set_name)
195    .fetch_one(&mut **tx)
196    .await
197    .with_context(|| format!("reading partition insert-time range for '{view_set_name}'"))?;
198    let begin: Option<DateTime<Utc>> = row.try_get("begin")?;
199    let end: Option<DateTime<Utc>> = row.try_get("end")?;
200    Ok(begin.zip(end))
201}