Add SQL mode and make it use the query! macro for type checking

.github/workflows/build.yml

@@ -45,6 +45,8 @@ jobs:
 
       - name: Build release binary
         run: cargo build --release
+        env:
+          SQLX_OFFLINE: true
 
       - name: Upload binary
         uses: actions/upload-artifact@v3

.sqlx/query-ac15bfcd1737751e27388ffddfe8ec47fd7277be74c47a59484e4fc993671d43.json

@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "with timelineItems as (select jsonb_array_elements(raw_files.json -> 'timelineItems') as timelineItem\n                       from raw_files\n                       where date = ANY ($1)),\n     max_last_saved as (select timelineItem ->> 'itemId'                          as itemId,\n                               max((timelineItem ->> 'lastSaved') :: timestamptz) as latest_last_saved\n                        from timelineItems\n                        group by timelineItem ->> 'itemId'),\n     unique_timline_items as (select distinct on (max_last_saved.itemId) *\n                              from max_last_saved\n                                       inner join timelineItems\n                                                  on timelineItems.timelineItem ->> 'itemId' = max_last_saved.itemId\n                                                      and (timelineItems.timelineItem ->> 'lastSaved') :: timestamptz =\n                                                          max_last_saved.latest_last_saved)\ninsert\ninto public.timeline_item (item_id, json, place_id, end_date, last_saved, server_last_updated)\nselect unique_timline_items.itemId :: uuid                                  as item_id,\n       unique_timline_items.timelineItem                                    as json,\n       (unique_timline_items.timelineItem -> 'place' ->> 'placeId') :: uuid as place_id,\n       (unique_timline_items.timelineItem ->> 'endDate') :: timestamptz     as end_date,\n       unique_timline_items.latest_last_saved :: timestamptz                as last_saved,\n       now()                                                                as server_last_updated\nfrom unique_timline_items\non conflict (item_id) do update set json                = excluded.json,\n                                    place_id            = excluded.place_id,\n                                    end_date            = excluded.end_date,\n                                    last_saved          = excluded.last_saved,\n                                    server_last_updated = excluded.server_last_updated\nwhere excluded.last_saved > public.timeline_item.last_saved\nreturning item_id;\n",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "item_id",
+        "type_info": "Uuid"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "TextArray"
+      ]
+    },
+    "nullable": [
+      false
+    ]
+  },
+  "hash": "ac15bfcd1737751e27388ffddfe8ec47fd7277be74c47a59484e4fc993671d43"
+}

.sqlx/query-f0402ce4a5c93837f39559bfb3358a7655f178080465f7abd7f7702d42474157.json

@@ -0,0 +1,14 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "with timelineItems as (select jsonb_array_elements(raw_files.json -> 'timelineItems') as timelineItem\n                       from raw_files\n                       where date = ANY ($1)),\n     places as (select distinct on (md5(timelineItem ->> 'place' :: text)) timelineItem -> 'place'                                  as place,\n                                                                           timelineItem -> 'place' ->> 'placeId'                    as placeId,\n                                                                           (timelineItem -> 'place' ->> 'lastSaved') :: timestamptz as lastSaved\n                from timelineItems\n                where timelineItem ->> 'place' is not null),\n     places_with_max_last_saved as (select place -> 'placeId'                          as placeId,\n                                           max((place ->> 'lastSaved') :: timestamptz) as latest_last_saved\n                                    from places\n                                    group by place -> 'placeId'),\n     latest_places as (select places.*\n                       from places_with_max_last_saved\n                                inner join places on places.place -> 'placeId' = places_with_max_last_saved.placeId and\n                                                     places.lastSaved =\n                                                     places_with_max_last_saved.latest_last_saved)\ninsert\ninto public.place (place_id, json, last_saved, server_last_updated)\nselect (placeId :: uuid) as place_id, place as json, lastSaved as last_saved, now() as server_last_updated\nfrom latest_places\non conflict (place_id) do update set json                = excluded.json,\n                                     last_saved          = excluded.last_saved,\n                                     server_last_updated = excluded.server_last_updated\nwhere excluded.last_saved > public.place.last_saved;\n",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "TextArray"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "f0402ce4a5c93837f39559bfb3358a7655f178080465f7abd7f7702d42474157"
+}

Cargo.lock

@@ -110,7 +110,7 @@ dependencies = [
 
 [[package]]
 name = "arc-ingester"
-version = "0.1.1"
+version = "0.1.2"
 dependencies = [
  "anyhow",
  "chrono",
@@ -118,7 +118,6 @@ dependencies = [
  "flate2",
  "futures",
  "itertools",
- "rayon",
  "serde",
  "serde_json",
  "sha256",
@@ -127,6 +126,8 @@ dependencies = [
  "testcontainers-modules",
  "thiserror",
  "tokio",
+ "tracing",
+ "tracing-subscriber",
  "uuid",
 ]
 
@@ -366,25 +367,6 @@ dependencies = [
  "cfg-if",
 ]
 
-[[package]]
-name = "crossbeam-deque"
-version = "0.8.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d"
-dependencies = [
- "crossbeam-epoch",
- "crossbeam-utils",
-]
-
-[[package]]
-name = "crossbeam-epoch"
-version = "0.9.18"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
-dependencies = [
- "crossbeam-utils",
-]
-
 [[package]]
 name = "crossbeam-queue"
 version = "0.3.11"
@@ -948,6 +930,16 @@ dependencies = [
  "minimal-lexical",
 ]
 
+[[package]]
+name = "nu-ansi-term"
+version = "0.46.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84"
+dependencies = [
+ "overload",
+ "winapi",
+]
+
 [[package]]
 name = "num-bigint-dig"
 version = "0.8.4"
@@ -1021,6 +1013,12 @@ version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
 
+[[package]]
+name = "overload"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.1"
@@ -1158,26 +1156,6 @@ dependencies = [
  "getrandom",
 ]
 
-[[package]]
-name = "rayon"
-version = "1.8.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fa7237101a77a10773db45d62004a272517633fbcc3df19d96455ede1122e051"
-dependencies = [
- "either",
- "rayon-core",
-]
-
-[[package]]
-name = "rayon-core"
-version = "1.12.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2"
-dependencies = [
- "crossbeam-deque",
- "crossbeam-utils",
-]
-
 [[package]]
 name = "redox_syscall"
 version = "0.4.1"
@@ -1380,6 +1358,15 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "sharded-slab"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
+dependencies = [
+ "lazy_static",
+]
+
 [[package]]
 name = "signal-hook-registry"
 version = "1.4.1"
@@ -1771,6 +1758,16 @@ dependencies = [
  "syn 2.0.48",
 ]
 
+[[package]]
+name = "thread_local"
+version = "1.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c"
+dependencies = [
+ "cfg-if",
+ "once_cell",
+]
+
 [[package]]
 name = "tinyvec"
 version = "1.6.0"
@@ -1857,6 +1854,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54"
 dependencies = [
  "once_cell",
+ "valuable",
+]
+
+[[package]]
+name = "tracing-log"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3"
+dependencies = [
+ "log",
+ "once_cell",
+ "tracing-core",
+]
+
+[[package]]
+name = "tracing-subscriber"
+version = "0.3.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b"
+dependencies = [
+ "nu-ansi-term",
+ "sharded-slab",
+ "smallvec",
+ "thread_local",
+ "tracing-core",
+ "tracing-log",
 ]
 
 [[package]]
@@ -1936,6 +1959,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "valuable"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d"
+
 [[package]]
 name = "vcpkg"
 version = "0.2.15"
@@ -2020,6 +2049,28 @@ version = "1.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "22fc3756b8a9133049b26c7f61ab35416c130e8c09b660f5b3958b446f52cc50"
 
+[[package]]
+name = "winapi"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
+
 [[package]]
 name = "windows-core"
 version = "0.52.0"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "arc-ingester"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 authors = ["Joshua Coles"]
 license = "MIT OR Apache-2.0"
@@ -23,7 +23,6 @@ thiserror = "1"
 sha256 = "1"
 serde = { version = "1.0.196", features = ["derive"] }
 serde_json = "1.0.113"
-rayon = "1.8.1"
 futures = "0.3.30"
 itertools = "0.12.1"
 chrono = { version = "0.4.33", features = ["serde"] }
@@ -31,3 +30,5 @@ uuid = { version = "1.7.0", features = ["serde"] }
 testcontainers = "0.15.0"
 testcontainers-modules = { version = "0.3.6", features = ["postgres"] }
 anyhow = { version = "1.0.82", features = ["backtrace"] }
+tracing = "0.1.40"
+tracing-subscriber = "0.3.18"

build.rs

@@ -2,4 +2,5 @@
 fn main() {
     // trigger recompilation when a new migration is added
     println!("cargo:rerun-if-changed=migrations");
+    println!("cargo:rustc-env=DATABASE_URL=postgres://joshuacoles@localhost/arc_test");
 }

functions/update_places.sql

@@ -1,5 +1,6 @@
 with timelineItems as (select jsonb_array_elements(raw_files.json -> 'timelineItems') as timelineItem
-                       from raw_files),
+                       from raw_files
+                       where date = ANY ($1)),
      places as (select distinct on (md5(timelineItem ->> 'place' :: text)) timelineItem -> 'place'                                  as place,
                                                                            timelineItem -> 'place' ->> 'placeId'                    as placeId,
                                                                            (timelineItem -> 'place' ->> 'lastSaved') :: timestamptz as lastSaved

functions/update_timeline_items.sql

@@ -1,5 +1,6 @@
 with timelineItems as (select jsonb_array_elements(raw_files.json -> 'timelineItems') as timelineItem
-                       from raw_files),
+                       from raw_files
+                       where date = ANY ($1)),
      max_last_saved as (select timelineItem ->> 'itemId'                          as itemId,
                                max((timelineItem ->> 'lastSaved') :: timestamptz) as latest_last_saved
                         from timelineItems

src/main.rs

@@ -5,8 +5,9 @@ use std::path::PathBuf;
 use chrono::{DateTime, Utc};
 use clap::Parser;
 use itertools::Itertools;
-use sqlx::FromRow;
+use sqlx::{Executor, FromRow, PgPool, Pool, Postgres};
 use serde_json::Value;
+use tracing::instrument;
 use uuid::Uuid;
 
 #[derive(Parser, Debug)]
@@ -18,6 +19,9 @@ struct Cli {
     /// psql connection string
     #[arg(long, short, env = "ARC_DB")]
     conn: String,
+
+    #[arg(long)]
+    sql_only: bool,
 }
 
 #[derive(Debug, FromRow)]
@@ -70,6 +74,7 @@ struct Place {
 fn hash_files(files: ReadDir) -> impl Iterator<Item=(String, PathBuf, String)> {
     files.into_iter()
         .map(|f| f.unwrap().path())
+        .filter(|p| p.extension().and_then(|ext| ext.to_str()).filter(|p| *p == "gz").is_some())
         .map(|path| {
             let date = {
                 let file_name = path.file_name().unwrap().to_str().unwrap();
@@ -82,7 +87,8 @@ fn hash_files(files: ReadDir) -> impl Iterator<Item=(String, PathBuf, String)> {
         })
 }
 
-async fn find_updated(db: &sqlx::PgPool, files: ReadDir) -> Vec<UpdatedFile> {
+#[instrument(skip(db, files))]
+async fn find_updated(db: &PgPool, files: ReadDir) -> Vec<UpdatedFile> {
     let date_hashes: Vec<DateHash> = sqlx::query_as("SELECT date, sha256 FROM raw_files")
         .fetch_all(db)
         .await
@@ -95,10 +101,19 @@ async fn find_updated(db: &sqlx::PgPool, files: ReadDir) -> Vec<UpdatedFile> {
     let new_hashes = hash_files(files);
 
     new_hashes.filter_map(|(date, path, new_hash)| {
+        let span = tracing::span!(tracing::Level::DEBUG, "considering_file", path = ?path, date = ?date);
+        let _enter = span.enter();
+
+        tracing::debug!("Considering file for updates {path:?} (read as date: {date})");
         if let Some(existing_hash) = date_hashes.get(&date) {
             if new_hash == *existing_hash {
+                tracing::debug!("Found and matched to to existing hash");
                 return None;
+            } else {
+                tracing::debug!("Found existing hash but file does not match");
             }
+        } else {
+            tracing::debug!("No existing hash found");
         }
 
         let bytes = std::fs::read(&path).unwrap();
@@ -122,7 +137,9 @@ async fn find_updated(db: &sqlx::PgPool, files: ReadDir) -> Vec<UpdatedFile> {
 }
 
 #[tokio::main]
-async fn main() {
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt::init();
+
     let cli = Cli::parse();
     let daily_exports = cli.root.join("Export/JSON/Daily");
     let files = std::fs::read_dir(daily_exports)
@@ -140,17 +157,43 @@ async fn main() {
     let need_refresh = find_updated(&db, files)
         .await;
 
-    // Refresh the database with the new files
-    for updated_file in &need_refresh {
-        sqlx::query("INSERT INTO raw_files (date, sha256, json) VALUES ($1, $2, $3 :: jsonb) ON CONFLICT (date) DO UPDATE SET sha256 = excluded.sha256, json = excluded.json")
-            .bind(&updated_file.date)
-            .bind(&updated_file.sha256)
-            .bind(&serde_json::to_value(&updated_file.json).unwrap())
-            .execute(&db)
-            .await
-            .unwrap();
+    upload_files(&db, &need_refresh)
+        .await?;
+
+    if cli.sql_only {
+        update_data_sql(&db, need_refresh)
+            .await?;
+    } else {
+        update_data(&db, need_refresh)
+            .await?;
     }
 
+    Ok(())
+}
+
+#[instrument(skip(db, need_refresh))]
+async fn update_data_sql(db: &PgPool, need_refresh: Vec<UpdatedFile>) -> anyhow::Result<()> {
+    let vec = need_refresh.iter()
+        .map(|d| d.date.clone())
+        .collect_vec();
+
+    let result = sqlx::query_file!("functions/update_places.sql", &vec)
+        .execute(db)
+        .await?;
+
+    tracing::info!("Updated {} places", result.rows_affected());
+
+    let updated = sqlx::query_file!("functions/update_timeline_items.sql", &vec)
+        .fetch_all(db)
+        .await?;
+
+    tracing::info!("Updated {} timeline items", updated.len());
+
+    Ok(())
+}
+
+#[instrument(skip(db, need_refresh))]
+async fn update_data(db: &Pool<Postgres>, need_refresh: Vec<UpdatedFile>) -> anyhow::Result<()> {
     // Take all the changed files' timeline items, and group them by item_id, then take the latest one.
     // If we are needing to update the database it will be with this one.
     let possibly_timeline_items = need_refresh.into_iter()
@@ -167,9 +210,8 @@ async fn main() {
 
     let existing_last_saved_at_map = sqlx::query_as("SELECT item_id, end_date, last_saved FROM timeline_item where item_id = ANY($1)")
         .bind(&possibly_timeline_item_ids)
-        .fetch_all(&db)
-        .await
-        .unwrap()
+        .fetch_all(db)
+        .await?
         .into_iter()
         .map(|row: TimelineItemUpdatedCheckRow| (row.item_id, row.last_saved))
         .collect::<HashMap<_, _>>();
@@ -190,18 +232,37 @@ async fn main() {
         if let Some(place) = &updated_timeline_item.place {
             sqlx::query("INSERT INTO place (place_id, json, last_saved, server_last_updated) VALUES ($1, $2 :: jsonb, $3, now()) ON CONFLICT (place_id) DO UPDATE SET json = $2 :: jsonb, last_saved = $3, server_last_updated = now()")
                 .bind(&place.place_id)
-                .bind(&serde_json::to_value(&place).unwrap())
+                .bind(&serde_json::to_value(&place)?)
                 .bind(&place.last_saved)
-                .execute(&db).await.unwrap();
+                .execute(db)
+                .await?;
         }
 
         // Then we can insert/update the timeline item.
         sqlx::query("INSERT INTO timeline_item (item_id, json, place_id, end_date, last_saved, server_last_updated) VALUES ($1, $2 :: jsonb, $3, $4, $5, now()) ON CONFLICT (item_id) DO UPDATE SET json = $2 :: jsonb, place_id = $3, end_date = $4, last_saved = $5, server_last_updated = now()")
             .bind(&updated_timeline_item.item_id)
-            .bind(&serde_json::to_value(&updated_timeline_item).unwrap())
+            .bind(&serde_json::to_value(&updated_timeline_item)?)
             .bind(&updated_timeline_item.place.map(|place| place.place_id))
             .bind(&updated_timeline_item.end_date)
             .bind(&updated_timeline_item.last_saved)
-            .execute(&db).await.unwrap();
+            .execute(db)
+            .await?;
     }
+
+    Ok(())
+}
+
+#[tracing::instrument(skip(db, need_refresh), err)]
+async fn upload_files(db: &PgPool, need_refresh: &[UpdatedFile]) -> anyhow::Result<()> {
+    // Refresh the database with the new files
+    for updated_file in need_refresh {
+        sqlx::query("INSERT INTO raw_files (date, sha256, json) VALUES ($1, $2, $3 :: jsonb) ON CONFLICT (date) DO UPDATE SET sha256 = excluded.sha256, json = excluded.json")
+            .bind(&updated_file.date)
+            .bind(&updated_file.sha256)
+            .bind(&serde_json::to_value(&updated_file.json)?)
+            .execute(db)
+            .await?;
+    }
+
+    Ok(())
 }