support composable templates

miguelgrinberg · Nov 11, 2024 · baaf601 · baaf601
1 parent 5511abe
commit baaf601
Show file tree

Hide file tree

Showing 6 changed files with 110 additions and 4 deletions.
diff --git a/elasticsearch_dsl/_async/index.py b/elasticsearch_dsl/_async/index.py
@@ -73,6 +73,47 @@ async def save(
         )
 
 
+class AsyncNewIndexTemplate:
+    def __init__(
+        self,
+        name: str,
+        template: str,
+        index: Optional["AsyncIndex"] = None,
+        priority: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        if index is None:
+            self._index = AsyncIndex(template, **kwargs)
+        else:
+            if kwargs:
+                raise ValueError(
+                    "You cannot specify options for Index when"
+                    " passing an Index instance."
+                )
+            self._index = index.clone()
+            self._index._name = template
+        self._template_name = name
+        self.priority = priority
+
+    def __getattr__(self, attr_name: str) -> Any:
+        return getattr(self._index, attr_name)
+
+    def to_dict(self) -> Dict[str, Any]:
+        d: Dict[str, Any] = {"template": self._index.to_dict()}
+        d["index_patterns"] = [self._index._name]
+        if self.priority is not None:
+            d["priority"] = self.priority
+        return d
+
+    async def save(
+        self, using: Optional[AsyncUsingType] = None
+    ) -> "ObjectApiResponse[Any]":
+        es = get_connection(using or self._index._using)
+        return await es.indices.put_index_template(
+            name=self._template_name, **self.to_dict()
+        )
+
+
 class AsyncIndex(IndexBase):
     _using: AsyncUsingType
 
@@ -109,6 +150,19 @@ def as_template(
             template_name, pattern or self._name, index=self, order=order
         )
 
+    def as_new_template(
+        self,
+        template_name: str,
+        pattern: Optional[str] = None,
+        priority: Optional[int] = None,
+    ) -> AsyncNewIndexTemplate:
+        # TODO: should we allow pattern to be a top-level arg?
+        # or maybe have an IndexPattern that allows for it and have
+        # Document._index be that?
+        return AsyncNewIndexTemplate(
+            template_name, pattern or self._name, index=self, priority=priority
+        )
+
     async def load_mappings(self, using: Optional[AsyncUsingType] = None) -> None:
         await self.get_or_create_mapping().update_from_es(
             self._name, using=using or self._using

diff --git a/elasticsearch_dsl/_sync/index.py b/elasticsearch_dsl/_sync/index.py
@@ -69,6 +69,43 @@ def save(self, using: Optional[UsingType] = None) -> "ObjectApiResponse[Any]":
         return es.indices.put_template(name=self._template_name, body=self.to_dict())
 
 
+class SyncNewIndexTemplate:
+    def __init__(
+        self,
+        name: str,
+        template: str,
+        index: Optional["Index"] = None,
+        priority: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        if index is None:
+            self._index = Index(template, **kwargs)
+        else:
+            if kwargs:
+                raise ValueError(
+                    "You cannot specify options for Index when"
+                    " passing an Index instance."
+                )
+            self._index = index.clone()
+            self._index._name = template
+        self._template_name = name
+        self.priority = priority
+
+    def __getattr__(self, attr_name: str) -> Any:
+        return getattr(self._index, attr_name)
+
+    def to_dict(self) -> Dict[str, Any]:
+        d: Dict[str, Any] = {"template": self._index.to_dict()}
+        d["index_patterns"] = [self._index._name]
+        if self.priority is not None:
+            d["priority"] = self.priority
+        return d
+
+    def save(self, using: Optional[UsingType] = None) -> "ObjectApiResponse[Any]":
+        es = get_connection(using or self._index._using)
+        return es.indices.put_index_template(name=self._template_name, **self.to_dict())
+
+
 class Index(IndexBase):
     _using: UsingType
 
@@ -103,6 +140,19 @@ def as_template(
             template_name, pattern or self._name, index=self, order=order
         )
 
+    def as_new_template(
+        self,
+        template_name: str,
+        pattern: Optional[str] = None,
+        priority: Optional[int] = None,
+    ) -> SyncNewIndexTemplate:
+        # TODO: should we allow pattern to be a top-level arg?
+        # or maybe have an IndexPattern that allows for it and have
+        # Document._index be that?
+        return SyncNewIndexTemplate(
+            template_name, pattern or self._name, index=self, priority=priority
+        )
+
     def load_mappings(self, using: Optional[UsingType] = None) -> None:
         self.get_or_create_mapping().update_from_es(
             self._name, using=using or self._using

diff --git a/examples/alias_migration.py b/examples/alias_migration.py
@@ -44,6 +44,7 @@
 
 ALIAS = "test-blog"
 PATTERN = ALIAS + "-*"
+PRIORITY = 100
 
 
 class BlogPost(Document):
@@ -81,7 +82,7 @@ def setup() -> None:
     deploy.
     """
     # create an index template
-    index_template = BlogPost._index.as_template(ALIAS, PATTERN)
+    index_template = BlogPost._index.as_new_template(ALIAS, PATTERN, priority=PRIORITY)
     # upload the template into elasticsearch
     # potentially overriding the one already there
     index_template.save()

diff --git a/examples/async/alias_migration.py b/examples/async/alias_migration.py
@@ -45,6 +45,7 @@
 
 ALIAS = "test-blog"
 PATTERN = ALIAS + "-*"
+PRIORITY = 100
 
 
 class BlogPost(AsyncDocument):
@@ -82,7 +83,7 @@ async def setup() -> None:
     deploy.
     """
     # create an index template
-    index_template = BlogPost._index.as_template(ALIAS, PATTERN)
+    index_template = BlogPost._index.as_new_template(ALIAS, PATTERN, priority=PRIORITY)
     # upload the template into elasticsearch
     # potentially overriding the one already there
     await index_template.save()

diff --git a/examples/async/parent_child.py b/examples/async/parent_child.py
@@ -226,7 +226,7 @@ async def save(self, **kwargs: Any) -> None:  # type: ignore[override]
 
 async def setup() -> None:
     """Create an IndexTemplate and save it into elasticsearch."""
-    index_template = Post._index.as_template("base")
+    index_template = Post._index.as_new_template("base", priority=100)
     await index_template.save()
 
 

diff --git a/examples/parent_child.py b/examples/parent_child.py
@@ -225,7 +225,7 @@ def save(self, **kwargs: Any) -> None:  # type: ignore[override]
 
 def setup() -> None:
     """Create an IndexTemplate and save it into elasticsearch."""
-    index_template = Post._index.as_template("base")
+    index_template = Post._index.as_new_template("base", priority=100)
     index_template.save()