refactor: jobs table to contain jsonb data col

Author: McPatate

generate_signature.py

@@ -3,7 +3,7 @@
 
 
 secret = "tmpsecret"
-payload_body_string = '{"action":"submitted","review":{"body":"test review","id":1234,"url":"https://github.com/huggingface/lor-e/5#comment-123"},"pull_request":{"title":"my great contribution to the world","body":"superb work, isnt it","id":4321,"number":5,"html_url":"https://github.com/huggingface/lor-e/5", "url":"https://github.com/api/huggingface/lor-e/5"}}'
+payload_body_string = '{"action":"created","comment":{"body":"test review","id":1234,"url":"https://github.com/huggingface/lor-e/5#comment-123"}, "issue":{"title":"my great contribution to the world","body":"superb work, isnt it","id":4321,"number":5,"html_url":"https://github.com/huggingface/lor-e/5", "url":"https://github.com/api/huggingface/lor-e/5"}}'
 
 
 def generate_signature(payload_body, secret_token):

init_db.sql

@@ -36,7 +36,7 @@ CREATE INDEX IF NOT EXISTS issues_embedding_hnsw_idx ON issues USING hnsw (embed
 CREATE TABLE IF NOT EXISTS jobs (
   id SERIAL PRIMARY KEY,
   repository_id VARCHAR NOT NULL UNIQUE,
-  page INT NOT NULL,
+  data JSONB NOT NULL,
   created_at timestamp with time zone NOT NULL DEFAULT (current_timestamp AT TIME ZONE 'UTC'),
   updated_at timestamp with time zone NOT NULL DEFAULT (current_timestamp AT TIME ZONE 'UTC')
 );

issue-bot/Cargo.toml

@@ -24,7 +24,7 @@ once_cell = "1.20"
 pgvector = { version = "0.4", features = ["sqlx"] }
 reqwest = { version = "0.12", features = ["json"] }
 serde = { version = "1.0", features = ["derive"] }
-serde_json = "1"
+serde_json = { version = "1", features = ["raw_value"] }
 sha2 = "0.10"
 sqlx = { version = "0.8", features = ["chrono", "postgres", "runtime-tokio"] }
 thiserror = "2"

issue-bot/src/github.rs

@@ -42,8 +42,9 @@ pub enum GithubApiError {
 }
 
 #[derive(Debug, Deserialize)]
+#[allow(unused)]
 struct PullRequest {
-    #[allow(unused)]
+    html_url: String,
     url: String,
 }
 
@@ -180,13 +181,13 @@ impl GithubApi {
 
     pub(crate) fn get_issues(
         &self,
-        from_page: Option<i32>,
+        from_page: i32,
         repository: RepositoryData,
     ) -> impl Stream<Item = Result<(IssueWithComments, Option<i32>), GithubApiError>> + use<'_> {
         try_stream! {
             let url = format!("https://api.github.com/repos/{}/issues", repository.repo_id);
             let client = self.client.clone();
-            let mut page = from_page.unwrap_or(1);
+            let mut page = from_page;
             loop {
                 let res = client
                     .get(&url)
@@ -206,7 +207,6 @@ impl GithubApi {
                 if get_next_page(link_header)?.is_none() {
                     break;
                 }
-                let issues: Vec<Issue> = issues.into_iter().filter(|i| i.pull_request.is_none()).collect();
                 let page_issue_count = issues.len();
                 for (i, issue) in issues.into_iter().enumerate() {
                     let res = client
@@ -226,14 +226,6 @@ impl GithubApi {
             }
         }
     }
-
-    // pub(crate) async fn get_prs(
-    //     &self,
-    //     repository: &RepositoryData,
-    // ) -> Result<Vec<IssueWithComments>, GithubApiError> {
-    //     let issues = Vec::new();
-    //     Ok(issues)
-    // }
 }
 
 async fn handle_ratelimit(remaining: Option<HeaderValue>, reset: Option<HeaderValue>) -> Result<(), GithubApiError> {

issue-bot/src/main.rs

@@ -17,7 +17,7 @@ use axum::{
 };
 use config::{load_config, IssueBotConfig, ServerConfig};
 use embeddings::inference_endpoints::EmbeddingApi;
-use futures::{pin_mut, StreamExt};
+use futures::{pin_mut, TryStreamExt};
 use github::GithubApi;
 use huggingface::HuggingfaceApi;
 use metrics::start_metrics_server;
@@ -30,6 +30,7 @@ use slack::Slack;
 use sqlx::{
     postgres::{PgConnectOptions, PgPoolOptions},
     prelude::FromRow,
+    types::Json,
     Pool, Postgres, QueryBuilder,
 };
 use tokio::{
@@ -244,11 +245,20 @@ struct ClosestIssue {
     title: String,
     number: i32,
     html_url: String,
-    body: String,
     #[allow(unused)]
     cosine_similarity: f64,
 }
 
+#[derive(Debug, Deserialize, FromRow)]
+struct JobData {
+    issues_page: i32,
+}
+
+#[derive(Debug, FromRow)]
+struct Job {
+    data: Json<JobData>,
+}
+
 async fn handle_webhooks_wrapper(
     rx: Receiver<EventData>,
     embedding_api: EmbeddingApi,
@@ -282,7 +292,7 @@ async fn handle_webhooks(
                             Vector::from(embedding_api.generate_embedding(issue_text).await?);
 
                         let closest_issues: Vec<ClosestIssue> = sqlx::query_as(
-                            "select title, number, html_url, body, 1 - (embedding <=> $1) as cosine_similarity from issues order by embedding <=> $1 LIMIT 3",
+                            "select title, number, html_url, 1 - (embedding <=> $1) as cosine_similarity from issues order by embedding <=> $1 LIMIT 3",
                         )
                             .bind(embedding.clone())
                             .fetch_all(&pool)
@@ -395,17 +405,17 @@ async fn handle_webhooks(
                     source = repository.source.to_string()
                 );
                 info!(parent: &span, "indexing started");
-                let from_page = sqlx::query!(
-                    "select page from jobs where repository_id = $1",
+                let job = sqlx::query_as!(
+                    Job,
+                    r#"select data as "data: Json<JobData>" from jobs where repository_id = $1"#,
                     repository.repo_id
                 )
-                .map(|row| row.page + 1)
                 .fetch_optional(&pool)
                 .await?;
-                let issues = github_api.get_issues(from_page, repository.clone());
+                let from_issues_page = job.as_ref().map(|j| j.data.issues_page + 1).unwrap_or(1);
+                let issues = github_api.get_issues(from_issues_page, repository.clone());
                 pin_mut!(issues);
-                while let Some(issue) = issues.next().await {
-                    let (issue, page) = issue?;
+                while let Some((issue, page)) = issues.try_next().await? {
                     let embedding_api = embedding_api.clone();
                     let pool = pool.clone();
                     let source = repository.source.to_string();
@@ -461,19 +471,19 @@ async fn handle_webhooks(
                         qb.build().execute(&pool).await?;
                     }
                     if let Some(page) = page {
-                        sqlx::query!(
-                            r#"insert into jobs (repository_id, page)
-                               values ($1, $2)
+                        sqlx::query(
+                                r#"insert into jobs (repository_id, data)
+                               values ($1, jsonb_build_object('issues_page', $2))
                                on conflict (repository_id)
                                do update
                                set
-                                   page = excluded.page,
+                                   data = jsonb_set(jobs.data, '{issues_page}', to_jsonb($2::int), true),
                                    updated_at = current_timestamp"#,
-                            repository.repo_id,
-                            page,
-                        )
-                        .execute(&pool)
-                        .await?;
+                            )
+                            .bind(&repository.repo_id)
+                            .bind(page)
+                            .execute(&pool)
+                            .await?;
                     }
                 }
                 sqlx::query!(

issue-bot/src/routes.rs

@@ -104,6 +104,7 @@ struct IssueData {
     html_url: String,
     id: i64,
     number: i32,
+    #[serde(default)]
     pull_request: Option<PullRequest>,
     title: String,
     url: String,
@@ -462,8 +463,8 @@ mod tests {
         };
         let mut app = app(state);
 
-        let payload_body = r#"{"action":"opened","pull_request":{"title":"my great contribution to the world","body":"superb work, isnt it","id":4321,"number":5,"html_url":"https://github.com/huggingface/lor-e/5", "url":"https://github.com/api/huggingface/lor-e/5"}}"#;
-        let sig = "sha256=0f66c678489c6ab39ba9b5a3dfcb957b2bf3b98aebf0872e8cebbc041ff71307";
+        let payload_body = r#"{"action":"opened","issue":{"title":"my great contribution to the world","body":"superb work, isnt it","id":4321,"number":5,"html_url":"https://github.com/huggingface/lor-e/5", "url":"https://github.com/api/huggingface/lor-e/5"}}"#;
+        let sig = "sha256=930e7b9a7cca3f85bc49a693f1d9105ca32bb15f13cda8871aaf79bba27c48cc";
 
         let response = app
             .borrow_mut()
@@ -480,8 +481,8 @@ mod tests {
 
         assert_eq!(response.status(), StatusCode::OK);
 
-        let payload_body = r#"{"action":"submitted","review":{"body":"test review","id":1234,"url":"https://github.com/huggingface/lor-e/5#comment-123"},"pull_request":{"title":"my great contribution to the world","body":"superb work, isnt it","id":4321,"number":5,"html_url":"https://github.com/huggingface/lor-e/5", "url":"https://github.com/api/huggingface/lor-e/5"}}"#;
-        let sig = "sha256=4a312de764757f5d18610c183602337f0c766791e8a886120302549c00988bba";
+        let payload_body = r#"{"action":"created","comment":{"body":"test review","id":1234,"url":"https://github.com/huggingface/lor-e/5#comment-123"}, "issue":{"title":"my great contribution to the world","body":"superb work, isnt it","id":4321,"number":5,"html_url":"https://github.com/huggingface/lor-e/5", "url":"https://github.com/api/huggingface/lor-e/5"}}"#;
+        let sig = "sha256=cb81358e382b98e54789541ec9deeef909d89437311b9c626b13b40662f5ef52";
 
         let response = app
             .oneshot(

issue-bot/src/slack.rs

@@ -59,14 +59,11 @@ impl Slack {
         closest_issues: &[ClosestIssue],
     ) -> Result<(), SlackError> {
         let mut msg = vec![format!(
-            "Closest issues for {} <{}|#{}>:\n```{}```",
-            issue.title, issue.number, issue.html_url, issue.body
+            "Closest issues for {} (<{}|#{}>):\n```{}```",
+            issue.title, issue.html_url, issue.number, issue.body
         )];
         for ci in closest_issues {
-            msg.push(format!(
-                "- {} (<{}|#{}>):\n```{}```",
-                ci.title, ci.html_url, ci.number, ci.body
-            ));
+            msg.push(format!("- {} (<{}|#{}>)", ci.title, ci.html_url, ci.number));
         }
         let body = SlackBody::new(&self.channel, msg.join("\n"));
         self.client