From bf3e417acc167583aaef365e92eed6c216e007ec Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Tue, 20 Feb 2024 19:42:05 +0100
Subject: [PATCH] fix(python): handle recursive python imports with loop

---
 .../windmill-parser-py-imports/src/lib.rs     |  8 ++++-
 .../windmill-parser-py-imports/tests/tests.rs | 31 +++++++++++++++++--
 .../windmill-worker/src/python_executor.rs    |  3 ++
 backend/windmill-worker/src/worker.rs         |  3 ++
 4 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/backend/parsers/windmill-parser-py-imports/src/lib.rs b/backend/parsers/windmill-parser-py-imports/src/lib.rs
index f1183fe812..ee66277f68 100644
--- a/backend/parsers/windmill-parser-py-imports/src/lib.rs
+++ b/backend/parsers/windmill-parser-py-imports/src/lib.rs
@@ -145,6 +145,7 @@ pub async fn parse_python_imports(
     w_id: &str,
     path: &str,
     db: &Pool<Postgres>,
+    already_visited: &mut Vec<String>,
 ) -> error::Result<Vec<String>> {
     let find_requirements = code
         .lines()
@@ -192,7 +193,12 @@ pub async fn parse_python_imports(
                 .fetch_optional(db)
                 .await?
                 .unwrap_or_else(|| "".to_string());
-                parse_python_imports(&code, w_id, &rpath, db).await?
+                if already_visited.contains(&rpath) {
+                    vec![]
+                } else {
+                    already_visited.push(rpath.clone());
+                    parse_python_imports(&code, w_id, &rpath, db, already_visited).await?
+                }
             } else {
                 vec![replace_import(n.to_string())]
             };
diff --git a/backend/parsers/windmill-parser-py-imports/tests/tests.rs b/backend/parsers/windmill-parser-py-imports/tests/tests.rs
index 3f2212d58f..436493e132 100644
--- a/backend/parsers/windmill-parser-py-imports/tests/tests.rs
+++ b/backend/parsers/windmill-parser-py-imports/tests/tests.rs
@@ -18,7 +18,15 @@ def main():
     pass
 
 ";
-        let r = parse_python_imports(code, "test-workspace", "f/foo/bar", &db).await?;
+        let mut already_visited = vec![];
+        let r = parse_python_imports(
+            code,
+            "test-workspace",
+            "f/foo/bar",
+            &db,
+            &mut already_visited,
+        )
+        .await?;
         // println!("{}", serde_json::to_string(&r)?);
         assert_eq!(r, vec!["matplotlib", "wmill", "zanzibar"]);
         Ok(())
@@ -42,7 +50,15 @@ def main():
     pass
 
 ";
-        let r = parse_python_imports(code, "test-workspace", "f/foo/bar", &db).await?;
+        let mut already_visited = vec![];
+        let r = parse_python_imports(
+            code,
+            "test-workspace",
+            "f/foo/bar",
+            &db,
+            &mut already_visited,
+        )
+        .await?;
         println!("{}", serde_json::to_string(&r)?);
         assert_eq!(r, vec!["burkina=0.4", "nigeria"]);
 
@@ -63,7 +79,16 @@ def main():
     pass
 
 ";
-        let r = parse_python_imports(code, "test-workspace", "f/foo/bar", &db).await?;
+        let mut already_visited = vec![];
+
+        let r = parse_python_imports(
+            code,
+            "test-workspace",
+            "f/foo/bar",
+            &db,
+            &mut already_visited,
+        )
+        .await?;
         println!("{}", serde_json::to_string(&r)?);
         assert_eq!(
             r,
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
index 2542c48b93..658f7f0208 100644
--- a/backend/windmill-worker/src/python_executor.rs
+++ b/backend/windmill-worker/src/python_executor.rs
@@ -631,11 +631,14 @@ async fn handle_python_deps(
     let requirements = match requirements_o {
         Some(r) => r,
         None => {
+            let mut already_visited = vec![];
+
             let requirements = windmill_parser_py_imports::parse_python_imports(
                 inner_content,
                 w_id,
                 script_path,
                 db,
+                &mut already_visited,
             )
             .await?
             .join("\n");
diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs
index 7bbb6c6ec9..82b11b4349 100644
--- a/backend/windmill-worker/src/worker.rs
+++ b/backend/windmill-worker/src/worker.rs
@@ -3964,11 +3964,14 @@ async fn capture_dependency_job(
             let reqs = if raw_deps {
                 job_raw_code.to_string()
             } else {
+                let mut already_visited = vec![];
+
                 windmill_parser_py_imports::parse_python_imports(
                     job_raw_code,
                     &w_id,
                     script_path,
                     &db,
+                    &mut already_visited,
                 )
                 .await?
                 .join("\n")