Add pre-commit hooks and scripts for async method checks in PGMQueue

Distinguish sync and async operations in PGMQueue

- Introduced a pre-commit hook to check for missing async methods in PGMQueue.
- Added scripts to identify and generate missing async methods.
- Created utility functions for AST manipulation and method transformation.
- Established configuration for project paths and console output.

Author: jason810496

.pre-commit-config.yaml

@@ -7,4 +7,10 @@ repos:
       - id: ruff
         args: [ --fix ]
       # Run the formatter.
-      - id: ruff-format
+      - id: ruff-format
+  - repo: local
+    hooks:
+      - id: check-sync-async-method-for-queue
+        name: Check sync/async method for queue
+        entry: ./scripts/ci/pre_commit/check_sync_async_method_for_queue.py
+        language: python

pgmq_sqlalchemy/queue.py

@@ -1,9 +1,11 @@
 import asyncio
-from typing import List, Optional
+from typing import List, Optional, TYPE_CHECKING
 
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.ext.asyncio import create_async_engine
+from asgiref.sync import async_to_sync
+
 
 from .schema import Message, QueueMetrics
 from ._types import ENGINE_TYPE, SESSION_TYPE
@@ -14,6 +16,10 @@
 )
 from .operation import PGMQOperation
 
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+    from sqlalchemy.ext.asyncio import AsyncSession
+
 
 class PGMQueue:
     engine: ENGINE_TYPE = None
@@ -118,47 +124,20 @@ def __init__(
         # create pgmq extension if not exists
         self._check_pgmq_ext()
 
-    async def _check_pgmq_ext_async(self) -> None:
-        """Check if the pgmq extension exists."""
-        async with self.session_maker() as session:
-            await PGMQOperation.check_pgmq_ext_async(session=session, commit=True)
-
-    def _check_pgmq_ext_sync(self) -> None:
-        """Check if the pgmq extension exists."""
-        with self.session_maker() as session:
-            PGMQOperation.check_pgmq_ext(session=session, commit=True)
-
     def _check_pgmq_ext(self) -> None:
         """Check if the pgmq extension exists."""
-        if self.is_async:
-            return self.loop.run_until_complete(self._check_pgmq_ext_async())
-        return self._check_pgmq_ext_sync()
-
-    async def _check_pg_partman_ext_async(self) -> None:
-        """Check if the pg_partman extension exists."""
-        async with self.session_maker() as session:
-            await PGMQOperation.check_pg_partman_ext_async(session=session, commit=True)
-
-    def _check_pg_partman_ext_sync(self) -> None:
-        """Check if the pg_partman extension exists."""
-        with self.session_maker() as session:
-            PGMQOperation.check_pg_partman_ext(session=session, commit=True)
+        self._execute_operation(PGMQOperation.check_pgmq_ext, session=None, commit=True)
 
     def _check_pg_partman_ext(self) -> None:
         """Check if the pg_partman extension exists."""
-        if self.is_pg_partman_ext_checked:
-            return
-        self.is_pg_partman_ext_checked
-
-        if self.is_async:
-            return self.loop.run_until_complete(self._check_pg_partman_ext_async())
-        return self._check_pg_partman_ext_sync()
+        self._execute_operation(
+            PGMQOperation.check_pg_partman_ext, session=None, commit=True
+        )
 
     def _execute_operation(
         self,
         op_sync,
-        op_async,
-        session: Optional[SESSION_TYPE],
+        session: Optional["Session"],
         commit: bool,
         *args,
         **kwargs,
@@ -167,7 +146,6 @@ def _execute_operation(
 
         Args:
             op_sync: The synchronous operation function from PGMQOperation
-            op_async: The asynchronous operation function from PGMQOperation
             session: Optional session to use (if None, creates a new one)
             commit: Whether to commit the transaction
             *args: Positional arguments to pass to the operation
@@ -176,23 +154,36 @@ def _execute_operation(
         Returns:
             The result from the operation
         """
-        if self.is_async:
-            if session is None:
-
-                async def _run():
-                    async with self.session_maker() as s:
-                        return await op_async(*args, session=s, commit=commit, **kwargs)
-
-                return self.loop.run_until_complete(_run())
-            return self.loop.run_until_complete(
-                op_async(*args, session=session, commit=commit, **kwargs)
-            )
-
         if session is None:
             with self.session_maker() as s:
                 return op_sync(*args, session=s, commit=commit, **kwargs)
         return op_sync(*args, session=session, commit=commit, **kwargs)
 
+    async def _execute_async_operation(
+        self,
+        op_async,
+        session: Optional["AsyncSession"],
+        commit: bool,
+        *args,
+        **kwargs,
+    ):
+        """Helper method to execute sync or async operations with session management.
+
+        Args:
+            op_async: The asynchronous operation function from PGMQOperation
+            session: Optional session to use (if None, creates a new one)
+            commit: Whether to commit the transaction
+            *args: Positional arguments to pass to the operation
+            **kwargs: Keyword arguments to pass to the operation
+
+        Returns:
+            The result from the operation
+        """
+        if session is None:
+            async with self.session_maker() as s:
+                return await op_async(*args, session=s, commit=commit, **kwargs)
+        return await op_async(*args, session=session, commit=commit, **kwargs)
+
     def create_queue(
         self,
         queue_name: str,
@@ -219,7 +210,6 @@ def create_queue(
         """
         return self._execute_operation(
             PGMQOperation.create_queue,
-            PGMQOperation.create_queue_async,
             session,
             commit,
             queue_name,
@@ -272,7 +262,6 @@ def create_partitioned_queue(
 
         return self._execute_operation(
             PGMQOperation.create_partitioned_queue,
-            PGMQOperation.create_partitioned_queue_async,
             session,
             commit,
             queue_name,
@@ -292,7 +281,6 @@ def validate_queue_name(
         """
         return self._execute_operation(
             PGMQOperation.validate_queue_name,
-            PGMQOperation.validate_queue_name_async,
             session,
             commit,
             queue_name,
@@ -329,7 +317,6 @@ def drop_queue(
 
         return self._execute_operation(
             PGMQOperation.drop_queue,
-            PGMQOperation.drop_queue_async,
             session,
             commit,
             queue,
@@ -351,7 +338,6 @@ def list_queues(
         """
         return self._execute_operation(
             PGMQOperation.list_queues,
-            PGMQOperation.list_queues_async,
             session,
             commit,
         )
@@ -385,7 +371,6 @@ def send(
         """
         return self._execute_operation(
             PGMQOperation.send,
-            PGMQOperation.send_async,
             session,
             commit,
             queue_name,
@@ -416,7 +401,6 @@ def send_batch(
         """
         return self._execute_operation(
             PGMQOperation.send_batch,
-            PGMQOperation.send_batch_async,
             session,
             commit,
             queue_name,
@@ -496,7 +480,6 @@ def read(
 
         return self._execute_operation(
             PGMQOperation.read,
-            PGMQOperation.read_async,
             session,
             commit,
             queue_name,
@@ -533,7 +516,6 @@ def read_batch(
 
         return self._execute_operation(
             PGMQOperation.read_batch,
-            PGMQOperation.read_batch_async,
             session,
             commit,
             queue_name,
@@ -605,7 +587,6 @@ def read_with_poll(
 
         return self._execute_operation(
             PGMQOperation.read_with_poll,
-            PGMQOperation.read_with_poll_async,
             session,
             commit,
             queue_name,
@@ -683,7 +664,6 @@ def consumer_with_backoff_retry(pgmq_client: PGMQueue, queue_name: str):
 
         return self._execute_operation(
             PGMQOperation.set_vt,
-            PGMQOperation.set_vt_async,
             session,
             commit,
             queue_name,
@@ -710,7 +690,6 @@ def pop(
         """
         return self._execute_operation(
             PGMQOperation.pop,
-            PGMQOperation.pop_async,
             session,
             commit,
             queue_name,
@@ -743,7 +722,6 @@ def delete(
         """
         return self._execute_operation(
             PGMQOperation.delete,
-            PGMQOperation.delete_async,
             session,
             commit,
             queue_name,
@@ -776,7 +754,6 @@ def delete_batch(
         """
         return self._execute_operation(
             PGMQOperation.delete_batch,
-            PGMQOperation.delete_batch_async,
             session,
             commit,
             queue_name,
@@ -813,7 +790,6 @@ def archive(
         """
         return self._execute_operation(
             PGMQOperation.archive,
-            PGMQOperation.archive_async,
             session,
             commit,
             queue_name,
@@ -843,7 +819,6 @@ def archive_batch(
         """
         return self._execute_operation(
             PGMQOperation.archive_batch,
-            PGMQOperation.archive_batch_async,
             session,
             commit,
             queue_name,
@@ -870,7 +845,6 @@ def purge(
         """
         return self._execute_operation(
             PGMQOperation.purge,
-            PGMQOperation.purge_async,
             session,
             commit,
             queue_name,
@@ -903,7 +877,6 @@ def metrics(
         """
         return self._execute_operation(
             PGMQOperation.metrics,
-            PGMQOperation.metrics_async,
             session,
             commit,
             queue_name,
@@ -950,7 +923,6 @@ def metrics_all(
         """
         return self._execute_operation(
             PGMQOperation.metrics_all,
-            PGMQOperation.metrics_all_async,
             session,
             commit,
         )

scripts/ci/pre_commit/check_sync_async_method_for_queue.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+# /// script
+# requires-python = ">=3.10,<3.11"
+# dependencies = [
+#   "rich>=13.6.0",
+# ]
+# ///
+"""
+Script to check for missing async methods in PGMQueue for per-commit.
+
+For each public sync method (not starting with _), checks if there's a corresponding
+async method with the same name plus '_async' suffix.
+"""
+
+import ast
+import sys
+from pathlib import Path
+
+
+sys.path.insert(0, str(Path(__name__).parent.parent.joinpath("scripts").resolve()))
+
+from scripts_utils.config import QUEUE_FILE  # noqa: E402
+from scripts_utils.console import console  # noqa: E402
+from scripts_utils.common_ast import parse_methods_info_from_target_class  # noqa: E402
+
+
+def main():
+    """Main function."""
+
+    module_tree = ast.parse(source=QUEUE_FILE.read_text(), filename=QUEUE_FILE)
+    _, missing_async = parse_methods_info_from_target_class(
+        module_tree, target_class="PGMQueue"
+    )
+
+    if not missing_async:
+        console.print(
+            "[bold green]SUCCESS:[/bold green] All public methods have corresponding async versions!"
+        )
+        sys.exit(0)
+
+    # log all the missing async methods
+    console.print()
+    console.print(
+        f"[bold yellow]WARNING:[/bold yellow] Found {len(missing_async)} missing async methods:",
+        style="bold",
+    )
+    for method in missing_async:
+        console.print(f"  [yellow]-[/yellow] {method}_async")
+    console.print()
+
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()