From 09564e95ac9670c775888620e2709e061efa9e8c Mon Sep 17 00:00:00 2001
From: Anders Kaseorg <anders@zulip.com>
Date: Mon, 9 Aug 2021 15:58:58 -0700
Subject: [PATCH] mypy: Add types-psycopg2.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
---
 analytics/lib/counts.py                           | 2 +-
 puppet/zulip/files/postgresql/process_fts_updates | 2 ++
 pyproject.toml                                    | 1 -
 requirements/dev.txt                              | 4 ++++
 requirements/mypy.in                              | 1 +
 requirements/mypy.txt                             | 4 ++++
 version.py                                        | 2 +-
 zerver/lib/retention.py                           | 4 ++--
 zerver/lib/test_helpers.py                        | 2 +-
 9 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/analytics/lib/counts.py b/analytics/lib/counts.py
index 3aad5e018c..761eaf094e 100644
--- a/analytics/lib/counts.py
+++ b/analytics/lib/counts.py
@@ -356,7 +356,7 @@ def do_pull_by_sql_query(
     group_by: Optional[Tuple[Type[models.Model], str]],
 ) -> int:
     if group_by is None:
-        subgroup = SQL("NULL")
+        subgroup: Composable = SQL("NULL")
         group_by_clause = SQL("")
     else:
         subgroup = Identifier(group_by[0]._meta.db_table, group_by[1])
diff --git a/puppet/zulip/files/postgresql/process_fts_updates b/puppet/zulip/files/postgresql/process_fts_updates
index 4954cfaf8a..472a6b5a12 100755
--- a/puppet/zulip/files/postgresql/process_fts_updates
+++ b/puppet/zulip/files/postgresql/process_fts_updates
@@ -33,6 +33,7 @@ import os
 import select
 import sys
 import time
+from typing import Optional
 
 import psycopg2
 import psycopg2.extensions
@@ -136,6 +137,7 @@ except ImportError:
         pg_args["host"] = "localhost"
     pg_args["dbname"] = get_config(config_file, "postgresql", "database_name", "zulip")
 
+conn: Optional[psycopg2.extensions.connection]
 
 if options.nagios_check:
     conn = psycopg2.connect(**pg_args)
diff --git a/pyproject.toml b/pyproject.toml
index 6efefc5666..f1f77d6c5a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -80,7 +80,6 @@ module = [
     "openapi_core.*",
     "openapi_schema_validator.*",
     "premailer.*",
-    "psycopg2.*",
     "pyinotify.*",
     "pyoembed.*",
     "pyuca.*",
diff --git a/requirements/dev.txt b/requirements/dev.txt
index 910786296c..2eb060b342 100644
--- a/requirements/dev.txt
+++ b/requirements/dev.txt
@@ -1696,6 +1696,10 @@ types-polib==0.1.3 \
     --hash=sha256:54fb4d5a9edc2b2f7e3cd7ebbf8631dbcecdeb69e9e31597e81c2b5b3bd3613b \
     --hash=sha256:bd79bc1780e5d7fb6fd0dce59bb8d118ca81f8b2f14d5d05e21d2716e2e1937b
     # via -r requirements/mypy.in
+types-psycopg2==2.9.0 \
+    --hash=sha256:6e15cdbae8944271ab5fe413a06b248b714666ba0af4162b62b35da6c7a74e8f \
+    --hash=sha256:c1cd002bb62c044a02a88c48a06b62242d4bd8732d8dd8bf064b29fdaf7771e5
+    # via -r requirements/mypy.in
 types-pygments==2.9.0 \
     --hash=sha256:04dd3054e579d739e61c48757962969285ebdd09aca1bdac7c6d210455fb8f98 \
     --hash=sha256:4dc3c9b1e96bce89b30b32503d9976b2a66dbcf0d285aabc59ffd7302dbf199f
diff --git a/requirements/mypy.in b/requirements/mypy.in
index 194325ac25..b38dfffa97 100644
--- a/requirements/mypy.in
+++ b/requirements/mypy.in
@@ -19,6 +19,7 @@ types-oauthlib
 types-orjson
 types-polib
 types-Pillow
+types-psycopg2
 types-Pygments
 types-python-dateutil
 types-pytz
diff --git a/requirements/mypy.txt b/requirements/mypy.txt
index 2b1bfb33fb..8956d4598f 100644
--- a/requirements/mypy.txt
+++ b/requirements/mypy.txt
@@ -137,6 +137,10 @@ types-polib==0.1.3 \
     --hash=sha256:54fb4d5a9edc2b2f7e3cd7ebbf8631dbcecdeb69e9e31597e81c2b5b3bd3613b \
     --hash=sha256:bd79bc1780e5d7fb6fd0dce59bb8d118ca81f8b2f14d5d05e21d2716e2e1937b
     # via -r requirements/mypy.in
+types-psycopg2==2.9.0 \
+    --hash=sha256:6e15cdbae8944271ab5fe413a06b248b714666ba0af4162b62b35da6c7a74e8f \
+    --hash=sha256:c1cd002bb62c044a02a88c48a06b62242d4bd8732d8dd8bf064b29fdaf7771e5
+    # via -r requirements/mypy.in
 types-pygments==2.9.0 \
     --hash=sha256:04dd3054e579d739e61c48757962969285ebdd09aca1bdac7c6d210455fb8f98 \
     --hash=sha256:4dc3c9b1e96bce89b30b32503d9976b2a66dbcf0d285aabc59ffd7302dbf199f
diff --git a/version.py b/version.py
index f184106284..0eb2a09dc7 100644
--- a/version.py
+++ b/version.py
@@ -48,4 +48,4 @@ API_FEATURE_LEVEL = 92
 #   historical commits sharing the same major version, in which case a
 #   minor version bump suffices.
 
-PROVISION_VERSION = "153.11"
+PROVISION_VERSION = "153.12"
diff --git a/zerver/lib/retention.py b/zerver/lib/retention.py
index cef45dc27b..02009a5409 100644
--- a/zerver/lib/retention.py
+++ b/zerver/lib/retention.py
@@ -93,7 +93,7 @@ models_with_message_key: List[Dict[str, Any]] = [
 @transaction.atomic(savepoint=False)
 def move_rows(
     base_model: Model,
-    raw_query: Composable,
+    raw_query: SQL,
     *,
     src_db_table: Optional[str] = None,
     returning_id: bool = False,
@@ -122,7 +122,7 @@ def move_rows(
 
 
 def run_archiving_in_chunks(
-    query: Composable,
+    query: SQL,
     type: int,
     realm: Optional[Realm] = None,
     chunk_size: int = MESSAGE_BATCH_SIZE,
diff --git a/zerver/lib/test_helpers.py b/zerver/lib/test_helpers.py
index 87c6773c3f..ba467b2386 100644
--- a/zerver/lib/test_helpers.py
+++ b/zerver/lib/test_helpers.py
@@ -149,7 +149,7 @@ def queries_captured(
 
     def wrapper_execute(
         self: TimeTrackingCursor,
-        action: Callable[[str, ParamsT], None],
+        action: Callable[[Query, ParamsT], None],
         sql: Query,
         params: ParamsT,
     ) -> None: