From 48975a43e170f8845d0bd8353598fd79cd6129f7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Tamarelle?=
 <jerome.tamarelle@mongodb.com>
Date: Wed, 18 Dec 2024 13:28:14 +0100
Subject: [PATCH] Add schema helpers to create search and vector indexes

---
 src/Schema/Blueprint.php | 36 ++++++++++++++++++++++++++
 tests/SchemaTest.php     | 56 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/src/Schema/Blueprint.php b/src/Schema/Blueprint.php
index f107bd7e5..fcd89df81 100644
--- a/src/Schema/Blueprint.php
+++ b/src/Schema/Blueprint.php
@@ -16,6 +16,14 @@
 use function is_string;
 use function key;
 
+/**
+ * @phpstan-type SearchIndexField array{type: 'boolean'|'date'|'dateFacet'|'objectId'|'stringFacet'|'uuid'} | array{type: 'autocomplete', analyzer?: string, maxGrams?: int, minGrams?: int, tokenization?: 'edgeGram'|'rightEdgeGram'|'nGram', foldDiacritics?: bool} | array{type: 'document'|'embeddedDocuments', dynamic?:bool, fields: array<string, array>} | array{type: 'geo', indexShapes?: bool} | array{type: 'number'|'numberFacet', representation?: 'int64'|'double', indexIntegers?: bool, indexDoubles?: bool} | array{type: 'token', normalizer?: 'lowercase'|'none'} | array{type: 'string', analyzer?: string, searchAnalyzer?: string, indexOptions?: 'docs'|'freqs'|'positions'|'offsets', store?: bool, ignoreAbove?: int, multi?: array<string, array>, norms?: 'include'|'omit'}
+ * @phpstan-type SearchIndexAnalyser array{name: string, charFilters?: list<array<string, mixed>>, tokenizer: array{type: string}, tokenFilters?: list<array<string, mixed>>}
+ * @phpstan-type SearchIndexStoredSource bool | array{includes: array<string>} | array{excludes: array<string>}
+ * @phpstan-type SearchIndexDefinition array{analyser?: string, analyzers?: SearchIndexAnalyser[], searchAnalyzer?: string, mappings: array{dynamic: true} | array{dynamic?: bool, fields: array<string, SearchIndexField>}, storedSource?: SearchIndexStoredSource}
+ * @phpstan-type VectorSearchIndexField array{type: 'vector', path: string, numDimensions: int, similarity: 'euclidean'|'cosine'|'dotProduct', quantization?: 'none'|'scalar'|'binary'}
+ * @phpstan-type VectorSearchIndexDefinition array{fields: array<string, VectorSearchIndexField>}
+ */
 class Blueprint extends SchemaBlueprint
 {
     /**
@@ -303,6 +311,34 @@ public function sparse_and_unique($columns = null, $options = [])
         return $this;
     }
 
+    /**
+     * Create an Atlas Search Index.
+     *
+     * @see https://www.mongodb.com/docs/atlas/atlas-search/
+     *
+     * @phpstan-param SearchIndexDefinition $definition
+     */
+    public function searchIndex(array $definition, string $name = 'default'): static
+    {
+        $this->collection->createSearchIndex($definition, ['name' => $name, 'type' => 'search']);
+
+        return $this;
+    }
+
+    /**
+     * Create an Atlas Vector Search Index.
+     *
+     * @see https://www.mongodb.com/docs/atlas/atlas-vector-search/
+     *
+     * @phpstan-param VectorSearchIndexDefinition $definition
+     */
+    public function vectorSearchIndex(array $definition, string $name = 'default'): static
+    {
+        $this->collection->createSearchIndex($definition, ['name' => $name, 'type' => 'vectorSearch']);
+
+        return $this;
+    }
+
     /**
      * Allow fluent columns.
      *
diff --git a/tests/SchemaTest.php b/tests/SchemaTest.php
index ff3dfe626..4b407a150 100644
--- a/tests/SchemaTest.php
+++ b/tests/SchemaTest.php
@@ -8,8 +8,10 @@
 use Illuminate\Support\Facades\Schema;
 use MongoDB\BSON\Binary;
 use MongoDB\BSON\UTCDateTime;
+use MongoDB\Collection;
 use MongoDB\Laravel\Schema\Blueprint;
 
+use function assert;
 use function collect;
 use function count;
 
@@ -502,9 +504,51 @@ public function testGetIndexes()
         $this->assertSame([], $indexes);
     }
 
+    /** @todo requires SearchIndex support */
+    public function testSearchIndex(): void
+    {
+        Schema::create('newcollection', function (Blueprint $collection) {
+            $collection->searchIndex([
+                'mappings' => [
+                    'dynamic' => false,
+                    'fields' => [
+                        'foo' => ['type' => 'string', 'analyzer' => 'lucene.whitespace'],
+                    ],
+                ],
+            ]);
+        });
+
+        $index = $this->getSearchIndex('newcollection', 'default');
+        self::assertNotFalse($index);
+
+        self::assertSame('default', $index['name']);
+        self::assertSame('search', $index['type']);
+        self::assertFalse($index['latestDefinition']['mappings']['dynamic']);
+        self::assertSame('lucene.whitespace', $index['latestDefinition']['mappings']['fields']['foo']['analyzer']);
+    }
+
+    public function testVectorSearchIndex()
+    {
+        Schema::create('newcollection', function (Blueprint $collection) {
+            $collection->vectorSearchIndex([
+                'fields' => [
+                    ['type' => 'vector', 'path' => 'foo', 'numDimensions' => 128, 'similarity' => 'euclidean', 'quantization' => 'none'],
+                ],
+            ], 'vector');
+        });
+
+        $index = $this->getSearchIndex('newcollection', 'vector');
+        self::assertNotFalse($index);
+
+        self::assertSame('vector', $index['name']);
+        self::assertSame('vectorSearch', $index['type']);
+        self::assertSame('vector', $index['latestDefinition']['fields'][0]['type']);
+    }
+
     protected function getIndex(string $collection, string $name)
     {
         $collection = DB::getCollection($collection);
+        assert($collection instanceof Collection);
 
         foreach ($collection->listIndexes() as $index) {
             if (isset($index['key'][$name])) {
@@ -514,4 +558,16 @@ protected function getIndex(string $collection, string $name)
 
         return false;
     }
+
+    protected function getSearchIndex(string $collection, string $name)
+    {
+        $collection = DB::getCollection($collection);
+        assert($collection instanceof Collection);
+
+        foreach ($collection->listSearchIndexes(['name' => $name]) as $index) {
+            return $index;
+        }
+
+        return false;
+    }
 }