Retake

docker-compose.uffizzi.yml

@@ -37,6 +37,16 @@ services:
       - POSTGRES_DB=grai
     ports:
       - 5432:5432
+    command:
+      [
+        "postgres",
+        "-c",
+        "wal_level=logical",
+        "-c",
+        "max_replication_slots=10",
+        "-c",
+        "max_wal_senders=5",
+      ]
 
   frontend:
     image: "${FRONTEND_IMAGE}"

grai-frontend/src/components/graph/drawer/GraphSearch.tsx

@@ -114,8 +114,6 @@ const GraphSearch: React.FC<GraphSearchProps> = ({
     }
   }, [onSearch, selected, tables, setTableCenter])
 
-  if (error) return <GraphError error={error} />
-
   return (
     <Box>
       <Box sx={{ p: 1 }}>
@@ -142,6 +140,7 @@ const GraphSearch: React.FC<GraphSearchProps> = ({
           }}
         />
       </Box>
+      {error && <GraphError error={error} />}
       {loading && <Loading />}
       <List
         disablePadding

grai-server/app/api/types.py

@@ -2,7 +2,7 @@
 import time
 from collections import defaultdict
 from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Tuple, Union
 
 import strawberry
 import strawberry_django
@@ -35,6 +35,7 @@
 from lineage.models import Node as NodeModel
 from lineage.models import Source as SourceModel
 from lineage.types import EdgeFilter, EdgeOrder, Filter, NodeFilter, NodeOrder
+from search.search import SearchClient
 from users.types import User, UserFilter
 from workspaces.models import Membership as MembershipModel
 from workspaces.models import Workspace as WorkspaceModel
@@ -877,8 +878,6 @@
             NodeModel.objects.filter(workspace=self).values_list("namespace", flat=True).distinct()
         )
 
-        print(namespaces)
-
         return DataWrapper(namespaces)
 
     # Graph
@@ -920,38 +919,14 @@
         self,
         search: Optional[str] = strawberry.UNSET,
     ) -> List[BaseTable]:
-        def get_tables(search: Optional[str]) -> List[Table]:
-            def get_words(word: str) -> List[str]:
-                from django.db import connection
-
-                with connection.cursor() as cursor:
-                    cursor.execute(
-                        "SELECT * FROM unique_lexeme WHERE levenshtein_less_equal(word, %s, 2) < 3",
-                        [word],
-                    )
-                    rows = cursor.fetchall()
-
-                return list([item[0] for item in rows])
+        def get_tables(
+            search: Optional[str],
+        ) -> List[Union[Table, BaseTable]]:
+            client = SearchClient()
 
-            result = []
+            tables = client.search(workspace=self, query=search)
 
-            for word in search.replace("_", " ").replace(".", " ").strip().split(" "):
-                result += get_words(word)
-
-            search_string = " ".join(result)
-
-            return list(
-                NodeModel.objects.raw(
-                    """
-                    SELECT *
-                    FROM lineage_node
-                    WHERE workspace_id=%s
-                    AND metadata->'grai'->>'node_type'='Table'
-                    AND ts_rank(search, websearch_to_tsquery('simple', replace(replace(%s, ' ', ' or '), '.', ' or '))) > 0
-                    ORDER BY ts_rank(search, websearch_to_tsquery('simple', replace(replace(%s, ' ', ' or '), '.', ' or '))) DESC""",
-                    [self.id, search_string, search_string],
-                )
-            )
+            return tables
 
         graph = GraphCache(workspace=self)
 
@@ -960,7 +935,13 @@
         if search:
             tables = await sync_to_async(get_tables)(search)
 
-            ids = [table.id for table in tables]
+            if len(tables) == 0:
+                return []
+
+            if isinstance(tables[0], BaseTable):
+                return tables
+
+            ids = [table["id"] for table in tables]
 
         return graph.get_tables(ids=ids)
 

grai-server/app/entrypoint.sh

@@ -21,6 +21,9 @@ fi
 if [ "$INSTANCE_TYPE" = "server" ]; then
   echo "Initializing database with default data..."
   bash /usr/src/app/initialize_db.sh
+
+  echo "Building retake index..."
+  python manage.py build_search
 fi
 
 exec "$@"

grai-server/app/pyproject.toml

@@ -64,6 +64,7 @@ django-postgres-extra = "^2.0.8"
 grandalf = "^0.8"
 drf-nested-routers = "^0.93.4"
 django-redis = "^5.3.0"
+retakesearch = "^0.1.32"
 
 [tool.poetry.group.dev.dependencies]
 isort = "^5.10.1"

grai-server/app/search/admin.py

@@ -0,0 +1,3 @@
+from django.contrib import admin
+
+# Register your models here.

grai-server/app/search/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class SearchConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "search"

grai-server/app/search/basic.py

@@ -0,0 +1,13 @@
+from typing import List, Optional
+
+from lineage.graph_cache import GraphCache
+from workspaces.models import Workspace
+
+from .search import SearchInterface
+
+
+class BasicSearch(SearchInterface):
+    def search(self, workspace: Workspace, query: Optional[str]) -> List:
+        graph = GraphCache(workspace=workspace)
+
+        return graph.get_tables(search=query)

grai-server/app/search/management/commands/build_search.py

@@ -0,0 +1,15 @@
+from django.core.management.base import BaseCommand
+
+from search.search import SearchClient
+
+
+class Command(BaseCommand):
+    help = "Build the search index"
+
+    def add_arguments(self, parser):
+        pass
+
+    def handle(self, *args, **options):
+        search = SearchClient()
+
+        search.build()

grai-server/app/search/models.py

@@ -0,0 +1,3 @@
+from django.db import models
+
+# Create your models here.

grai-server/app/search/retake.py

@@ -0,0 +1,74 @@
+from typing import List, Optional
+
+from decouple import config
+from retakesearch import Client, Database, Search, Table
+
+from workspaces.models import Workspace
+
+from .search import SearchInterface
+
+client = Client(
+    api_key=config("RETAKE_API_KEY"),
+    url=config("RETAKE_API_URL"),
+)
+
+
+class RetakeSearch(SearchInterface):
+    index_name = "nodes"
+    columns = [
+        "id",
+        "name",
+        "namespace",
+        "display_name",
+        "workspace_id",
+        "metadata->grai->node_type",
+    ]
+
+    def build(self):
+        database = Database(
+            dbname=config("DB_NAME", default="grai"),
+            host=config("DB_HOST", default="db"),
+            port=int(config("DB_PORT", default="5432")),
+            user=config("DB_USER", default="grai"),
+            password=config("DB_PASSWORD", default="grai"),
+        )
+
+        table = Table(
+            name="lineage_node",
+            columns=self.columns,
+            transform={
+                "mapping": {"workspace_id": "keyword"},
+                "rename": {
+                    "metadata_grai_node_type": "node_type",
+                },
+            },
+        )
+
+        try:
+            index = client.get_index(index_name=self.index_name)
+        except:
+            index = client.create_index(index_name=self.index_name)
+
+        if not index:
+            raise ValueError("Table failed to index due to an unexpected error")
+
+        index.vectorize(self.columns)
+        index.add_source(database=database, table=table)
+
+    def search(self, workspace: Workspace, query: Optional[str]) -> List:
+        index = client.get_index(self.index_name)
+
+        dsl = {"query": {"query_string": {"query": f"workspace_id.keyword:{str(workspace.id)}"}}}
+        bm25_search_query = (
+            Search()
+            .from_dict(dsl)
+            # .query("fuzzy", display_name={"value": query, "fuzziness": 10})
+            .filter("term", node_type="table")
+            .with_neural(query=query, fields=["display_name"])
+        )
+
+        result = index.search(bm25_search_query)
+
+        hits = result.get("hits", {"hits": []}).get("hits", [])
+
+        return [hit["_source"] for hit in hits]

grai-server/app/search/search.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from decouple import config
+
+from workspaces.models import Workspace
+
+
+class SearchInterface(ABC):
+    @abstractmethod
+    def search(self, workspace: Workspace, query: str) -> List:
+        pass  # pragma: no cover
+
+    def build(self) -> None:
+        pass
+
+
+class SearchClient(SearchInterface):
+    @property
+    def client(self):
+        if config("RETAKE_API_URL", None):
+            from .retake import RetakeSearch
+
+            return RetakeSearch()
+
+        from .basic import BasicSearch
+
+        return BasicSearch()
+
+    def search(self, workspace: Workspace, query: Optional[str]) -> List:
+        return self.client.search(workspace, query)
+
+    def build(self):
+        return self.client.build()

grai-server/app/search/tests.py

@@ -0,0 +1,3 @@
+from django.test import TestCase
+
+# Create your tests here.

grai-server/app/search/tests/test_search.py

@@ -0,0 +1,7 @@
+from search.search import SearchClient
+
+
+async def test_build():
+    client = SearchClient()
+
+    client.build()

grai-server/app/search/views.py

@@ -0,0 +1,3 @@
+from django.shortcuts import render
+
+# Create your views here.

grai-server/app/the_guide/settings/base.py

@@ -152,6 +152,7 @@ def inner(value: str | bool) -> bool:
     "notifications",
     "workspaces",
     "users",
+    "search",
     "telemetry",
 ]