🐛 Support with_sqlalchemy_conn to wrap generator

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ ci:
   autoupdate_schedule: quarterly
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.15.11"
+    rev: "v0.15.12"
     hooks:
       - id: ruff-check
         args: [--fix, --exit-non-zero-on-fix]

CHANGELOG.md

@@ -21,6 +21,8 @@ All notable changes to this library are documented in this file.
 
 ### Bug Fixes
 
+- Support `with_sqlalchemy_conn` to decorate a generator.
+
 ## **1.27.0** (14 Apr 2026)
 
 ### API Changes

src/pyiem/database.py

@@ -1,11 +1,11 @@
 """Database helpers."""
 
-# stdlib
 import getpass
+import inspect
 from contextlib import contextmanager
+from functools import wraps
 from typing import Generator
 
-# third party
 import numpy as np
 import psycopg
 from psycopg.adapt import Dumper
@@ -213,6 +213,16 @@ def foo(args, conn=None, **kwargs):
     """
 
     def decorator(func):
+        if inspect.isgeneratorfunction(func):
+
+            @wraps(func)
+            def generator_wrapper(*args, **kwds):
+                with get_sqlalchemy_conn(name, **kwargs) as conn:
+                    yield from func(*args, **kwds, conn=conn)
+
+            return generator_wrapper
+
+        @wraps(func)
         def wrapper(*args, **kwds):
             with get_sqlalchemy_conn(name, **kwargs) as conn:
                 return func(*args, **kwds, conn=conn)

tests/test_database.py

@@ -1,6 +1,8 @@
 """Testing of pyiem.database"""
 
 # third party
+from contextlib import contextmanager
+
 import numpy as np
 import pytest
 
@@ -43,6 +45,38 @@ def foo(arg1, kwarg1=None, conn=None):
     foo(1, kwarg1=2)
 
 
+def test_with_sqlalchemy_conn_generator(monkeypatch):
+    """Test decorator keeps connection open while generator is consumed."""
+    state = {"entered": 0, "exited": 0}
+
+    @contextmanager
+    def fake_get_sqlalchemy_conn(_name, **_kwargs):
+        state["entered"] += 1
+        try:
+            yield object()
+        finally:
+            state["exited"] += 1
+
+    monkeypatch.setattr(
+        "pyiem.database.get_sqlalchemy_conn", fake_get_sqlalchemy_conn
+    )
+
+    @with_sqlalchemy_conn("coop")
+    def foo(conn=None):
+        assert conn is not None
+        yield 1
+        assert state["exited"] == 0
+        yield 2
+
+    gen = foo()
+    assert state["entered"] == 0
+    assert next(gen) == 1
+    assert state["entered"] == 1
+    assert state["exited"] == 0
+    assert list(gen) == [2]
+    assert state["exited"] == 1
+
+
 def test_sql_helper():
     """Test that we can do a contextmanager with this API."""
     sql = "SELECT * FROM {table} WHERE {limiter} state_abbr = :state"