Merge pull request #45 from lmnr-ai/dev

input costs, docker buildx, project api keys

.github/workflows/build-push.yml

@@ -34,25 +34,29 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Log in to the Container registry
-        uses: docker/login-action@65b78e6e13532edd9afa3aa52ac7964289d1a9c1
+        uses: docker/login-action@v3
         with:
           registry: ${{ env.REGISTRY_GH }}
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
 
       - name: Extract metadata (tags, labels) for Docker
         id: meta
-        uses: docker/metadata-action@9ec57ed1fcdbf14dcef7dfbe97b2010124a938b7
+        uses: docker/metadata-action@v5
         with:
           images: ${{ matrix.image }}
 
       - name: Build and push Docker images
         id: push
-        uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4
+        uses: docker/build-push-action@v6
         with:
           context: ${{ matrix.context }}
           file: ${{ matrix.dockerfile }}
           push: true
+          platforms: linux/amd64,linux/arm64 
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
 

README.md

@@ -28,7 +28,7 @@ The easiest way to get started is with a generous free tier on our managed platf
 
 Start local version with docker compose.
 ```sh
-git clone git@github.com:lmnr-ai/lmnr
+git clone https://github.com/lmnr-ai/lmnr
 cd lmnr
 docker compose up
 ```

app-server/Cargo.toml

@@ -64,6 +64,7 @@ time = "0.3.36"
 rustls = { version = "0.23.12", features = ["ring"] }
 serde_repr = "0.1.19"
 num_cpus = "1.16.0"
+sha3 = "0.10.8"
 
 [build-dependencies]
 tonic-build = "0.12.3"

app-server/src/api/utils.rs

@@ -1,9 +1,11 @@
 use std::sync::Arc;
 
+use sqlx::PgPool;
 use uuid::Uuid;
 
 use crate::db::pipelines::PipelineVersion;
 use crate::pipeline::utils::get_target_pipeline_version_cache_key;
+use crate::routes::api_keys::hash_api_key;
 use crate::routes::error;
 use crate::{
     cache::Cache,
@@ -37,3 +39,25 @@ pub async fn query_target_pipeline_version(
         }
     }
 }
+
+pub async fn get_api_key_from_raw_value(
+    pool: &PgPool,
+    cache: Arc<Cache>,
+    raw_api_key: String,
+) -> anyhow::Result<db::api_keys::ProjectApiKey> {
+    let api_key_hash = hash_api_key(&raw_api_key);
+    let cache_res = cache
+        .get::<db::api_keys::ProjectApiKey>(&api_key_hash)
+        .await;
+    match cache_res {
+        Ok(Some(api_key)) => Ok(api_key),
+        Ok(None) | Err(_) => {
+            let api_key = db::api_keys::get_api_key(pool, &api_key_hash).await?;
+            let _ = cache
+                .insert::<db::api_keys::ProjectApiKey>(api_key_hash, &api_key)
+                .await;
+
+            Ok(api_key)
+        }
+    }
+}

app-server/src/auth/mod.rs

@@ -11,8 +11,9 @@ use actix_web::{FromRequest, HttpMessage, HttpRequest};
 use actix_web_httpauth::extractors::bearer::{BearerAuth, Config};
 use actix_web_httpauth::extractors::AuthenticationError;
 
+use crate::api::utils::get_api_key_from_raw_value;
 use crate::cache::Cache;
-use crate::db::api_keys::{get_api_key, ProjectApiKey};
+use crate::db::api_keys::ProjectApiKey;
 use crate::db::user::{get_user_from_api_key, User};
 use crate::db::DB;
 
@@ -96,7 +97,7 @@ pub async fn project_validator(
         .unwrap()
         .into_inner();
 
-    match get_api_key(&db.pool, &credentials.token().to_string(), cache.clone()).await {
+    match get_api_key_from_raw_value(&db.pool, cache, credentials.token().to_string()).await {
         Ok(api_key) => {
             req.extensions_mut().insert(api_key);
             Ok(req)

app-server/src/cache/cache.rs

@@ -4,8 +4,6 @@ use std::result::Result;
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use sqlx::postgres::PgRow;
-use sqlx::FromRow;
 
 #[derive(thiserror::Error, Debug)]
 pub enum CacheError {
@@ -25,7 +23,7 @@ pub trait CacheTrait: Sync + Send {
 #[async_trait]
 impl<T> CacheTrait for moka::future::Cache<String, T>
 where
-    T: for<'a> FromRow<'a, PgRow> + 'static + Send + Sync + Clone,
+    T: 'static + Send + Sync + Clone,
 {
     async fn get(&self, key: &str) -> Option<Box<dyn Any>> {
         self.get(key)

app-server/src/db/api_keys.rs

@@ -1,48 +1,57 @@
-use std::sync::Arc;
-
 use anyhow::Result;
-use serde::{Deserialize, Serialize};
+use serde::Serialize;
 use sqlx::{FromRow, PgPool};
 use uuid::Uuid;
 
-use crate::cache::Cache;
-
-#[derive(Debug, Clone, Deserialize, Serialize, FromRow)]
-#[serde(rename_all = "camelCase")]
+#[derive(Debug, Clone, FromRow)]
 pub struct ProjectApiKey {
-    pub value: String,
     pub project_id: Uuid,
     pub name: Option<String>,
+    pub hash: String,
+    pub shorthand: String,
 }
 
-pub async fn create_project_api_key(
-    db: &PgPool,
-    api_key: &ProjectApiKey,
-    cache: Arc<Cache>,
-) -> Result<()> {
-    sqlx::query("INSERT INTO project_api_keys (value, project_id, name) VALUES ($1, $2, $3);")
-        .bind(&api_key.value)
-        .bind(api_key.project_id)
-        .bind(&api_key.name)
-        .execute(db)
-        .await?;
+#[derive(Serialize, FromRow)]
+pub struct ProjectApiKeyResponse {
+    pub id: Uuid,
+    pub project_id: Uuid,
+    pub name: Option<String>,
+    pub shorthand: String,
+}
 
-    let _ = cache
-        .insert::<ProjectApiKey>(api_key.value.clone(), api_key)
-        .await;
+pub async fn create_project_api_key(
+    pool: &PgPool,
+    project_id: &Uuid,
+    name: &Option<String>,
+    hash: &String,
+    shorthand: &String,
+) -> Result<ProjectApiKey> {
+    let key_info = sqlx::query_as::<_, ProjectApiKey>(
+        "INSERT
+        INTO project_api_keys (shorthand, project_id, name, hash)
+        VALUES ($1, $2, $3, $4)
+        RETURNING id, project_id, name, hash, shorthand",
+    )
+    .bind(&shorthand)
+    .bind(&project_id)
+    .bind(&name)
+    .bind(&hash)
+    .fetch_one(pool)
+    .await?;
 
-    Ok(())
+    Ok(key_info)
 }
 
 pub async fn get_api_keys_for_project(
     db: &PgPool,
     project_id: &Uuid,
-) -> Result<Vec<ProjectApiKey>> {
-    let api_keys = sqlx::query_as::<_, ProjectApiKey>(
+) -> Result<Vec<ProjectApiKeyResponse>> {
+    let api_keys = sqlx::query_as::<_, ProjectApiKeyResponse>(
         "SELECT
-            project_api_keys.value,
             project_api_keys.project_id,
-            project_api_keys.name
+            project_api_keys.name,
+            project_api_keys.id,
+            project_api_keys.shorthand
         FROM
             project_api_keys
         WHERE
@@ -55,50 +64,47 @@ pub async fn get_api_keys_for_project(
     Ok(api_keys)
 }
 
-pub async fn get_api_key(
-    db: &PgPool,
-    api_key: &String,
-    cache: Arc<Cache>,
-) -> Result<ProjectApiKey> {
-    let cache_res = cache.get::<ProjectApiKey>(api_key).await;
-    match cache_res {
-        Ok(Some(api_key)) => return Ok(api_key),
-        Ok(None) => {}
-        Err(e) => log::error!("Error getting project API key from cache: {}", e),
-    }
-
+pub async fn get_api_key(pool: &PgPool, hash: &String) -> Result<ProjectApiKey> {
     let api_key = match sqlx::query_as::<_, ProjectApiKey>(
         "SELECT
-            project_api_keys.value,
+            project_api_keys.hash,
             project_api_keys.project_id,
-            project_api_keys.name
+            project_api_keys.name,
+            project_api_keys.id,
+            project_api_keys.shorthand
         FROM
             project_api_keys
         WHERE
-            project_api_keys.value = $1",
+            project_api_keys.hash = $1",
     )
-    .bind(api_key)
-    .fetch_optional(db)
+    .bind(hash)
+    .fetch_optional(pool)
     .await
     {
         Ok(None) => Err(anyhow::anyhow!("invalid project API key")),
-        Ok(Some(api_key_meta)) => {
-            let _ = cache
-                .insert::<ProjectApiKey>(api_key.clone(), &api_key_meta)
-                .await;
-            Ok(api_key_meta)
-        }
+        Ok(Some(api_key)) => Ok(api_key),
         Err(e) => Err(e.into()),
     }?;
 
     Ok(api_key)
 }
 
-pub async fn delete_api_key(pool: &PgPool, api_key: &String, project_id: &Uuid) -> Result<()> {
-    sqlx::query("DELETE FROM project_api_keys WHERE value = $1 AND project_id = $2")
-        .bind(api_key)
-        .bind(project_id)
-        .execute(pool)
-        .await?;
-    Ok(())
+#[derive(FromRow)]
+struct ProjectApiKeyHash {
+    hash: String,
+}
+
+pub async fn delete_api_key(pool: &PgPool, id: &Uuid, project_id: &Uuid) -> Result<String> {
+    let row = sqlx::query_as::<_, ProjectApiKeyHash>(
+        "DELETE
+        FROM project_api_keys
+        WHERE id = $1 AND project_id = $2
+        RETURNING hash",
+    )
+    .bind(id)
+    .bind(project_id)
+    .fetch_one(pool)
+    .await?;
+
+    Ok(row.hash)
 }