From 56f9300d0d73787546b3dbcc4a589c98f3164993 Mon Sep 17 00:00:00 2001
From: PIG208 <39874143+PIG208@users.noreply.github.com>
Date: Sun, 24 Jul 2022 03:40:13 -0400
Subject: [PATCH] Improve stubs for contrib.postgres.operations. (#1071)

* Improve stubs for contrib.postgres.operations.

This adds a bunch of missing operations in contrib.postgres.operations.

Documentation: https://docs.djangoproject.com/en/4.0/ref/contrib/postgres/operations/#managing-collations-using-migrations
Source code: https://github.com/django/django/blob/2fac0a18081dcc77fc860c801e5d727dc90435b3/django/contrib/postgres/operations.py

Signed-off-by: Zixuan James Li <p359101898@gmail.com>

* Add migration_name_fragment to Operation.

Signed-off-by: Zixuan James Li <p359101898@gmail.com>

* Remove redefinition of methods in subclasses.

Signed-off-by: Zixuan James Li <p359101898@gmail.com>
---
 django-stubs/contrib/postgres/operations.pyi  | 35 +++++++++++++++++++
 .../db/migrations/operations/base.pyi         |  2 ++
 2 files changed, 37 insertions(+)

diff --git a/django-stubs/contrib/postgres/operations.pyi b/django-stubs/contrib/postgres/operations.pyi
index 94fb10a..b901495 100644
--- a/django-stubs/contrib/postgres/operations.pyi
+++ b/django-stubs/contrib/postgres/operations.pyi
@@ -1,9 +1,17 @@
+from typing import Dict, Literal, Tuple
+
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.migrations import AddConstraint, AddIndex, RemoveIndex
 from django.db.migrations.operations.base import Operation
 
 class CreateExtension(Operation):
     reversible: bool = ...
     name: str = ...
     def __init__(self, name: str) -> None: ...
+    def extension_exists(self, schema_editor: BaseDatabaseSchemaEditor, extension: str) -> bool: ...
+
+class BloomExtension(CreateExtension):
+    def __init__(self) -> None: ...
 
 class BtreeGinExtension(CreateExtension):
     def __init__(self) -> None: ...
@@ -25,3 +33,30 @@ class TrigramExtension(CreateExtension):
 
 class UnaccentExtension(CreateExtension):
     def __init__(self) -> None: ...
+
+class NotInTransactionMixin:
+    def _ensure_not_in_transaction(self, schema_editor: BaseDatabaseSchemaEditor) -> None: ...
+
+class AddIndexConcurrently(NotInTransactionMixin, AddIndex):
+    atomic: Literal[False] = ...
+
+class RemoveIndexConcurrently(NotInTransactionMixin, RemoveIndex):
+    atomic: Literal[False] = ...
+
+class CollationOperation(Operation):
+    name: str
+    locale: str
+    provider: str
+    deterministic: bool
+    def __init__(self, name: str, locale: str, *, provider: str = ..., deterministic: bool = ...) -> None: ...
+    def create_collation(self, schema_editor: BaseDatabaseSchemaEditor) -> None: ...
+    def remove_collation(self, schema_editor: BaseDatabaseSchemaEditor) -> None: ...
+
+class CreateCollation(CollationOperation): ...
+class RemoveCollation(CollationOperation): ...
+class AddConstraintNotValid(AddConstraint): ...
+
+class ValidateConstraint(Operation):
+    model_name: str
+    name: str
+    def __init__(self, model_name: str, name: str) -> None: ...
diff --git a/django-stubs/db/migrations/operations/base.pyi b/django-stubs/db/migrations/operations/base.pyi
index 29e273a..2896df3 100644
--- a/django-stubs/db/migrations/operations/base.pyi
+++ b/django-stubs/db/migrations/operations/base.pyi
@@ -25,3 +25,5 @@ class Operation:
     def references_field(self, model_name: str, name: str, app_label: str) -> bool: ...
     def allow_migrate_model(self, connection_alias: Union[BaseDatabaseWrapper, str], model: Type[Model]) -> bool: ...
     def reduce(self, operation: Operation, app_label: str) -> Union[bool, List[Operation]]: ...
+    @property
+    def migration_name_fragment(self) -> str: ...