From bac3595e99abd2afb094493dee9ea62bf4c7f43d Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 5 Jul 2025 23:58:41 +0200
Subject: [PATCH] feat(pq): add libpq OAuth support structures and functions

---
 psycopg/psycopg/pq/_enums.py      | 12 +++++++
 psycopg/psycopg/pq/_pq_ctypes.py  | 60 +++++++++++++++++++++++++++++++
 psycopg/psycopg/pq/_pq_ctypes.pyi |  7 ++++
 3 files changed, 79 insertions(+)

diff --git a/psycopg/psycopg/pq/_enums.py b/psycopg/psycopg/pq/_enums.py
index 3f82c489e..54ea62116 100644
--- a/psycopg/psycopg/pq/_enums.py
+++ b/psycopg/psycopg/pq/_enums.py
@@ -257,3 +257,15 @@ class Trace(IntFlag):
 
     REGRESS_MODE = 2
     """Redact some fields, e.g. OIDs, from messages."""
+
+
+class AuthData(IntEnum):
+    """
+    Enum to represent the available OAuth hook types.
+    """
+
+    PROMPT_OAUTH_DEVICE = 0
+    """user must visit a device-authorization URL."""
+
+    OAUTH_BEARER_TOKEN = auto()
+    """Server requests an OAuth Bearer token."""
diff --git a/psycopg/psycopg/pq/_pq_ctypes.py b/psycopg/psycopg/pq/_pq_ctypes.py
index 846b29559..549192e15 100644
--- a/psycopg/psycopg/pq/_pq_ctypes.py
+++ b/psycopg/psycopg/pq/_pq_ctypes.py
@@ -759,6 +759,66 @@ PQinitOpenSSL.argtypes = [c_int, c_int]
 PQinitOpenSSL.restype = None
 
 
+# 32.20. OAuth support
+
+
+class PGpromptOAuthDevice_struct(Structure):
+    _fields_ = [
+        ("verification_uri", c_char_p),
+        ("user_code", c_char_p),
+        ("verification_uri_complete", c_char_p),
+        ("expires_in", c_int),
+    ]
+
+
+class PGoauthBearerRequest_struct(Structure):
+    # Struct forward declaration to allow to define the callbacks
+    pass
+
+
+PGpromptOAuthDevice_ptr = POINTER(PGpromptOAuthDevice_struct)
+PGoauthBearerRequest_ptr = POINTER(PGoauthBearerRequest_struct)
+
+
+# TODO: not sure - uintptr_t on win32 in libpq source
+SOCKTYPE = c_void_p if sys.platform == "win32" else c_int
+
+PGoauthBearerRequestAsyncCB = CFUNCTYPE(
+    c_int, PGconn_ptr, PGoauthBearerRequest_ptr, POINTER(SOCKTYPE)
+)
+PGoauthBearerRequestCleanupCB = CFUNCTYPE(
+    None, PGconn_ptr, POINTER(PGoauthBearerRequest_struct)
+)
+
+PGoauthBearerRequest_struct._fields_ = [
+    ("openid_configuration", c_char_p),
+    ("scope", c_char_p),
+    ("async", PGoauthBearerRequestAsyncCB),
+    ("cleanup", PGoauthBearerRequestCleanupCB),
+    ("token", c_char_p),
+    ("user", c_void_p),
+]
+
+PQauthDataHook_type = CFUNCTYPE(c_int, c_int, PGconn_ptr, c_void_p)
+
+if libpq_version >= 180000:
+    PQsetAuthDataHook = pq.PQsetAuthDataHook
+    PQsetAuthDataHook.argtypes = [PQauthDataHook_type]
+    PQsetAuthDataHook.restype = None
+
+    PQgetAuthDataHook = pq.PQgetAuthDataHook
+    PQgetAuthDataHook.argtypes = []
+    PQgetAuthDataHook.restype = PQauthDataHook_type
+
+    PQdefaultAuthDataHook = pq.PQdefaultAuthDataHook
+    PQdefaultAuthDataHook.argtypes = [c_int, PGconn_ptr, c_void_p]
+    PQdefaultAuthDataHook.restype = c_int
+else:
+    PQsetAuthDataHook = not_supported_before("PQsetAuthDataHook", 180000)
+    PQgetAuthDataHook = not_supported_before("PQgetAuthDataHook", 180000)
+    PQdefaultAuthDataHook = not_supported_before("PQdefaultAuthDataHook", 180000)
+
+
 def generate_stub() -> None:
     import re
     from ctypes import _CFuncPtr  # type: ignore[attr-defined]
diff --git a/psycopg/psycopg/pq/_pq_ctypes.pyi b/psycopg/psycopg/pq/_pq_ctypes.pyi
index 3936d08db..ebb42808c 100644
--- a/psycopg/psycopg/pq/_pq_ctypes.pyi
+++ b/psycopg/psycopg/pq/_pq_ctypes.pyi
@@ -134,6 +134,13 @@ def PQenterPipelineMode(pgconn: PGconn_struct | None) -> int: ...
 def PQexitPipelineMode(pgconn: PGconn_struct | None) -> int: ...
 def PQpipelineSync(pgconn: PGconn_struct | None) -> int: ...
 def PQsendFlushRequest(pgconn: PGconn_struct | None) -> int: ...
+def PQsetAuthDataHook(
+    hook: Callable[[int, PGconn_struct | None, Any], int],
+) -> None: ...
+def PQgetAuthDataHook() -> Callable[[int, PGconn_struct | None, Any], int]: ...
+def PQdefaultAuthDataHook(
+    type: int, pgconn: PGconn_struct | None, data: Any
+) -> int: ...
 
 # Autogenerated section.
 # In order to refresh, run:
-- 
2.47.2