feat: Add 'using' and 'database' arguments to management command

Author: qcoumes

.github/workflows/tests_and_publish.yml

@@ -124,9 +124,9 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ 3.9, '3.10', '3.11', '3.12', '3.13' ]
-        django-version: [ 42, 50, 51, 52 ]
-        opensearch-version: [ 10, 20, 30 ]
+        python-version: [ '3.9', '3.10', '3.11', '3.12', '3.13' ]
+        django-version: [ '42', '50', '51', '52' ]
+        opensearch-version: [ '10', '20', '30' ]
         exclude:
           - python-version: 3.9
             django-version: 50
@@ -149,7 +149,7 @@ jobs:
 
       - name: Run Opensearch in docker
         run: |
-          docker compose up -d opensearch_test_${{ matrix.opensearch-version }}
+          docker compose up -d opensearch_test_${{ matrix.opensearch-version }}_0 opensearch_test_${{ matrix.opensearch-version }}_1
           sleep 30
 
       - name: Install Tox and any other packages

django_opensearch_dsl/documents.py

@@ -67,9 +67,11 @@ def search(cls, using: str = None, index: str = None) -> opensearchpy.Search:
             model=cls.django.model,
         )
 
-    def get_queryset(self, filter_: Optional[Q] = None, exclude: Optional[Q] = None, count: int = None) -> QuerySet:
+    def get_queryset(
+        self, filter_: Optional[Q] = None, exclude: Optional[Q] = None, count: int = None, alias: str = None
+    ) -> QuerySet:
         """Return the queryset that should be indexed by this doc type."""
-        qs = self.django.model.objects.all()
+        qs = self.django.model.objects.using(alias).all()
 
         if filter_:
             qs = qs.filter(filter_)
@@ -95,13 +97,14 @@ def get_indexing_queryset(
         verbose: bool = False,
         filter_: Optional[Q] = None,
         exclude: Optional[Q] = None,
+        alias: str = None,
         count: int = None,
         action: CommandAction = CommandAction.INDEX,
         stdout: TextIO = sys.stdout,
     ) -> Iterable:
         """Divide the queryset into chunks."""
         chunk_size = self.django.queryset_pagination
-        qs = self.get_queryset(filter_=filter_, exclude=exclude, count=count)
+        qs = self.get_queryset(filter_=filter_, exclude=exclude, count=count, alias=alias)
         qs = qs.order_by("pk") if not qs.query.is_sliced else qs
         total = qs.count()
         model = self.django.model.__name__

django_opensearch_dsl/management/commands/opensearch.py

@@ -8,15 +8,27 @@
 
 import opensearchpy
 from django.core.exceptions import FieldError
-from django.core.management import BaseCommand
+from django.core.management import BaseCommand, CommandError
 from django.db.models import Q
+from opensearchpy import OpenSearch
+from opensearchpy.connection.connections import connections
 
 from ...apps import DODConfig
 from ...enums import CommandAction
 from ...registries import registry
 from ..types import Values, parse
 
 
+def connection(using: str = "default") -> opensearchpy.OpenSearch:
+    """Return the OpenSearch connection for the given alias."""
+    try:
+        return connections.get_connection(using)
+    except KeyError:
+        raise CommandError(
+            f"No OpenSearch connection found for alias '{using}', known connections are: {list(connections._kwargs.keys())}"
+        )
+
+
 class Command(BaseCommand):
     """Manage indices and documents."""
 
@@ -46,20 +58,20 @@ def wrap(value: str) -> tuple[str, Values]:
                     f"manage.py index: error: invalid filter: '{value}' (filter must be formatted as "
                     f"'[Field Lookups]=[value]')\n",
                 )
-                exit(1)
+                raise CommandError
             return lookup, v  # noqa
 
         return wrap
 
-    def __list_index(self, **options: Any) -> None:  # noqa pragma: no cover
+    def __list_index(self, using: OpenSearch, **options: Any) -> None:  # noqa pragma: no cover
         """List all known index and indicate whether they are created or not."""
         indices = registry.get_indices()
         result = defaultdict(list)
         for index in indices:
             module = index._doc_types[0].__module__.split(".")[-2]  # noqa
-            exists = index.exists()
+            exists = index.exists(using=using)
             checkbox = f"[{'X' if exists else ' '}]"
-            count = f" ({index.search().count()} documents)" if exists else ""
+            count = f" ({index.search(using=using).count()} documents)" if exists else ""
             result[module].append(f"{checkbox} {index._name}{count}")
         for app, indice_names in result.items():
             self.stdout.write(self.style.MIGRATE_LABEL(app))
@@ -72,6 +84,7 @@ def _manage_index(
         force: bool,
         verbosity: int,
         ignore_error: bool,
+        using: OpenSearch,
         **options: Any,
     ) -> None:
         """Manage the creation and deletion of indices."""
@@ -85,7 +98,7 @@ def _manage_index(
             unknown = set(indices) - set(known_name)
             if unknown:
                 self.stderr.write(f"Unknown indices '{list(unknown)}', choices are: '{known_name}'")
-                exit(1)
+                raise CommandError
 
             # Only keep given indices
             given_indices = [i for i in known if i._name in indices]
@@ -107,7 +120,7 @@ def _manage_index(
                     self.stdout.write("")
                     break
                 elif p.lower() in ["no", "n"]:
-                    exit(1)
+                    raise CommandError
 
         pp = action.present_participle.title()
         for index in given_indices:
@@ -119,24 +132,24 @@ def _manage_index(
                 self.stdout.flush()
             try:
                 if action == CommandAction.CREATE:
-                    index.create()
+                    index.create(using=using)
                 elif action == CommandAction.DELETE:
-                    index.delete()
+                    index.delete(using=using)
                 elif action == CommandAction.UPDATE:
-                    index.put_mapping(body=index.to_dict()["mappings"])
+                    index.put_mapping(using=using, body=index.to_dict()["mappings"])
                 else:
                     try:
-                        index.delete()
+                        index.delete(using=using)
                     except opensearchpy.exceptions.NotFoundError:
                         pass
-                    index.create()
+                    index.create(using=using)
             except opensearchpy.exceptions.TransportError as e:
                 if verbosity or not ignore_error:
                     error = self.style.ERROR(f"Error: {e.error} - {e.info}")
                     self.stderr.write(f"{pp} index '{index._name}'...\n{error}")  # noqa
                 if not ignore_error:
                     self.stderr.write("exiting...")
-                    exit(1)
+                    raise CommandError
             else:
                 if verbosity:
                     self.stdout.write(f"{pp} index '{index._name}'... {self.style.SUCCESS('OK')}")  # noqa
@@ -153,6 +166,8 @@ def _manage_document(
         count: bool,
         refresh: bool,
         missing: bool,
+        using: OpenSearch,
+        database: str,
         **options: Any,
     ) -> None:
         """Manage the creation and deletion of indices."""
@@ -168,19 +183,19 @@ def _manage_document(
             unknown = set(indices) - set(known_name)
             if unknown:
                 self.stderr.write(f"Unknown indices '{list(unknown)}', choices are: '{known_name}'")
-                exit(1)
+                raise CommandError
 
             # Only keep given indices
             given_indices = list(filter(lambda i: i._name in indices, known))  # type: ignore[arg-type]
         else:
             given_indices = list(known)
 
         # Ensure every indices needed are created
-        not_created = [i._name for i in given_indices if not i.exists()]  # noqa
+        not_created = [i._name for i in given_indices if not i.exists(using=using)]  # noqa
         if not_created:
             self.stderr.write(f"The following indices are not created : {not_created}")
             self.stderr.write("Use 'python3 manage.py opensearch list' to list indices' state.")
-            exit(1)
+            raise CommandError
 
         # Check field, preparing to display expected actions
         s = f"The following documents will be {action.past}:"
@@ -189,17 +204,17 @@ def _manage_document(
             # Handle --missing
             exclude_ = exclude
             if missing and action == CommandAction.INDEX:
-                q = Q(pk__in=[h.meta.id for h in index.search().extra(stored_fields=[]).scan()])
+                q = Q(pk__in=[h.meta.id for h in index.search(using=using).extra(stored_fields=[]).scan()])
                 exclude_ = exclude_ & q if exclude_ is not None else q
 
             document = index._doc_types[0]()  # noqa
             try:
                 kwargs_list.append({"filter_": filter_, "exclude": exclude_, "count": count})
-                qs = document.get_queryset(filter_=filter_, exclude=exclude_, count=count).count()
+                qs = document.get_queryset(filter_=filter_, exclude=exclude_, count=count, alias=database).count()
             except FieldError as e:
                 model = index._doc_types[0].django.model.__name__  # noqa
                 self.stderr.write(f"Error while filtering on '{model}' (from index '{index._name}'):\n{e}'")  # noqa
-                exit(1)
+                raise CommandError
             else:
                 s += f"\n\t- {qs} {document.django.model.__name__}."
 
@@ -215,14 +230,16 @@ def _manage_document(
                     self.stdout.write("\n")
                     break
                 elif p.lower() in ["no", "n"]:
-                    exit(1)
+                    raise CommandError
 
         result = "\n"
         for index, kwargs in zip(given_indices, kwargs_list):
             document = index._doc_types[0]()  # noqa
-            qs = document.get_indexing_queryset(stdout=self.stdout._out, verbose=verbosity, action=action, **kwargs)
+            qs = document.get_indexing_queryset(
+                stdout=self.stdout._out, verbose=verbosity, action=action, alias=database, **kwargs
+            )
             success, errors = document.update(
-                qs, parallel=parallel, refresh=refresh, action=action, raise_on_error=False
+                qs, parallel=parallel, refresh=refresh, action=action, raise_on_error=False, using=using
             )
 
             success_str = self.style.SUCCESS(success) if success else success
@@ -256,6 +273,12 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
             description="Show all available indices (and their state) for the current project.",
         )
         subparser.set_defaults(func=self.__list_index)
+        subparser.add_argument(
+            "--using",
+            type=connection,
+            default=connection("default"),
+            help="Alias of the OpenSearch connection to use. Default to 'default'.",
+        )
 
         # 'manage' subcommand
         subparser = subparsers.add_parser(
@@ -264,6 +287,12 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
             description="Manage the creation an deletion of indices.",
         )
         subparser.set_defaults(func=self._manage_index)
+        subparser.add_argument(
+            "--using",
+            type=connection,
+            default=connection("default"),
+            help="Alias of the OpenSearch connection to use. Default to 'default'.",
+        )
         subparser.add_argument(
             "action",
             type=str,
@@ -308,6 +337,13 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
                 CommandAction.UPDATE.value,
             ],
         )
+        subparser.add_argument(
+            "--using",
+            type=connection,
+            default=connection("default"),
+            help="Alias of the OpenSearch connection to use. Default to 'default'.",
+        )
+        subparser.add_argument("-d", "--database", default=None, dest="database", help="Nominates a database.")
         subparser.add_argument(
             "-f",
             "--filters",
@@ -383,6 +419,6 @@ def handle(self, *args: Any, **options: Any) -> None:
         if "func" not in options:  # pragma: no cover
             self.stderr.write(self.usage)
             self.stderr.write(f"manage.py opensearch: error: no subcommand provided.")
-            exit(1)
+            raise CommandError
 
         options["func"](**options)

docker-compose.yml

@@ -1,7 +1,7 @@
 services:
 
-  opensearch_test_10:
-    container_name: opensearch_test_10
+  opensearch_test_10_0:
+    container_name: opensearch_test_10_0
     image: opensearchproject/opensearch:1
     ports:
       - "9210:9200"
@@ -11,8 +11,19 @@ services:
       - discovery.type=single-node
       - "ES_JAVA_OPTS=-Xms128m -Xmx128m"
 
-  opensearch_test_20:
-    container_name: opensearch_test_20
+  opensearch_test_10_1:
+    container_name: opensearch_test_10_1
+    image: opensearchproject/opensearch:1
+    ports:
+      - "9211:9200"
+      - "9611:9600"
+    environment:
+      - plugins.security.disabled=true
+      - discovery.type=single-node
+      - "ES_JAVA_OPTS=-Xms128m -Xmx128m"
+
+  opensearch_test_20_0:
+    container_name: opensearch_test_20_0
     image: opensearchproject/opensearch:2
     ports:
       - "9220:9200"
@@ -23,8 +34,20 @@ services:
       - "ES_JAVA_OPTS=-Xms128m -Xmx128m"
       - OPENSEARCH_INITIAL_ADMIN_PASSWORD="XJ67NCmLj4yMPPz0wthVUvVGV0cQiq"
 
-  opensearch_test_30:
-    container_name: opensearch_test_30
+  opensearch_test_20_1:
+    container_name: opensearch_test_20_1
+    image: opensearchproject/opensearch:2
+    ports:
+      - "9221:9200"
+      - "9621:9600"
+    environment:
+      - plugins.security.disabled=true
+      - discovery.type=single-node
+      - "ES_JAVA_OPTS=-Xms128m -Xmx128m"
+      - OPENSEARCH_INITIAL_ADMIN_PASSWORD="XJ67NCmLj4yMPPz0wthVUvVGV0cQiq"
+
+  opensearch_test_30_0:
+    container_name: opensearch_test_30_0
     image: opensearchproject/opensearch:3
     ports:
       - "9230:9200"
@@ -34,3 +57,15 @@ services:
       - discovery.type=single-node
       - "ES_JAVA_OPTS=-Xms128m -Xmx128m"
       - OPENSEARCH_INITIAL_ADMIN_PASSWORD="XJ67NCmLj4yMPPz0wthVUvVGV0cQiq"
+
+  opensearch_test_30_1:
+    container_name: opensearch_test_30_1
+    image: opensearchproject/opensearch:3
+    ports:
+      - "9231:9200"
+      - "9631:9600"
+    environment:
+      - plugins.security.disabled=true
+      - discovery.type=single-node
+      - "ES_JAVA_OPTS=-Xms128m -Xmx128m"
+      - OPENSEARCH_INITIAL_ADMIN_PASSWORD="XJ67NCmLj4yMPPz0wthVUvVGV0cQiq"

docs/document.md

@@ -106,11 +106,12 @@ queryset used to fetch the data.
 
 ---
 
-* `def get_queryset(self, filter_=None, exclude=None, count=None)`
+* `def get_queryset(self, filter_=None, exclude=None, count=None, alias=None)`
 
     * `filter_` (`Optional[Q]`) - `Q` object given to the queryset's `filter()` method.
     * `exclude` (`Optional[Q]`) - `Q` object given to the queryset's `exclude()` method.
     * `count` (`Optional[int]`) - Limit the queryset with the given number.
+    * `alias` (`Optional[str]`) - Alias of the database to use.
 
 By default, this method retrieves all objects from the model associated with the `Document`, optionally filtering and
 excluding elements according to the given arguments. You can also limit the number of results using the `count` argument.
@@ -131,7 +132,7 @@ class EventDocument(Document):
 
     country = fields.ObjectField(doc_class=CountryDocument)
 
-    def get_queryset(self, filter_: Optional[Q] = None, exclude: Optional[Q] = None, count: int = 0) -> QuerySet:
+    def get_queryset(self, filter_: Optional[Q]=None, exclude: Optional[Q]=None, count: int=0, alias: str=None) -> QuerySet:
         """Select country to improve indexing performance."""
         return super().get_queryset(filter_=filter_, exclude=exclude, count=count).select_related('country')
 ```

docs/fields.md

@@ -68,7 +68,6 @@ directly. Change the `CarDocument` to look like this:
 ```python
 # documents.py
 
-@registry.register_document
 class CarDocument(Document):
     # add a string field to the Opensearch mapping called type, the
     # value of which is derived from the model's type_to_string attribute
@@ -192,7 +191,7 @@ class CarDocument(Document):
       'color',
     ]
 
-  def get_queryset(self, filters: Optional[Dict[str, Any]] = None, count: int = 0) -> 'QuerySet':
+  def get_queryset(self, filters: Optional[Dict[str, Any]]=None, count: int=0, alias: str=None) -> 'QuerySet':
     """Not mandatory but to improve performance we can select related in one sql request"""
     return super().get_queryset(count=count).select_related(
       'manufacturer'