feat: introduce deep models chain detection

richardhapb · richardhapb · commit 71a6b22fd6fa · 2026-02-06T22:16:39.000-03:00
prefetchs like `relation__child__childchild` are supported
where `child` is a child relation from `relation` and when
we access to relation.child, the prefetch is detected
for the relation, no warnings are triggered, same case
for child.childchild
diff --git a/README.md b/README.md
@@ -178,17 +178,13 @@ for user in users:
 
 * Interprocedural analysis requires type hints on QuerySet parameters
 * Limited understanding of:
-  * `Prefetch` objects
+  * `Prefetch` objects with basic access support (e.g. `Prefetch("related_model")`)
   * `annotate`, `aggregate`
   * complex custom managers
-* No complete capture of implicit prefetchs in django chains when has many. e.g.
-  `prefetch_releated(chat.users__profile)`, then iterate over `chat.users` and
-  access to `profile`. This will raise a warning.
 
 ## Roadmap
 
-* Prefetch object support
-* Capture of implicits prefetchs in a root instance
+* Complete Prefetch object support
 * Custom queryset method summaries
 * Templates integration
 
diff --git a/crates/django-check_semantic/src/ir/binding.rs b/crates/django-check_semantic/src/ir/binding.rs
@@ -1,6 +1,8 @@
 //! Binding-level intermediate representation for tracking variable states.
 
-use std::collections::HashSet;
+use std::{collections::HashMap, str::Split};
+
+use crate::{ModelGraph, ir::model::ModelDef};
 
 #[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)]
 pub struct DjangoSymbolId(pub u32);
@@ -14,21 +16,21 @@ impl DjangoSymbolId {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct QuerySetState {
     pub model_name: String,
-    pub prefetched_relations: HashSet<String>,
+    pub prefetched_relations: HashMap<String, QuerySetState>,
     pub is_values_query: bool,
 }
 
 impl QuerySetState {
     pub fn new(model_name: String) -> Self {
         Self {
             model_name,
-            prefetched_relations: HashSet::new(),
+            prefetched_relations: HashMap::new(),
             is_values_query: false,
         }
     }
 
     pub fn is_access_safe(&self, relation: &str) -> bool {
-        self.is_values_query || self.prefetched_relations.contains(relation)
+        self.is_values_query || self.prefetched_relations.contains_key(relation)
     }
 }
 
@@ -50,46 +52,78 @@ pub enum DjangoSymbolKind {
     Unknown,
 }
 
-/// Parse Django relation field syntax (e.g., "author__profile" -> ["author", "profile"])
-pub fn parse_relation_fields(fields: &[String]) -> Vec<String> {
-    fields
-        .iter()
-        .flat_map(|field| field.split("__"))
-        .map(|s| s.to_string())
-        .collect()
+/// Parse Django relation field syntax (e.g., "author__profile" ->
+/// [PrefetchedRelation(QuerySetState)] with `profile` prefetched)
+pub fn parse_relation_fields(
+    model: &ModelDef,
+    model_graph: &ModelGraph,
+    fields: &[String],
+) -> HashMap<String, QuerySetState> {
+    let mut relations = HashMap::new();
+    for field in fields.iter() {
+        let parts = field.split("__");
+        if let Some((literal, relation)) = bind_relation(model, model_graph, parts) {
+            relations.insert(literal, relation);
+        }
+    }
+
+    relations
+}
+
+fn bind_relation(
+    model: &ModelDef,
+    model_graph: &ModelGraph,
+    mut parts: Split<'_, &str>,
+) -> Option<(String, QuerySetState)> {
+    if let Some(base) = parts.next()
+        && let Some(relation) = model_graph.get_relation(&model.name, base)
+        && let Some(related_model) = model_graph.get(relation)
+    {
+        let mut qs = QuerySetState::new(related_model.name.to_string());
+
+        // Insert child relation
+        if let Some((literal, relation)) = bind_relation(related_model, model_graph, parts) {
+            qs.prefetched_relations.insert(literal, relation);
+        };
+        return Some((base.to_string(), qs));
+    }
+
+    None
 }
 
 #[cfg(test)]
 mod tests {
+    use crate::Parser;
+
     use super::*;
 
-    #[test]
-    fn parse_single_field() {
-        let fields = vec!["ticker".to_string()];
-        assert_eq!(parse_relation_fields(&fields), vec!["ticker"]);
-    }
+    fn get_graph() -> ModelGraph {
+        let source = r#"
+class User(Model):
+    pass
 
-    #[test]
-    fn parse_nested_field() {
-        let fields = vec!["theoanalysis__ticker".to_string()];
-        assert_eq!(
-            parse_relation_fields(&fields),
-            vec!["theoanalysis", "ticker"]
-        );
-    }
+class Photo(Model):
+    user = models.ForeignKey(User, related_name="photos")
 
-    #[test]
-    fn parse_deeply_nested() {
-        let fields = vec!["a__b__c__d".to_string()];
-        assert_eq!(parse_relation_fields(&fields), vec!["a", "b", "c", "d"]);
-    }
 
-    #[test]
-    fn parse_multiple_fields() {
-        let fields = vec!["ticker__sector".to_string(), "analysis__report".to_string()];
-        assert_eq!(
-            parse_relation_fields(&fields),
-            vec!["ticker", "sector", "analysis", "report"]
-        );
+class Order(Model):
+    user = models.ForeignKey(User, related_name="orders")
+
+class Order(Model):
+    user = models.ForeignKey(User, related_name="orders")
+
+class Sale(Model):
+    order = models.ForeignKey(Sale, related_name="sales")
+
+class Transaction(Model):
+    sale = models.ForeignKey(Sale, related_name="transactions")
+
+
+users = User.objects.all().prefetch_related(Prefetch("orders"))
+print([user.orders for user in users])
+        "#;
+
+        let parser = Parser::new();
+        parser.build_graph(source, "test.py").unwrap()
     }
 }
diff --git a/crates/django-check_semantic/src/parser.rs b/crates/django-check_semantic/src/parser.rs
@@ -90,29 +90,33 @@ impl Parser {
 
         for entry in Self::python_files(dir) {
             let filename = Self::relative_path(dir, entry.path());
-            let source = match fs::read_to_string(entry.path()) {
+            let source = fs::read_to_string(entry.path())?;
+            let graph = match self.build_graph(&source, &filename) {
                 Ok(source) => source,
                 Err(err) => {
                     warn!(%err, source=%filename, "parsing graph");
                     continue;
                 }
             };
-            let parsed = match self.parse_module(&source) {
-                Ok(parsed) => parsed,
-                Err(err) => {
-                    warn!(%err, source=%filename, "parsing graph");
-                    continue;
-                }
-            };
-
-            let mut pass = ModelGraphPass::new(&filename, &source);
-            let graph = pass.run(parsed.syntax());
             combined_graph.merge(graph);
         }
 
         Ok(combined_graph)
     }
 
+    pub fn build_graph(
+        &self,
+        source: &str,
+        filename: &str,
+    ) -> Result<ModelGraph, SourceParseError> {
+        let parsed = self
+            .parse_module(source)
+            .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidInput, "parsing module"))?;
+
+        let mut pass = ModelGraphPass::new(&filename, &source);
+        Ok(pass.run(parsed.syntax()))
+    }
+
     pub fn extract_functions(&self, dir: &Path) -> Result<Vec<QueryFunction>, SourceParseError> {
         let mut functions = Vec::new();
         for entry in Self::python_files(dir) {
diff --git a/crates/django-check_semantic/src/passes/n_plus_one.rs b/crates/django-check_semantic/src/passes/n_plus_one.rs
@@ -154,21 +154,25 @@ impl<'a> QuerySetResolver<'a> {
                     // Accessing `user.orders` -> Returns a QuerySet for the related model
                     // Accessing an attribute on a queryset (e.g., `orders.users`)
                     // If it's a relation on the model, return a QuerySet of the related model
-                    if let Some(related_model) = self
-                        .model_graph
-                        .get_relation(&state.model_name, &attr.attr.id)
-                    {
-                        let mut new_state = QuerySetState::new(related_model.to_string());
-                        // FIXME: Currently we use the current relations to be inherited by the related
-                        // model, probably a better idea is to detect dynamically the appropiates
-                        // prefetched relations but for now this is ok.
-                        new_state
-                            .prefetched_relations
-                            .extend(state.prefetched_relations);
-                        return Some(DjangoSymbol::QuerySet(new_state));
+                    // firs try to get the prefetched query set state, if it is not prefetched, create
+                    // an empty query set state
+                    match state.prefetched_relations.get(&attr.attr.id.to_string()) {
+                        Some(related_state) => {
+                            return Some(DjangoSymbol::QuerySet(related_state.clone()));
+                        }
+                        None => {
+                            if let Some(related_model) = self
+                                .model_graph
+                                .get_relation(&state.model_name, &attr.attr.id)
+                            {
+                                // Empty state, no prefetch related detected before for this model
+                                let new_state = QuerySetState::new(related_model.to_string());
+                                return Some(DjangoSymbol::QuerySet(new_state));
+                            }
+                            // Could be accessing a regular field, return None
+                            return None;
+                        }
                     }
-                    // Could be accessing a regular field, return None
-                    return None;
                 }
             }
         }
@@ -223,9 +227,15 @@ impl<'a> QuerySetResolver<'a> {
                 }
             }
 
-            // Helper from binding.rs
-            let relations = crate::ir::binding::parse_relation_fields(&literal_fields);
-            state.prefetched_relations.extend(relations);
+            // Resolve prefetched relations
+            if let Some(model) = self.model_graph.get(&state.model_name) {
+                let relations = crate::ir::binding::parse_relation_fields(
+                    model,
+                    self.model_graph,
+                    &literal_fields,
+                );
+                state.prefetched_relations.extend(relations);
+            }
         }
     }
 
@@ -838,8 +848,11 @@ for item in qs:
     #[test]
     fn safe_with_prefetch() {
         let source = r#"
+class App(Model):
+    pass
+
 class User(Model):
-    related_field = models.ForeignKey("some.App")
+    related_field = models.ForeignKey("App")
 
 qs = User.objects.filter(active=True).prefetch_related('related_field')
 for item in qs:
@@ -986,6 +999,29 @@ for p in performances:
         assert_eq!(diags.len(), 1);
     }
 
+    #[test]
+    fn allow_using_deep_chained_prefetch() {
+        let source = r#"
+class Analysis(Model):
+    pattern = models.ForeignKey("Pattern", related_name="analyses")
+
+class Performance(Model):
+    analysis = models.ForeignKey("Analysis", related_name="performances")
+
+class Pattern(Model):
+    name = models.CharField(max_length=20)
+
+
+patterns = Pattern.objects.prefetch_related("analyses__performances").all()
+
+for p in patterns:
+    for a in p.analyses:
+       print(a.performances)  # safe
+"#;
+        let diags = run_pass(source);
+        assert!(diags.is_empty());
+    }
+
     #[test]
     fn second_level_related_access() {
         let source = r#"
