meerkat_mobkit/runtime/
session_store.rs

1//! Session store subsystem — persistence backends and session lifecycle operations.
2
3use super::*;
4
5#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
6#[serde(rename_all = "snake_case")]
7pub enum SessionStoreKind {
8    BigQuery,
9    JsonFile,
10}
11
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct SessionStoreContract {
14    pub store: SessionStoreKind,
15    pub latest_row_per_session: bool,
16    pub tombstones_supported: bool,
17    pub dedup_read_path: bool,
18    pub file_locking: bool,
19    pub crash_recovery: bool,
20    pub bigquery_dataset: Option<String>,
21    pub bigquery_table: Option<String>,
22}
23
24#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
25pub struct SessionPersistenceRow {
26    #[serde(default)]
27    pub session_id: String,
28    #[serde(default)]
29    pub updated_at_ms: u64,
30    #[serde(default)]
31    pub deleted: bool,
32    #[serde(default)]
33    pub payload: Value,
34    #[serde(default)]
35    pub labels: BTreeMap<String, String>,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct JsonStoreLockRecord {
40    pub owner_pid: u32,
41    pub created_at_ms: u64,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub enum JsonFileSessionStoreError {
46    Io(String),
47    Serialize(String),
48    InvalidStoreData(String),
49    LockHeld { lock_path: String },
50    StaleLockRecoveryFailed(String),
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct JsonFileSessionStore {
55    data_path: PathBuf,
56    lock_path: PathBuf,
57    stale_lock_threshold: Duration,
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum BigQuerySessionStoreError {
62    Io(String),
63    Serialize(String),
64    Configuration(String),
65    Http(String),
66    Api(String),
67    InvalidQueryResponse(String),
68    ProcessFailed { command: String, stderr: String },
69}
70
71/// One process-wide HTTP client for the BigQuery adapter. First
72/// `reqwest::Client::build()` in a process performs the system-proxy scan
73/// (macOS SystemConfiguration — ~700ms observed once the `system-proxy`
74/// feature is unified in by the realtime provider chain); sharing the built
75/// client pays that exactly once, off the per-RPC path after first use, and
76/// per-request timeouts replace per-adapter client rebuilds.
77pub(crate) fn shared_bigquery_http_client() -> reqwest::Client {
78    static CLIENT: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
79    CLIENT
80        .get_or_init(|| reqwest::Client::builder().build().unwrap_or_default())
81        .clone()
82}
83
84#[derive(Debug, Clone)]
85pub struct BigQuerySessionStoreAdapter {
86    dataset: String,
87    table: String,
88    project_id: Option<String>,
89    api_base_url: String,
90    access_token: Option<String>,
91    http_timeout: Duration,
92    client: reqwest::Client,
93}
94
95struct JsonFileLockGuard {
96    lock_path: PathBuf,
97}
98
99impl Drop for JsonFileLockGuard {
100    fn drop(&mut self) {
101        let _ = fs::remove_file(&self.lock_path);
102    }
103}
104
105pub fn session_store_contracts(decisions: &RuntimeDecisionState) -> Vec<SessionStoreContract> {
106    vec![
107        SessionStoreContract {
108            store: SessionStoreKind::BigQuery,
109            latest_row_per_session: true,
110            tombstones_supported: true,
111            dedup_read_path: true,
112            file_locking: false,
113            crash_recovery: false,
114            bigquery_dataset: Some(decisions.bigquery.dataset.clone()),
115            bigquery_table: Some(decisions.bigquery.table.clone()),
116        },
117        SessionStoreContract {
118            store: SessionStoreKind::JsonFile,
119            latest_row_per_session: true,
120            tombstones_supported: true,
121            dedup_read_path: true,
122            file_locking: true,
123            crash_recovery: true,
124            bigquery_dataset: None,
125            bigquery_table: None,
126        },
127    ]
128}
129
130pub fn materialize_latest_session_rows(
131    rows: &[SessionPersistenceRow],
132) -> Vec<SessionPersistenceRow> {
133    let mut latest_by_session: BTreeMap<String, SessionPersistenceRow> = BTreeMap::new();
134    for row in rows {
135        let should_replace = match latest_by_session.get(&row.session_id) {
136            Some(existing) => row.updated_at_ms >= existing.updated_at_ms,
137            None => true,
138        };
139        if should_replace {
140            latest_by_session.insert(row.session_id.clone(), row.clone());
141        }
142    }
143    latest_by_session.into_values().collect()
144}
145
146pub fn materialize_live_session_rows(rows: &[SessionPersistenceRow]) -> Vec<SessionPersistenceRow> {
147    materialize_latest_session_rows(rows)
148        .into_iter()
149        .filter(|row| !row.deleted)
150        .collect()
151}
152
153impl JsonFileSessionStore {
154    pub fn new(data_path: impl AsRef<Path>) -> Self {
155        let data_path = data_path.as_ref().to_path_buf();
156        let lock_path = data_path.with_extension("lock");
157        Self {
158            data_path,
159            lock_path,
160            stale_lock_threshold: Duration::from_secs(30),
161        }
162    }
163
164    pub fn with_lock_path(mut self, lock_path: impl AsRef<Path>) -> Self {
165        self.lock_path = lock_path.as_ref().to_path_buf();
166        self
167    }
168
169    pub fn with_stale_lock_threshold(mut self, threshold: Duration) -> Self {
170        self.stale_lock_threshold = threshold;
171        self
172    }
173
174    pub fn data_path(&self) -> &Path {
175        &self.data_path
176    }
177
178    pub fn lock_path(&self) -> &Path {
179        &self.lock_path
180    }
181
182    pub fn append_rows(
183        &self,
184        rows: &[SessionPersistenceRow],
185    ) -> Result<(), JsonFileSessionStoreError> {
186        let _guard = self.acquire_lock()?;
187        if let Some(parent) = self.data_path.parent() {
188            fs::create_dir_all(parent)
189                .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
190        }
191
192        let mut persisted = self.read_rows()?;
193        persisted.extend(rows.iter().cloned());
194
195        let tmp_path = self.data_path.with_extension("tmp");
196        let json = serde_json::to_vec_pretty(&persisted)
197            .map_err(|err| JsonFileSessionStoreError::Serialize(err.to_string()))?;
198        fs::write(&tmp_path, json).map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
199        fs::rename(&tmp_path, &self.data_path)
200            .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
201        Ok(())
202    }
203
204    pub fn read_rows(&self) -> Result<Vec<SessionPersistenceRow>, JsonFileSessionStoreError> {
205        if !self.data_path.exists() {
206            return Ok(vec![]);
207        }
208        let bytes = fs::read(&self.data_path)
209            .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
210        serde_json::from_slice::<Vec<SessionPersistenceRow>>(&bytes)
211            .map_err(|err| JsonFileSessionStoreError::InvalidStoreData(err.to_string()))
212    }
213
214    pub fn read_latest_rows(
215        &self,
216    ) -> Result<Vec<SessionPersistenceRow>, JsonFileSessionStoreError> {
217        let rows = self.read_rows()?;
218        Ok(materialize_latest_session_rows(&rows))
219    }
220
221    pub fn read_live_rows(&self) -> Result<Vec<SessionPersistenceRow>, JsonFileSessionStoreError> {
222        let rows = self.read_rows()?;
223        Ok(materialize_live_session_rows(&rows))
224    }
225
226    fn acquire_lock(&self) -> Result<JsonFileLockGuard, JsonFileSessionStoreError> {
227        if let Some(parent) = self.lock_path.parent() {
228            fs::create_dir_all(parent)
229                .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
230        }
231
232        let mut attempts = 0_u8;
233        loop {
234            attempts += 1;
235            match OpenOptions::new()
236                .create_new(true)
237                .write(true)
238                .open(&self.lock_path)
239            {
240                Ok(mut file) => {
241                    let lock_record = JsonStoreLockRecord {
242                        owner_pid: std::process::id(),
243                        created_at_ms: current_time_ms(),
244                    };
245                    let lock_bytes = serde_json::to_vec(&lock_record)
246                        .map_err(|err| JsonFileSessionStoreError::Serialize(err.to_string()))?;
247                    file.write_all(&lock_bytes)
248                        .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
249                    return Ok(JsonFileLockGuard {
250                        lock_path: self.lock_path.clone(),
251                    });
252                }
253                Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
254                    if attempts >= 2 {
255                        return Err(JsonFileSessionStoreError::LockHeld {
256                            lock_path: self.lock_path.display().to_string(),
257                        });
258                    }
259                    if self.is_lock_stale()? {
260                        fs::remove_file(&self.lock_path).map_err(|remove_err| {
261                            JsonFileSessionStoreError::StaleLockRecoveryFailed(
262                                remove_err.to_string(),
263                            )
264                        })?;
265                        continue;
266                    }
267                    return Err(JsonFileSessionStoreError::LockHeld {
268                        lock_path: self.lock_path.display().to_string(),
269                    });
270                }
271                Err(err) => return Err(JsonFileSessionStoreError::Io(err.to_string())),
272            }
273        }
274    }
275
276    fn is_lock_stale(&self) -> Result<bool, JsonFileSessionStoreError> {
277        let bytes = fs::read(&self.lock_path)
278            .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
279        let stale_threshold_ms = self.stale_lock_threshold.as_millis() as u64;
280        if let Ok(record) = serde_json::from_slice::<JsonStoreLockRecord>(&bytes) {
281            let age_ms = current_time_ms().saturating_sub(record.created_at_ms);
282            if age_ms < stale_threshold_ms {
283                return Ok(false);
284            }
285            return Ok(!is_process_alive(record.owner_pid));
286        }
287
288        let modified = fs::metadata(&self.lock_path)
289            .and_then(|meta| meta.modified())
290            .map_err(|err| JsonFileSessionStoreError::Io(err.to_string()))?;
291        let age = SystemTime::now()
292            .duration_since(modified)
293            .unwrap_or_default();
294        Ok(age >= self.stale_lock_threshold)
295    }
296}
297
298fn is_process_alive(pid: u32) -> bool {
299    if pid == 0 {
300        return false;
301    }
302    let status = Command::new("kill")
303        .arg("-0")
304        .arg(pid.to_string())
305        .stdout(Stdio::null())
306        .stderr(Stdio::null())
307        .status();
308    match status {
309        Ok(exit_status) => exit_status.success(),
310        // If liveness probing is unavailable, avoid evicting potentially active locks.
311        Err(_) => true,
312    }
313}
314
315impl BigQuerySessionStoreAdapter {
316    pub const DEFAULT_API_BASE_URL: &'static str = "https://bigquery.googleapis.com/bigquery/v2";
317
318    pub fn new(
319        _legacy_bq_binary: impl AsRef<Path>,
320        dataset: impl Into<String>,
321        table: impl Into<String>,
322    ) -> Self {
323        Self::new_native(dataset, table)
324    }
325
326    pub fn new_native(dataset: impl Into<String>, table: impl Into<String>) -> Self {
327        Self {
328            dataset: dataset.into(),
329            table: table.into(),
330            project_id: None,
331            api_base_url: Self::DEFAULT_API_BASE_URL.to_string(),
332            access_token: None,
333            client: shared_bigquery_http_client(),
334            http_timeout: Duration::from_secs(30),
335        }
336    }
337
338    pub fn with_project_id(mut self, project_id: impl Into<String>) -> Self {
339        self.project_id = Some(project_id.into());
340        self
341    }
342
343    pub fn with_api_base_url(mut self, api_base_url: impl Into<String>) -> Self {
344        self.api_base_url = api_base_url.into();
345        self
346    }
347
348    pub fn with_access_token(mut self, access_token: impl Into<String>) -> Self {
349        self.access_token = Some(access_token.into());
350        self
351    }
352
353    pub fn with_http_timeout(mut self, timeout: Duration) -> Self {
354        // The timeout applies per REQUEST (see `authorized_request`) over one
355        // shared client: reqwest's first `Client::build()` in a process pays
356        // the macOS system-proxy scan (~700ms observed) since the realtime
357        // feature unification enabled reqwest's `system-proxy`; rebuilding a
358        // client per adapter put that scan — and a connector build — inside
359        // every RPC's latency window.
360        self.http_timeout = timeout;
361        self
362    }
363
364    pub fn with_bearer_token(self, access_token: impl Into<String>) -> Self {
365        self.with_access_token(access_token)
366    }
367
368    pub fn table_ref(&self) -> String {
369        format!("{}.{}", self.dataset, self.table)
370    }
371
372    pub async fn stream_insert_rows(
373        &self,
374        rows: &[SessionPersistenceRow],
375    ) -> Result<(), BigQuerySessionStoreError> {
376        if rows.is_empty() {
377            return Ok(());
378        }
379
380        let project_id = self.resolve_project_id()?;
381        let access_token = self.resolve_access_token()?;
382        let endpoint = format!(
383            "{}/projects/{project_id}/datasets/{}/tables/{}/insertAll",
384            self.api_base_url(),
385            self.dataset,
386            self.table
387        );
388
389        let mut request_rows = Vec::with_capacity(rows.len());
390        for row in rows {
391            let payload_json = serde_json::to_string(&row.payload)
392                .map_err(|err| BigQuerySessionStoreError::Serialize(err.to_string()))?;
393            let mut row_json = serde_json::json!({
394                "session_id": row.session_id,
395                "updated_at_ms": row.updated_at_ms.to_string(),
396                "deleted": row.deleted,
397                "payload": payload_json,
398            });
399            if !row.labels.is_empty() {
400                let labels_json = serde_json::to_string(&row.labels)
401                    .map_err(|err| BigQuerySessionStoreError::Serialize(err.to_string()))?;
402                row_json["labels_json"] = serde_json::Value::String(labels_json);
403            }
404            request_rows.push(serde_json::json!({ "json": row_json }));
405        }
406        let request = serde_json::json!({
407            "ignoreUnknownValues": false,
408            "skipInvalidRows": false,
409            "rows": request_rows,
410        });
411
412        let response = self
413            .send_json_request(
414                reqwest::Method::POST,
415                &endpoint,
416                &access_token,
417                Some(&request),
418            )
419            .await?;
420        if let Some(errors) = response.get("insertErrors").and_then(Value::as_array)
421            && !errors.is_empty()
422        {
423            let detail =
424                serde_json::to_string(errors).unwrap_or_else(|_| "<serialize_error>".to_string());
425            return Err(BigQuerySessionStoreError::Api(format!(
426                "BigQuery insertAll returned row errors: {detail}"
427            )));
428        }
429
430        Ok(())
431    }
432
433    pub async fn read_rows(&self) -> Result<Vec<SessionPersistenceRow>, BigQuerySessionStoreError> {
434        let fq_table = self.fully_qualified_table()?;
435        let query = format!(
436            "SELECT session_id, updated_at_ms, deleted, payload, labels_json FROM `{fq_table}` ORDER BY updated_at_ms ASC"
437        );
438        self.execute_query(&query).await
439    }
440
441    pub async fn read_latest_rows(
442        &self,
443    ) -> Result<Vec<SessionPersistenceRow>, BigQuerySessionStoreError> {
444        let fq_table = self.fully_qualified_table()?;
445        let query = format!(
446            "SELECT session_id, updated_at_ms, deleted, payload, labels_json \
447             FROM `{fq_table}` \
448             QUALIFY ROW_NUMBER() OVER (PARTITION BY session_id ORDER BY updated_at_ms DESC) = 1"
449        );
450        self.execute_query(&query).await
451    }
452
453    pub async fn read_live_rows(
454        &self,
455    ) -> Result<Vec<SessionPersistenceRow>, BigQuerySessionStoreError> {
456        let fq_table = self.fully_qualified_table()?;
457        let query = format!(
458            "SELECT session_id, updated_at_ms, deleted, payload, labels_json FROM (\
459               SELECT session_id, updated_at_ms, deleted, payload, labels_json \
460               FROM `{fq_table}` \
461               QUALIFY ROW_NUMBER() OVER (PARTITION BY session_id ORDER BY updated_at_ms DESC) = 1\
462             ) WHERE deleted = false"
463        );
464        self.execute_query(&query).await
465    }
466
467    fn fully_qualified_table(&self) -> Result<String, BigQuerySessionStoreError> {
468        let project_id = self.resolve_project_id()?;
469        Ok(format!("{}.{}", project_id, self.table_ref()))
470    }
471
472    async fn execute_query(
473        &self,
474        query: &str,
475    ) -> Result<Vec<SessionPersistenceRow>, BigQuerySessionStoreError> {
476        let project_id = self.resolve_project_id()?;
477        let access_token = self.resolve_access_token()?;
478        let endpoint = format!("{}/projects/{project_id}/queries", self.api_base_url());
479        let request = serde_json::json!({
480            "query": query,
481            "useLegacySql": false,
482            "maxResults": 10000,
483        });
484        let response = self
485            .send_json_request(
486                reqwest::Method::POST,
487                &endpoint,
488                &access_token,
489                Some(&request),
490            )
491            .await?;
492        parse_bigquery_query_rows(&response)
493    }
494
495    pub async fn gc_superseded_rows(&self) -> Result<u64, BigQuerySessionStoreError> {
496        let project_id = self.resolve_project_id()?;
497        let access_token = self.resolve_access_token()?;
498        let table_ref = self.table_ref();
499        let endpoint = format!("{}/projects/{project_id}/queries", self.api_base_url());
500        let query = format!(
501            "DELETE FROM `{project_id}.{table_ref}` AS t \
502             WHERE STRUCT(t.session_id, t.updated_at_ms) NOT IN ( \
503               SELECT AS STRUCT session_id, MAX(updated_at_ms) \
504               FROM `{project_id}.{table_ref}` \
505               GROUP BY session_id \
506             )"
507        );
508        let request = serde_json::json!({ "query": query, "useLegacySql": false });
509        let response = self
510            .send_json_request(
511                reqwest::Method::POST,
512                &endpoint,
513                &access_token,
514                Some(&request),
515            )
516            .await?;
517        let affected = response
518            .get("numDmlAffectedRows")
519            .and_then(|v| v.as_str())
520            .and_then(|s| s.parse::<u64>().ok())
521            .unwrap_or(0);
522        Ok(affected)
523    }
524
525    pub async fn truncate_sessions(&self) -> Result<(), BigQuerySessionStoreError> {
526        let project_id = self.resolve_project_id()?;
527        let access_token = self.resolve_access_token()?;
528        let table_ref = self.table_ref();
529        let endpoint = format!("{}/projects/{project_id}/queries", self.api_base_url());
530        let query = format!("TRUNCATE TABLE `{project_id}.{table_ref}`");
531        let request = serde_json::json!({ "query": query, "useLegacySql": false });
532        self.send_json_request(
533            reqwest::Method::POST,
534            &endpoint,
535            &access_token,
536            Some(&request),
537        )
538        .await?;
539        Ok(())
540    }
541
542    fn api_base_url(&self) -> &str {
543        self.api_base_url.trim_end_matches('/')
544    }
545
546    fn resolve_project_id(&self) -> Result<String, BigQuerySessionStoreError> {
547        if let Some(project_id) = self.project_id.as_deref() {
548            let project = project_id.trim();
549            if !project.is_empty() {
550                return Ok(project.to_string());
551            }
552        }
553
554        if let Ok(project_id) = std::env::var("BIGQUERY_PROJECT_ID") {
555            let project = project_id.trim();
556            if !project.is_empty() {
557                return Ok(project.to_string());
558            }
559        }
560
561        Err(BigQuerySessionStoreError::Configuration(
562            "missing BigQuery project_id: call with_project_id(...) or set BIGQUERY_PROJECT_ID"
563                .to_string(),
564        ))
565    }
566
567    fn resolve_access_token(&self) -> Result<String, BigQuerySessionStoreError> {
568        if let Some(token) = self.access_token.as_deref() {
569            let token = token.trim();
570            if !token.is_empty() {
571                return Ok(token.to_string());
572            }
573        }
574
575        for key in [
576            "BIGQUERY_ACCESS_TOKEN",
577            "GOOGLE_OAUTH_ACCESS_TOKEN",
578            "GOOGLE_ACCESS_TOKEN",
579        ] {
580            if let Ok(token) = std::env::var(key) {
581                let token = token.trim();
582                if !token.is_empty() {
583                    return Ok(token.to_string());
584                }
585            }
586        }
587
588        Err(BigQuerySessionStoreError::Configuration(
589            "missing BigQuery access token: call with_access_token(...) or set BIGQUERY_ACCESS_TOKEN"
590                .to_string(),
591        ))
592    }
593
594    async fn send_json_request(
595        &self,
596        method: reqwest::Method,
597        endpoint: &str,
598        access_token: &str,
599        body: Option<&Value>,
600    ) -> Result<Value, BigQuerySessionStoreError> {
601        let mut request = self
602            .client
603            .request(method, endpoint)
604            .timeout(self.http_timeout)
605            .bearer_auth(access_token)
606            .header("accept", "application/json");
607        if let Some(body) = body {
608            request = request
609                .header("content-type", "application/json")
610                .json(body);
611        }
612
613        let response = request
614            .send()
615            .await
616            .map_err(|err| BigQuerySessionStoreError::Http(format!("{err:?}")))?;
617        let status = response.status();
618        let text = response
619            .text()
620            .await
621            .map_err(|err| BigQuerySessionStoreError::Http(format!("{err:?}")))?;
622
623        if !status.is_success() {
624            return Err(BigQuerySessionStoreError::Api(format!(
625                "BigQuery API request failed (status {}): {}",
626                status.as_u16(),
627                text
628            )));
629        }
630
631        if text.trim().is_empty() {
632            return Ok(serde_json::json!({}));
633        }
634
635        serde_json::from_str::<Value>(&text)
636            .map_err(|err| BigQuerySessionStoreError::InvalidQueryResponse(err.to_string()))
637    }
638}
639
640#[derive(Debug, Clone)]
641pub struct BigQueryGcConfig {
642    pub interval: Duration,
643}
644
645impl Default for BigQueryGcConfig {
646    fn default() -> Self {
647        Self {
648            interval: Duration::from_hours(6),
649        }
650    }
651}
652
653/// Callback for GC error reporting. Receives the error message.
654pub type GcErrorCallback = std::sync::Arc<dyn Fn(String) + Send + Sync>;
655
656pub fn run_periodic_gc(
657    adapter: BigQuerySessionStoreAdapter,
658    config: BigQueryGcConfig,
659) -> impl FnOnce() {
660    run_periodic_gc_with_error_callback(adapter, config, None)
661}
662
663pub fn run_periodic_gc_with_error_callback(
664    adapter: BigQuerySessionStoreAdapter,
665    config: BigQueryGcConfig,
666    on_error: Option<GcErrorCallback>,
667) -> impl FnOnce() {
668    move || {
669        let rt = match tokio::runtime::Builder::new_current_thread()
670            .enable_all()
671            .build()
672        {
673            Ok(rt) => rt,
674            Err(err) => {
675                let msg = format!("failed to create async runtime for BQ GC: {err}");
676                eprintln!("[mobkit-gc] {msg}");
677                if let Some(ref cb) = on_error {
678                    cb(msg);
679                }
680                return;
681            }
682        };
683        loop {
684            std::thread::sleep(config.interval);
685            match rt.block_on(adapter.gc_superseded_rows()) {
686                Ok(deleted) => {
687                    if deleted > 0 {
688                        eprintln!("[mobkit-gc] deleted {deleted} superseded session rows");
689                    }
690                }
691                Err(err) => {
692                    let msg = format!("gc_superseded_rows failed: {err:?}");
693                    eprintln!("[mobkit-gc] {msg}");
694                    if let Some(ref cb) = on_error {
695                        cb(msg);
696                    }
697                }
698            }
699        }
700    }
701}
702
703fn parse_bigquery_query_rows(
704    response: &Value,
705) -> Result<Vec<SessionPersistenceRow>, BigQuerySessionStoreError> {
706    if response.is_array() {
707        return serde_json::from_value::<Vec<SessionPersistenceRow>>(response.clone())
708            .map_err(|err| BigQuerySessionStoreError::InvalidQueryResponse(err.to_string()));
709    }
710
711    let rows = response
712        .get("rows")
713        .and_then(Value::as_array)
714        .cloned()
715        .unwrap_or_default();
716    let mut parsed = Vec::with_capacity(rows.len());
717    for row in rows {
718        parsed.push(parse_bigquery_query_row(&row)?);
719    }
720
721    Ok(parsed)
722}
723
724fn parse_bigquery_query_row(
725    row: &Value,
726) -> Result<SessionPersistenceRow, BigQuerySessionStoreError> {
727    let fields = row.get("f").and_then(Value::as_array).ok_or_else(|| {
728        BigQuerySessionStoreError::InvalidQueryResponse(
729            "missing row.f cell array in query response".to_string(),
730        )
731    })?;
732    if fields.len() < 4 {
733        return Err(BigQuerySessionStoreError::InvalidQueryResponse(
734            "query response row has fewer than 4 columns".to_string(),
735        ));
736    }
737
738    let session_id = parse_bigquery_string_cell(&fields[0], "session_id")?;
739    let updated_at_ms = parse_bigquery_u64_cell(&fields[1], "updated_at_ms")?;
740    let deleted = parse_bigquery_bool_cell(&fields[2], "deleted")?;
741    let payload = parse_bigquery_payload_cell(&fields[3], "payload")?;
742    let labels = if fields.len() > 4 {
743        parse_bigquery_labels_cell(&fields[4])?
744    } else {
745        BTreeMap::new()
746    };
747
748    Ok(SessionPersistenceRow {
749        session_id,
750        updated_at_ms,
751        deleted,
752        payload,
753        labels,
754    })
755}
756
757fn parse_bigquery_string_cell(
758    cell: &Value,
759    column: &str,
760) -> Result<String, BigQuerySessionStoreError> {
761    let value = bigquery_cell_value(cell);
762    match value {
763        Value::String(s) => Ok(s.clone()),
764        _ => Err(BigQuerySessionStoreError::InvalidQueryResponse(format!(
765            "query column {column} is not a string"
766        ))),
767    }
768}
769
770fn parse_bigquery_u64_cell(cell: &Value, column: &str) -> Result<u64, BigQuerySessionStoreError> {
771    let value = bigquery_cell_value(cell);
772    match value {
773        Value::Number(num) => num.as_u64().ok_or_else(|| {
774            BigQuerySessionStoreError::InvalidQueryResponse(format!(
775                "query column {column} is not a u64 number"
776            ))
777        }),
778        Value::String(s) => s.parse::<u64>().map_err(|_| {
779            BigQuerySessionStoreError::InvalidQueryResponse(format!(
780                "query column {column} is not a u64 string"
781            ))
782        }),
783        _ => Err(BigQuerySessionStoreError::InvalidQueryResponse(format!(
784            "query column {column} is not a u64 value"
785        ))),
786    }
787}
788
789fn parse_bigquery_bool_cell(cell: &Value, column: &str) -> Result<bool, BigQuerySessionStoreError> {
790    let value = bigquery_cell_value(cell);
791    match value {
792        Value::Bool(flag) => Ok(*flag),
793        Value::String(s) => match s.as_str() {
794            "true" | "TRUE" | "1" => Ok(true),
795            "false" | "FALSE" | "0" => Ok(false),
796            _ => Err(BigQuerySessionStoreError::InvalidQueryResponse(format!(
797                "query column {column} is not a bool string"
798            ))),
799        },
800        _ => Err(BigQuerySessionStoreError::InvalidQueryResponse(format!(
801            "query column {column} is not a bool value"
802        ))),
803    }
804}
805
806fn parse_bigquery_payload_cell(
807    cell: &Value,
808    column: &str,
809) -> Result<Value, BigQuerySessionStoreError> {
810    let value = bigquery_cell_value(cell);
811    match value {
812        Value::Null => Ok(serde_json::json!({})),
813        Value::String(s) => {
814            if s.trim().is_empty() {
815                return Ok(serde_json::json!({}));
816            }
817            serde_json::from_str::<Value>(s).map_err(|_| {
818                BigQuerySessionStoreError::InvalidQueryResponse(format!(
819                    "query column {column} payload JSON parse failed"
820                ))
821            })
822        }
823        _ => Ok(value.clone()),
824    }
825}
826
827fn parse_bigquery_labels_cell(
828    cell: &Value,
829) -> Result<BTreeMap<String, String>, BigQuerySessionStoreError> {
830    let value = bigquery_cell_value(cell);
831    match value {
832        Value::Null => Ok(BTreeMap::new()),
833        Value::String(s) if s.trim().is_empty() => Ok(BTreeMap::new()),
834        Value::String(s) => serde_json::from_str::<BTreeMap<String, String>>(s).map_err(|_| {
835            BigQuerySessionStoreError::InvalidQueryResponse(
836                "query column labels_json parse failed".to_string(),
837            )
838        }),
839        _ => Ok(BTreeMap::new()),
840    }
841}
842
843fn bigquery_cell_value(cell: &Value) -> &Value {
844    cell.get("v").unwrap_or(cell)
845}
meerkat_mobkit/runtime/session_store.rs

meerkat_mobkit/runtime/
session_store.rs
