Add SearchIndex and VectorSearchIndex

Co-authored-by: Tim Graham <timograham@gmail.com>

Author: WaVEV

.github/workflows/mongodb_settings.py

@@ -17,10 +17,13 @@
  "default": {
  "ENGINE": "django_mongodb_backend",
  "NAME": "djangotests",
+ # Required when connecting to the Atlas image in Docker.
+ "OPTIONS": {"directConnection": True},
  },
  "other": {
  "ENGINE": "django_mongodb_backend",
  "NAME": "djangotests-other",
+ "OPTIONS": {"directConnection": True},
  },
  }
 

.github/workflows/start_local_atlas.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+set -eu
+
+echo "Starting the container"
+
+IMAGE=${1:-mongodb/mongodb-atlas-local:latest}
+DOCKER=$(which docker || which podman)
+
+$DOCKER pull $IMAGE
+
+$DOCKER kill mongodb_atlas_local || true
+
+CONTAINER_ID=$($DOCKER run --rm -d --name mongodb_atlas_local -p 27017:27017 $IMAGE)
+
+function wait() {
+ CONTAINER_ID=$1
+ echo "waiting for container to become healthy..."
+ $DOCKER logs mongodb_atlas_local
+}
+
+wait "$CONTAINER_ID"
+
+# Sleep for a bit to let all services start.
+sleep 5

.github/workflows/test-python-atlas.yml

@@ -0,0 +1,56 @@
+name: Python Tests on Atlas
+
+on:
+ pull_request:
+ paths:
+ - '**.py'
+ - '!setup.py'
+ - '.github/workflows/test-python-atlas.yml'
+ workflow_dispatch:
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}
+ cancel-in-progress: true
+
+defaults:
+ run:
+ shell: bash -eux {0}
+
+jobs:
+ build:
+ name: Django Test Suite
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout django-mongodb-backend
+ uses: actions/checkout@v4
+ with:
+ persist-credentials: false
+ - name: install django-mongodb-backend
+ run: |
+ pip3 install --upgrade pip
+ pip3 install -e .
+ - name: Checkout Django
+ uses: actions/checkout@v4
+ with:
+ repository: 'mongodb-forks/django'
+ ref: 'mongodb-5.2.x'
+ path: 'django_repo'
+ persist-credentials: false
+ - name: Install system packages for Django's Python test dependencies
+ run: |
+ sudo apt-get update
+ sudo apt-get install libmemcached-dev
+ - name: Install Django and its Python test dependencies
+ run: |
+ cd django_repo/tests/
+ pip3 install -e ..
+ pip3 install -r requirements/py3.txt
+ - name: Copy the test settings file
+ run: cp .github/workflows/mongodb_settings.py django_repo/tests/
+ - name: Copy the test runner file
+ run: cp .github/workflows/runtests.py django_repo/tests/runtests_.py
+ - name: Start local Atlas
+ working-directory: .
+ run: bash .github/workflows/start_local_atlas.sh mongodb/mongodb-atlas-local:7
+ - name: Run tests
+ run: python3 django_repo/tests/runtests_.py

django_mongodb_backend/__init__.py

@@ -7,6 +7,7 @@
 check_django_compatability()
 
 from .aggregates import register_aggregates # noqa: E402
+from .checks import register_checks # noqa: E402
 from .expressions import register_expressions # noqa: E402
 from .fields import register_fields # noqa: E402
 from .functions import register_functions # noqa: E402
@@ -17,6 +18,7 @@
 __all__ = ["parse_uri"]
 
 register_aggregates()
+register_checks()
 register_expressions()
 register_fields()
 register_functions()

django_mongodb_backend/checks.py

@@ -0,0 +1,32 @@
+from itertools import chain
+
+from django.apps import apps
+from django.core.checks import Tags, register
+from django.db import connections, router
+
+
+def check_indexes(app_configs, databases=None, **kwargs): # noqa: ARG001
+ """
+ Call Index.check() on all model indexes.
+
+ This function will be obsolete when Django calls Index.check() after
+ https://code.djangoproject.com/ticket/36273.
+ """
+ errors = []
+ if app_configs is None:
+ models = apps.get_models()
+ else:
+ models = chain.from_iterable(app_config.get_models() for app_config in app_configs)
+ for model in models:
+ for db in databases or ():
+ if not router.allow_migrate_model(db, model):
+ continue
+ connection = connections[db]
+ for model_index in model._meta.indexes:
+ if hasattr(model_index, "check"):
+ errors.extend(model_index.check(model, connection))
+ return errors
+
+
+def register_checks():
+ register(check_indexes, Tags.models)

django_mongodb_backend/features.py

@@ -1,5 +1,6 @@
 from django.db.backends.base.features import BaseDatabaseFeatures
 from django.utils.functional import cached_property
+from pymongo.errors import OperationFailure
 
 
 class DatabaseFeatures(BaseDatabaseFeatures):
@@ -548,3 +549,19 @@ def django_test_expected_failures(self):
  @cached_property
  def is_mongodb_6_3(self):
  return self.connection.get_database_version() >= (6, 3)
+
+ @cached_property
+ def supports_atlas_search(self):
+ """Does the server support Atlas search queries and search indexes?"""
+ try:
+ # An existing collection must be used on MongoDB 6, otherwise
+ # the operation will not error when unsupported.
+ self.connection.get_collection("django_migrations").list_search_indexes()
+ except OperationFailure:
+ # It would be best to check the error message or error code to
+ # avoid hiding some other exception, but the message/code varies
+ # across MongoDB versions. Example error message:
+ # "$listSearchIndexes stage is only allowed on MongoDB Atlas".
+ return False
+ else:
+ return True

django_mongodb_backend/fields/embedded_model.py

@@ -20,6 +20,9 @@ def __init__(self, embedded_model, *args, **kwargs):
  self.embedded_model = embedded_model
  super().__init__(*args, **kwargs)
 
+ def db_type(self, connection):
+ return "embeddedDocuments"
+
  def check(self, **kwargs):
  from ..models import EmbeddedModel
 

django_mongodb_backend/indexes.py

@@ -1,12 +1,16 @@
+import itertools
 from collections import defaultdict
 
+from django.core.checks import Error, Warning
 from django.db import NotSupportedError
-from django.db.models import Index
+from django.db.models import FloatField, Index, IntegerField
 from django.db.models.lookups import BuiltinLookup
 from django.db.models.sql.query import Query
 from django.db.models.sql.where import AND, XOR, WhereNode
 from pymongo import ASCENDING, DESCENDING
-from pymongo.operations import IndexModel
+from pymongo.operations import IndexModel, SearchIndexModel
+
+from django_mongodb_backend.fields import ArrayField
 
 from .query_utils import process_rhs
 
@@ -101,6 +105,181 @@ def where_node_idx(self, compiler, connection):
  return mql
 
 
+class SearchIndex(Index):
+ suffix = "six"
+ _error_id_prefix = "django_mongodb_backend.indexes.SearchIndex"
+
+ def __init__(self, *, fields=(), name=None):
+ super().__init__(fields=fields, name=name)
+
+ def check(self, model, connection):
+ errors = []
+ if not connection.features.supports_atlas_search:
+ errors.append(
+ Warning(
+ f"This MongoDB server does not support {self.__class__.__name__}.",
+ hint=(
+ "The index won't be created. Use an Atlas-enabled version of MongoDB, "
+ "or silence this warning if you don't care about it."
+ ),
+ obj=model,
+ id=f"{self._error_id_prefix}.W001",
+ )
+ )
+ return errors
+
+ def search_index_data_types(self, db_type):
+ """
+ Map a model field's type to search index type.
+ https://www.mongodb.com/docs/atlas/atlas-search/define-field-mappings/#data-types
+ """
+ if db_type in {"double", "int", "long"}:
+ return "number"
+ if db_type == "binData":
+ return "string"
+ if db_type == "bool":
+ return "boolean"
+ if db_type == "object":
+ return "document"
+ if db_type == "array":
+ return "embeddedDocuments"
+ return db_type
+
+ def get_pymongo_index_model(
+ self, model, schema_editor, field=None, unique=False, column_prefix=""
+ ):
+ if not schema_editor.connection.features.supports_atlas_search:
+ return None
+ fields = {}
+ for field_name, _ in self.fields_orders:
+ field = model._meta.get_field(field_name)
+ type_ = self.search_index_data_types(field.db_type(schema_editor.connection))
+ field_path = column_prefix + model._meta.get_field(field_name).column
+ fields[field_path] = {"type": type_}
+ return SearchIndexModel(
+ definition={"mappings": {"dynamic": False, "fields": fields}}, name=self.name
+ )
+
+
+class VectorSearchIndex(SearchIndex):
+ suffix = "vsi"
+ _error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
+ VALID_FIELD_TYPES = frozenset(("boolean", "date", "number", "objectId", "string", "uuid"))
+ VALID_SIMILARITIES = frozenset(("cosine", "dotProduct", "euclidean"))
+
+ def __init__(self, *, fields=(), name=None, similarities):
+ super().__init__(fields=fields, name=name)
+ self.similarities = similarities
+ self._multiple_similarities = isinstance(similarities, tuple | list)
+ for func in similarities if self._multiple_similarities else (similarities,):
+ if func not in self.VALID_SIMILARITIES:
+ raise ValueError(
+ f"'{func}' isn't a valid similarity function "
+ f"({', '.join(sorted(self.VALID_SIMILARITIES))})."
+ )
+ seen_fields = set()
+ for field_name, _ in self.fields_orders:
+ if field_name in seen_fields:
+ raise ValueError(f"Field '{field_name}' is duplicated in fields.")
+ seen_fields.add(field_name)
+
+ def check(self, model, connection):
+ errors = super().check(model, connection)
+ num_arrayfields = 0
+ for field_name, _ in self.fields_orders:
+ field = model._meta.get_field(field_name)
+ if isinstance(field, ArrayField):
+ num_arrayfields += 1
+ try:
+ int(field.size)
+ except (ValueError, TypeError):
+ errors.append(
+ Error(
+ f"VectorSearchIndex requires 'size' on field '{field_name}'.",
+ obj=model,
+ id=f"{self._error_id_prefix}.E002",
+ )
+ )
+ if not isinstance(field.base_field, FloatField | IntegerField):
+ errors.append(
+ Error(
+ "VectorSearchIndex requires the base field of "
+ f"ArrayField '{field.name}' to be FloatField or "
+ "IntegerField but is "
+ f"{field.base_field.get_internal_type()}.",
+ obj=model,
+ id=f"{self._error_id_prefix}.E003",
+ )
+ )
+ else:
+ search_type = self.search_index_data_types(field.db_type(connection))
+ if search_type not in self.VALID_FIELD_TYPES:
+ errors.append(
+ Error(
+ "VectorSearchIndex does not support field "
+ f"'{field_name}' ({field.get_internal_type()}).",
+ obj=model,
+ id=f"{self._error_id_prefix}.E004",
+ hint=f"Allowed types are {', '.join(sorted(self.VALID_FIELD_TYPES))}.",
+ )
+ )
+ if self._multiple_similarities and num_arrayfields != len(self.similarities):
+ errors.append(
+ Error(
+ f"VectorSearchIndex requires the same number of similarities "
+ f"and vector fields; {model._meta.object_name} has "
+ f"{num_arrayfields} ArrayField(s) but similarities "
+ f"has {len(self.similarities)} element(s).",
+ obj=model,
+ id=f"{self._error_id_prefix}.E005",
+ )
+ )
+ if num_arrayfields == 0:
+ errors.append(
+ Error(
+ "VectorSearchIndex requires at least one ArrayField to " "store vector data.",
+ obj=model,
+ id=f"{self._error_id_prefix}.E006",
+ hint="If you want to perform search operations without vectors, "
+ "use SearchIndex instead.",
+ )
+ )
+ return errors
+
+ def deconstruct(self):
+ path, args, kwargs = super().deconstruct()
+ kwargs["similarities"] = self.similarities
+ return path, args, kwargs
+
+ def get_pymongo_index_model(
+ self, model, schema_editor, field=None, unique=False, column_prefix=""
+ ):
+ if not schema_editor.connection.features.supports_atlas_search:
+ return None
+ similarities = (
+ itertools.cycle([self.similarities])
+ if not self._multiple_similarities
+ else iter(self.similarities)
+ )
+ fields = []
+ for field_name, _ in self.fields_orders:
+ field_ = model._meta.get_field(field_name)
+ field_path = column_prefix + model._meta.get_field(field_name).column
+ mappings = {"path": field_path}
+ if isinstance(field_, ArrayField):
+ mappings.update(
+ {
+ "type": "vector",
+ "numDimensions": int(field_.size),
+ "similarity": next(similarities),
+ }
+ )
+ else:
+ mappings["type"] = "filter"
+ fields.append(mappings)
+ return SearchIndexModel(definition={"fields": fields}, name=self.name, type="vectorSearch")
+
+
 def register_indexes():
  BuiltinLookup.as_mql_idx = builtin_lookup_idx
  Index._get_condition_mql = _get_condition_mql