feat(fmt): add quick-fix for missing indexes when relationMode = "pri…

…sma" (#3431)

Co-authored-by: Tom tom@tomhoule.com

prisma-fmt/src/code_actions.rs

@@ -1,10 +1,7 @@
 mod relations;
 
 use lsp_types::{CodeActionOrCommand, CodeActionParams, Diagnostic};
-use psl::{
-    parser_database::{ast, walkers::RefinedRelationWalker, ParserDatabase, SourceFile},
-    Diagnostics,
-};
+use psl::parser_database::{ast, walkers::RefinedRelationWalker, SourceFile};
 use std::sync::Arc;
 
 pub(crate) fn empty_code_actions() -> Vec<CodeActionOrCommand> {
@@ -16,22 +13,38 @@ pub(crate) fn available_actions(schema: String, params: CodeActionParams) -> Vec
 
     let file = SourceFile::new_allocated(Arc::from(schema.into_boxed_str()));
 
-    let db = {
-        let mut diag = Diagnostics::new();
-        ParserDatabase::new(file.clone(), &mut diag)
-    };
+    let validated_schema = psl::validate(file);
 
-    for relation in db.walk_relations() {
+    for relation in validated_schema.db.walk_relations() {
         if let RefinedRelationWalker::Inline(relation) = relation.refine() {
             let complete_relation = match relation.as_complete() {
                 Some(relation) => relation,
                 None => continue,
             };
 
-            relations::add_referenced_side_unique(&mut actions, &params, file.as_str(), complete_relation);
+            relations::add_referenced_side_unique(
+                &mut actions,
+                &params,
+                validated_schema.db.source(),
+                complete_relation,
+            );
 
             if relation.is_one_to_one() {
-                relations::add_referencing_side_unique(&mut actions, &params, file.as_str(), complete_relation);
+                relations::add_referencing_side_unique(
+                    &mut actions,
+                    &params,
+                    validated_schema.db.source(),
+                    complete_relation,
+                );
+            }
+
+            if validated_schema.relation_mode().is_prisma() {
+                relations::add_index_for_relation_fields(
+                    &mut actions,
+                    &params,
+                    validated_schema.db.source(),
+                    complete_relation.referencing_field(),
+                );
             }
         }
     }

prisma-fmt/src/code_actions/relations.rs

@@ -1,7 +1,9 @@
-use lsp_types::{CodeAction, CodeActionKind, CodeActionOrCommand, CodeActionParams, Range, TextEdit, WorkspaceEdit};
+use lsp_types::{
+    CodeAction, CodeActionKind, CodeActionOrCommand, CodeActionParams, Diagnostic, Range, TextEdit, WorkspaceEdit,
+};
 use psl::parser_database::{
     ast::WithSpan,
-    walkers::{CompleteInlineRelationWalker, ModelWalker, ScalarFieldWalker},
+    walkers::{CompleteInlineRelationWalker, ModelWalker, RelationFieldWalker, ScalarFieldWalker},
 };
 use std::collections::HashMap;
 
@@ -65,7 +67,13 @@ pub(super) fn add_referencing_side_unique(
         _ => (),
     }
 
-    let text = create_missing_unique(schema, relation.referencing_model(), relation.referencing_fields());
+    let attribute_name = "unique";
+    let text = create_missing_attribute(
+        schema,
+        relation.referencing_model(),
+        relation.referencing_fields(),
+        attribute_name,
+    );
 
     let mut changes = HashMap::new();
     changes.insert(params.text_document.uri.clone(), vec![text]);
@@ -151,7 +159,13 @@ pub(super) fn add_referenced_side_unique(
         _ => (),
     }
 
-    let text = create_missing_unique(schema, relation.referenced_model(), relation.referenced_fields());
+    let attribute_name = "unique";
+    let text = create_missing_attribute(
+        schema,
+        relation.referenced_model(),
+        relation.referenced_fields(),
+        attribute_name,
+    );
 
     let mut changes = HashMap::new();
     changes.insert(params.text_document.uri.clone(), vec![text]);
@@ -180,13 +194,109 @@ pub(super) fn add_referenced_side_unique(
     actions.push(CodeActionOrCommand::CodeAction(action));
 }
 
-fn create_missing_unique<'a>(
+/// For schema's with emulated relations,
+/// If the referenced side of the relation does not point to a unique
+/// constraint, the action adds the attribute.
+///
+/// If referencing a single field:
+///
+/// ```ignore
+/// model A {
+///     id      Int @id
+///     field1  B   @relation(fields: [bId], references: [id])
+///                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ // Warn
+///     bId     Int
+///
+///     // <- suggest @@index([bId]) here
+/// }
+///
+/// model B {
+///     id Int @id
+///     as A[]
+/// }
+/// ```
+///
+/// If referencing multiple fields:
+///
+/// ```ignore
+/// model A {
+///     id      Int @id
+///     field1  B   @relation(fields: [bId1, bId2, bId3], references: [id1, id2, id3])
+///                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ // Warn
+///     bId1    Int
+///     bId2    Int
+///     bId3    Int
+///
+///     // <- suggest @@index([bId1, bId2, bId3]) here
+/// }
+///
+/// model B {
+///     id1 Int
+///     id2 Int
+///     id3 Int
+///     as  A[]
+///
+///     @@id([id1, id2, id3])
+/// }
+/// ```
+pub(super) fn add_index_for_relation_fields(
+    actions: &mut Vec<CodeActionOrCommand>,
+    params: &CodeActionParams,
+    schema: &str,
+    relation: RelationFieldWalker<'_>,
+) {
+    let Some(fields) = relation.fields() else { return; };
+    if relation.model().indexes().any(|index| {
+        index
+            .fields()
+            .zip(fields.clone())
+            .all(|(index_field, relation_field)| index_field.field_id() == relation_field.field_id())
+    }) {
+        return;
+    }
+
+    let attribute_name = "index";
+    let (new_text, range) = create_block_attribute(schema, relation.model(), fields, attribute_name);
+    let text = TextEdit { range, new_text };
+
+    let mut changes = HashMap::new();
+    changes.insert(params.text_document.uri.clone(), vec![text]);
+
+    let edit = WorkspaceEdit {
+        changes: Some(changes),
+        ..Default::default()
+    };
+
+    let Some(span_diagnostics) = super::diagnostics_for_span(
+        schema,
+        &params.context.diagnostics,
+        relation.relation_attribute().unwrap().span(),
+    ) else { return; };
+
+    let diagnostics = span_diagnostics
+        .into_iter()
+        .filter(|diag| diag.message.contains("relationMode = \"prisma\""))
+        .collect::<Vec<Diagnostic>>();
+
+    let action = CodeAction {
+        title: String::from("Add an index for the relation's field(s)"),
+        kind: Some(CodeActionKind::QUICKFIX),
+        edit: Some(edit),
+        diagnostics: Some(diagnostics),
+        ..Default::default()
+    };
+
+    actions.push(CodeActionOrCommand::CodeAction(action))
+}
+
+fn create_missing_attribute<'a>(
     schema: &str,
     model: ModelWalker<'a>,
     mut fields: impl ExactSizeIterator<Item = ScalarFieldWalker<'a>> + 'a,
+    attribute_name: &str,
 ) -> TextEdit {
     let (new_text, range) = if fields.len() == 1 {
-        let new_text = String::from(" @unique");
+        let new_text = format!(" @{attribute_name}");
 
         let field = fields.next().unwrap();
         let position = crate::position_after_span(field.ast_field().span(), schema);
@@ -198,26 +308,34 @@ fn create_missing_unique<'a>(
 
         (new_text, range)
     } else {
-        let fields = fields.map(|f| f.name()).collect::<Vec<_>>().join(", ");
-
-        let indentation = model.indentation();
-        let newline = model.newline();
+        let (new_text, range) = create_block_attribute(schema, model, fields, attribute_name);
+        (new_text, range)
+    };
 
-        let separator = if model.ast_model().attributes.is_empty() {
-            newline.as_ref()
-        } else {
-            ""
-        };
+    TextEdit { range, new_text }
+}
 
-        let new_text = format!("{separator}{indentation}@@unique([{fields}]){newline}}}");
+fn create_block_attribute<'a>(
+    schema: &str,
+    model: ModelWalker<'a>,
+    fields: impl ExactSizeIterator<Item = ScalarFieldWalker<'a>> + 'a,
+    attribute_name: &str,
+) -> (String, Range) {
+    let fields = fields.map(|f| f.name()).collect::<Vec<_>>().join(", ");
 
-        let start = crate::offset_to_position(model.ast_model().span().end - 1, schema);
-        let end = crate::offset_to_position(model.ast_model().span().end, schema);
+    let indentation = model.indentation();
+    let newline = model.newline();
+    let separator = if model.ast_model().attributes.is_empty() {
+        newline.as_ref()
+    } else {
+        ""
+    };
+    let new_text = format!("{separator}{indentation}@@{attribute_name}([{fields}]){newline}}}");
 
-        let range = Range { start, end };
+    let start = crate::offset_to_position(model.ast_model().span().end - 1, schema);
+    let end = crate::offset_to_position(model.ast_model().span().end, schema);
 
-        (new_text, range)
-    };
+    let range = Range { start, end };
 
-    TextEdit { range, new_text }
+    (new_text, range)
 }

prisma-fmt/tests/code_actions/scenarios/relation_mode_prisma_missing_index/diagnostics.json

@@ -0,0 +1,16 @@
+[
+    {
+        "range": {
+            "start": {
+                "line": 19,
+                "character": 20
+            },
+            "end": {
+                "line": 19,
+                "character": 63
+            }
+        },
+        "severity": 1,
+        "message": "With `relationMode = \"prisma\"`, no foreign keys are used, so relation fields will not benefit from the index usually created by the relational database under the hood. This can lead to poor performance when querying these fields. We recommend adding an index manually. Learn more at https://pris.ly/d/relation-mode#indexes"
+    }
+]

prisma-fmt/tests/code_actions/scenarios/relation_mode_prisma_missing_index/result.json

@@ -0,0 +1,41 @@
+[
+  {
+    "title": "Add an index for the relation's field(s)",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 19,
+            "character": 20
+          },
+          "end": {
+            "line": 19,
+            "character": 63
+          }
+        },
+        "severity": 1,
+        "message": "With `relationMode = \"prisma\"`, no foreign keys are used, so relation fields will not benefit from the index usually created by the relational database under the hood. This can lead to poor performance when querying these fields. We recommend adding an index manually. Learn more at https://pris.ly/d/relation-mode#indexes"
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file:///path/to/schema.prisma": [
+          {
+            "range": {
+              "start": {
+                "line": 20,
+                "character": 0
+              },
+              "end": {
+                "line": 20,
+                "character": 1
+              }
+            },
+            "newText": "\n    @@index([userId])\n}"
+          }
+        ]
+      }
+    }
+  }
+]

prisma-fmt/tests/code_actions/scenarios/relation_mode_prisma_missing_index/schema.prisma

@@ -0,0 +1,21 @@
+generator client {
+    provider = "prisma-client-js"
+    previewFeatures = ["referentialIntegrity"]
+}
+
+datasource db {
+    provider = "mysql"
+    url = env("TEST_DATABASE_URL")
+    relationMode = "prisma"
+}
+
+model SomeUser {
+    id Int @id
+    posts Post[]
+}
+
+model Post {
+    id Int @id
+    userId Int
+    user SomeUser @relation(fields: [userId], references: [id])
+}

prisma-fmt/tests/code_actions/tests.rs

@@ -23,4 +23,5 @@ scenarios! {
     one_to_one_referencing_side_misses_unique_single_field
     one_to_one_referencing_side_misses_unique_compound_field
     one_to_one_referencing_side_misses_unique_compound_field_indentation_four_spaces
+    relation_mode_prisma_missing_index
 }