micromegas_analytics/lakehouse/
view_definition_store.rs1use 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#[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 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#[async_trait]
78pub trait ViewDefinitionStore: Send + Sync + Debug {
79 async fn list(&self) -> Result<Vec<ViewDefinitionRow>>;
80}
81
82#[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
109pub 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
127pub 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
168pub 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
181pub 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}