Use trigram search for single word queries

ChrisPenner · ChrisPenner · commit 1ae58eb54018 · 2024-07-29T22:52:10.000-07:00
diff --git a/src/Share/Postgres/Search/DefinitionSearch/Queries.hs b/src/Share/Postgres/Search/DefinitionSearch/Queries.hs
@@ -6,8 +6,9 @@ module Share.Postgres.Search.DefinitionSearch.Queries
     claimUnsyncedRelease,
     insertDefinitionDocuments,
     cleanIndexForRelease,
-    defNameInfixSearch,
-    definitionSearch,
+    defNameCompletionSearch,
+    definitionTokenSearch,
+    definitionNameSearch,
     DefnNameSearchFilter (..),
     -- Exported for logging/debugging
     searchTokensToTsQuery,
@@ -268,9 +269,9 @@ data DefnNameSearchFilter
   | ReleaseFilter ReleaseId
   | UserFilter UserId
 
--- | Find definitions whose name contains the query.
-defNameInfixSearch :: Maybe UserId -> Maybe DefnNameSearchFilter -> Query -> Limit -> Transaction e [(ProjectId, ReleaseId, Name, TermOrTypeTag)]
-defNameInfixSearch mayCaller mayFilter (Query query) limit = do
+-- | Find names which would be valid completions for the given query.
+defNameCompletionSearch :: Maybe UserId -> Maybe DefnNameSearchFilter -> Query -> Limit -> Transaction e [(ProjectId, ReleaseId, Name, TermOrTypeTag)]
+defNameCompletionSearch mayCaller mayFilter (Query query) limit = do
   let filters = case mayFilter of
         Just (ProjectFilter projId) -> [sql| AND doc.project_id = #{projId} |]
         Just (ReleaseFilter relId) -> [sql| AND doc.release_id = #{relId} |]
@@ -309,8 +310,9 @@ defNameInfixSearch mayCaller mayFilter (Query query) limit = do
     -- Names are stored in absolute form, but we usually work with them in relative form.
     <&> over (traversed . _3) Name.makeRelative
 
-definitionSearch :: Maybe UserId -> Maybe DefnNameSearchFilter -> Limit -> Set (DefnSearchToken (Either Name ShortHash)) -> Maybe Arity -> Transaction e [(ProjectId, ReleaseId, Name, TermOrTypeSummary)]
-definitionSearch mayCaller mayFilter limit searchTokens preferredArity = do
+-- | Perform a type search for the given tokens.
+definitionTokenSearch :: Maybe UserId -> Maybe DefnNameSearchFilter -> Limit -> Set (DefnSearchToken (Either Name ShortHash)) -> Maybe Arity -> Transaction e [(ProjectId, ReleaseId, Name, TermOrTypeSummary)]
+definitionTokenSearch mayCaller mayFilter limit searchTokens preferredArity = do
   let filters = case mayFilter of
         Just (ProjectFilter projId) -> [sql| AND doc.project_id = #{projId} |]
         Just (ReleaseFilter relId) -> [sql| AND doc.release_id = #{relId} |]
@@ -355,3 +357,40 @@ definitionSearch mayCaller mayFilter limit searchTokens preferredArity = do
             Aeson.Error err -> unrecoverableError $ FailedToDecodeMetadata v (Text.pack err)
             Aeson.Success summary -> pure summary
       )
+
+-- | Perform a fuzzy trigram search on definition names
+definitionNameSearch :: Maybe UserId -> Maybe DefnNameSearchFilter -> Limit -> Query -> Transaction e [(ProjectId, ReleaseId, Name, TermOrTypeSummary)]
+definitionNameSearch mayCaller mayFilter limit (Query query) = do
+  let filters = case mayFilter of
+        Just (ProjectFilter projId) -> [sql| AND doc.project_id = #{projId} |]
+        Just (ReleaseFilter relId) -> [sql| AND doc.release_id = #{relId} |]
+        Just (UserFilter userId) -> [sql| AND p.owner_user_id = #{userId} |]
+        Nothing -> mempty
+  rows <-
+    queryListRows @(ProjectId, ReleaseId, Name, Hasql.Jsonb)
+      [sql|
+    WITH matches_deduped_by_project(project_id, release_id, name, metadata) AS (
+      SELECT DISTINCT ON (doc.project_id, doc.name) doc.project_id, doc.release_id, doc.name, doc.metadata FROM global_definition_search_docs doc
+        JOIN projects p ON p.id = doc.project_id
+        JOIN project_releases r ON r.id = doc.release_id
+        WHERE
+          -- Adjust the similarity threshold as needed (between 0 and 1)
+          word_similarity(#{query}, doc.name) >= 0.5
+        AND (NOT p.private OR (#{mayCaller} IS NOT NULL AND EXISTS (SELECT FROM accessible_private_projects pp WHERE pp.user_id = #{mayCaller} AND pp.project_id = p.id)))
+          ^{filters}
+        ORDER BY doc.project_id, doc.name, r.major_version, r.minor_version, r.patch_version
+    ) SELECT m.project_id, m.release_id, m.name, m.metadata
+        FROM matches_deduped_by_project m
+        -- Score matches by name similarity
+        ORDER BY word_similarity(#{query}, m.name) DESC
+        LIMIT #{limit}
+  |]
+  rows
+    & over (traversed . _3) Name.makeRelative
+    & traverseOf
+      (traversed . _4)
+      ( \(Hasql.Jsonb v) -> do
+          case fromJSON v of
+            Aeson.Error err -> unrecoverableError $ FailedToDecodeMetadata v (Text.pack err)
+            Aeson.Success summary -> pure summary
+      )
diff --git a/src/Share/Web/Share/Impl.hs b/src/Share/Web/Share/Impl.hs
@@ -383,7 +383,7 @@ searchDefinitionNamesEndpoint ::
 searchDefinitionNamesEndpoint callerUserId query@(Query queryText) mayLimit userFilter projectFilter releaseFilter = do
   filter <- runMaybeT $ resolveProjectAndReleaseFilter projectFilter releaseFilter <|> resolveUserFilter (IDs.unPrefix <$> userFilter)
   matches <-
-    (PG.runTransaction $ DDQ.defNameInfixSearch callerUserId filter query limit)
+    (PG.runTransaction $ DDQ.defNameCompletionSearch callerUserId filter query limit)
       <&> ordNubOn (view _3) . (mapMaybe $ traverseOf _3 (rewriteMatches queryText))
   let response = matches <&> \(_projId, _releaseId, name, tag) -> DefinitionNameSearchResult name tag
   pure response
@@ -438,27 +438,39 @@ searchDefinitionsEndpoint ::
 searchDefinitionsEndpoint callerUserId (Query query) mayLimit userFilter projectFilter releaseFilter = do
   Logging.logInfoText $ "definition-search-query: " <> query
   filter <- runMaybeT $ resolveProjectAndReleaseFilter projectFilter releaseFilter <|> resolveUserFilter (IDs.unPrefix <$> userFilter)
-  case DefinitionSearch.queryToTokens query of
-    Left _err -> do
-      Logging.logErrorText $ "Failed to parse query: " <> query
-      pure $ DefinitionSearchResults []
-    Right (searchTokens, mayArity) -> do
-      Logging.logInfoText $ "definition-search-tokens: " <> DSQ.searchTokensToTsQuery searchTokens
-      matches <-
-        PG.runTransactionMode PG.ReadCommitted PG.Read $
-          DDQ.definitionSearch callerUserId filter limit searchTokens mayArity
-            >>= PQ.expectProjectShortHandsOf (traversed . _1)
-            >>= RQ.expectReleaseVersionsOf (traversed . _2)
-            <&> over (traversed . _2) IDs.ReleaseShortHand
-      let results =
-            matches <&> \(project, release, fqn, summary) ->
+  matches <- case Text.words query of
+    [] -> pure $ []
+    [":"] -> pure $ []
+    -- If the query is a single word, and it doesn't contain "->", we treat it as a name search
+    [name]
+      | not (Text.isInfixOf "->" name) ->
+          PG.runTransactionMode PG.ReadCommitted PG.Read $
+            DDQ.definitionNameSearch callerUserId filter limit (Query name)
+    _ -> do
+      case DefinitionSearch.queryToTokens query of
+        Left _err -> do
+          Logging.logErrorText $ "Failed to parse query: " <> query
+          pure $ []
+        Right (searchTokens, mayArity) -> do
+          Logging.logInfoText $ "definition-search-tokens: " <> DSQ.searchTokensToTsQuery searchTokens
+          PG.runTransactionMode PG.ReadCommitted PG.Read $
+            DDQ.definitionTokenSearch callerUserId filter limit searchTokens mayArity
+  results <-
+    PG.runTransactionMode PG.ReadCommitted PG.Read $
+      do
+        PQ.expectProjectShortHandsOf (traversed . _1) matches
+        >>= RQ.expectReleaseVersionsOf (traversed . _2)
+        <&> over (traversed . _2) IDs.ReleaseShortHand
+        <&> fmap
+          ( \(project, release, fqn, summary) ->
               DefinitionSearchResult
                 { fqn,
                   summary,
                   project,
                   release
                 }
-      pure $ DefinitionSearchResults results
+          )
+  pure $ DefinitionSearchResults results
   where
     limit = fromMaybe (Limit 20) mayLimit
 
diff --git a/transcripts/share-apis/search/defn-search-fuzzy.json b/transcripts/share-apis/search/defn-search-fuzzy.json
@@ -0,0 +1,119 @@
+{
+  "body": {
+    "results": [
+      {
+        "branchRef": "releases/1.2.3",
+        "definition": {
+          "displayName": "data.Text",
+          "hash": "##Text",
+          "summary": {
+            "contents": [
+              {
+                "annotation": null,
+                "segment": "data.Text"
+              }
+            ],
+            "tag": "BuiltinObject"
+          },
+          "tag": "Data"
+        },
+        "fqn": "data.Text",
+        "kind": "type",
+        "projectRef": "@transcripts/search"
+      },
+      {
+        "branchRef": "releases/1.2.3",
+        "definition": {
+          "displayName": "Nat.fromText",
+          "hash": "#he2h9hm6c2obcmnvp2g17dcg1l65h9t0q5mamjlrl2lk0flmjioc5j350l0386r88ho1u1ncqklh31qdqgrmbltrincivfcjcsmdeo0",
+          "summary": {
+            "contents": [
+              {
+                "annotation": {
+                  "contents": "##Text",
+                  "tag": "TypeReference"
+                },
+                "segment": "Text"
+              },
+              {
+                "annotation": null,
+                "segment": " "
+              },
+              {
+                "annotation": {
+                  "tag": "TypeOperator"
+                },
+                "segment": "->"
+              },
+              {
+                "annotation": null,
+                "segment": " "
+              },
+              {
+                "annotation": {
+                  "contents": "##Nat",
+                  "tag": "TypeReference"
+                },
+                "segment": "Nat"
+              }
+            ],
+            "tag": "UserObject"
+          },
+          "tag": "Plain"
+        },
+        "fqn": "Nat.fromText",
+        "kind": "term",
+        "projectRef": "@transcripts/search"
+      },
+      {
+        "branchRef": "releases/1.2.3",
+        "definition": {
+          "displayName": "Nat.toText",
+          "hash": "#9da94h7t7rhg55vdur11jle4b1r8tn8ptpqc4m4hg8tve9vb57g64av66ghc0dpoc4r5e4f7ippu9qja7nqcq1c73cvd5fuvl0o6efo",
+          "summary": {
+            "contents": [
+              {
+                "annotation": {
+                  "contents": "##Nat",
+                  "tag": "TypeReference"
+                },
+                "segment": "Nat"
+              },
+              {
+                "annotation": null,
+                "segment": " "
+              },
+              {
+                "annotation": {
+                  "tag": "TypeOperator"
+                },
+                "segment": "->"
+              },
+              {
+                "annotation": null,
+                "segment": " "
+              },
+              {
+                "annotation": {
+                  "contents": "##Text",
+                  "tag": "TypeReference"
+                },
+                "segment": "Text"
+              }
+            ],
+            "tag": "UserObject"
+          },
+          "tag": "Plain"
+        },
+        "fqn": "Nat.toText",
+        "kind": "term",
+        "projectRef": "@transcripts/search"
+      }
+    ]
+  },
+  "status": [
+    {
+      "status_code": 200
+    }
+  ]
+}
diff --git a/transcripts/share-apis/search/run.zsh b/transcripts/share-apis/search/run.zsh
@@ -55,6 +55,8 @@ fetch "$transcript_user" GET 'name-search-release-filter-sad' '/search-names?que
 # "b -> a -> a"
 fetch "$transcript_user" GET 'type-var-search' '/search-definitions?query=b%20-%3E%20a%20-%3E%20a'
 
+# Fuzzy search, should find 'Text' and 'Nat.toText' and 'Nat.fromText'
+fetch "$transcript_user" GET 'defn-search-fuzzy' '/search-definitions?query=Text'
 
 # (a -> b) -> List a -> List b
 fetch "$transcript_user" GET 'complex-type-mention-search' '/search-definitions?query=(a%20-%3E%20b)%20-%3E%20List%20a%20-%3E%20List%20b'
