From 189039b9d38343b482f1b077bbcf6f6ae99cbacd Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Wed, 10 May 2023 11:08:07 -0400
Subject: add full parameter types for ORM with_for_update

Fixed typing for the :paramref:`_orm.Session.get.with_for_update` parameter
of :meth:`_orm.Session.get` and :meth:`_orm.Session.refresh` (as well as
corresponding methods on :class:`_asyncio.AsyncSession`) to accept boolean
``True`` and all other argument forms accepted by the parameter at runtime.

Fixes: #9762
Change-Id: Ied4d37a269906b3d9be5ab7d31a2fa863360cced
---
 test/ext/mypy/plain_files/session.py | 44 ++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

(limited to 'test')

diff --git a/test/ext/mypy/plain_files/session.py b/test/ext/mypy/plain_files/session.py
index 9106b9016..dfebdd5a9 100644
--- a/test/ext/mypy/plain_files/session.py
+++ b/test/ext/mypy/plain_files/session.py
@@ -1,14 +1,20 @@
 from __future__ import annotations
 
+import asyncio
 from typing import List
 
 from sqlalchemy import create_engine
 from sqlalchemy import ForeignKey
+from sqlalchemy.ext.asyncio import async_scoped_session
+from sqlalchemy.ext.asyncio import async_sessionmaker
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
 from sqlalchemy.orm import relationship
+from sqlalchemy.orm import scoped_session
 from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 
 
 class Base(DeclarativeBase):
@@ -94,3 +100,41 @@ with Session(e) as sess:
     ).offset(User.id)
 
 # more result tests in typed_results.py
+
+
+def test_with_for_update() -> None:
+    """test #9762"""
+    sess = Session()
+    ss = scoped_session(sessionmaker())
+
+    sess.get(User, 1)
+    sess.get(User, 1, with_for_update=True)
+    ss.get(User, 1)
+    ss.get(User, 1, with_for_update=True)
+
+    u1 = User()
+    sess.refresh(u1)
+    sess.refresh(u1, with_for_update=True)
+    ss.refresh(u1)
+    ss.refresh(u1, with_for_update=True)
+
+
+async def test_with_for_update_async() -> None:
+    """test #9762"""
+    sess = AsyncSession()
+    ss = async_scoped_session(
+        async_sessionmaker(), scopefunc=asyncio.current_task
+    )
+
+    await sess.get(User, 1)
+    await sess.get(User, 1, with_for_update=True)
+
+    await ss.get(User, 1)
+    await ss.get(User, 1, with_for_update=True)
+
+    u1 = User()
+    await sess.refresh(u1)
+    await sess.refresh(u1, with_for_update=True)
+
+    await ss.refresh(u1)
+    await ss.refresh(u1, with_for_update=True)
-- 
cgit v1.2.1