Allow tuple-valued params in read_sql[_query]

pandas-dev · Sep 12, 2024 · 39a96fb · 39a96fb
1 parent fa7e444
commit 39a96fb
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 4 deletions.
diff --git a/pandas-stubs/io/sql.pyi b/pandas-stubs/io/sql.pyi
@@ -66,7 +66,14 @@ def read_sql_query(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[Scalar] | tuple[Scalar, ...] | Mapping[str, Scalar] | None = ...,
+    params: (
+        list[Scalar]
+        | tuple[Scalar, ...]
+        | tuple[tuple[Scalar, ...], ...]
+        | Mapping[str, Scalar]
+        | Mapping[str, tuple[Scalar, ...]]
+        | None
+    ) = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     *,
     chunksize: int,
@@ -79,7 +86,14 @@ def read_sql_query(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[Scalar] | tuple[Scalar, ...] | Mapping[str, Scalar] | None = ...,
+    params: (
+        list[Scalar]
+        | tuple[Scalar, ...]
+        | tuple[tuple[Scalar, ...], ...]
+        | Mapping[str, Scalar]
+        | Mapping[str, tuple[Scalar, ...]]
+        | None
+    ) = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     chunksize: None = ...,
     dtype: DtypeArg | None = ...,
@@ -91,7 +105,14 @@ def read_sql(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[Scalar] | tuple[Scalar, ...] | Mapping[str, Scalar] | None = ...,
+    params: (
+        list[Scalar]
+        | tuple[Scalar, ...]
+        | tuple[tuple[Scalar, ...], ...]
+        | Mapping[str, Scalar]
+        | Mapping[str, tuple[Scalar, ...]]
+        | None
+    ) = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     columns: list[str] = ...,
     *,
@@ -105,7 +126,14 @@ def read_sql(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[Scalar] | tuple[Scalar, ...] | Mapping[str, Scalar] | None = ...,
+    params: (
+        list[Scalar]
+        | tuple[Scalar, ...]
+        | tuple[tuple[Scalar, ...], ...]
+        | Mapping[str, Scalar]
+        | Mapping[str, tuple[Scalar, ...]]
+        | None
+    ) = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     columns: list[str] = ...,
     chunksize: None = ...,

diff --git a/tests/test_io.py b/tests/test_io.py
@@ -1238,6 +1238,39 @@ def test_read_sql_query_via_sqlalchemy_engine_with_params():
         engine.dispose()
 
 
+@pytest.mark.skip(
+    reason="Only works in Postgres (and MySQL, but with different query syntax)"
+)
+def test_read_sql_query_via_sqlalchemy_engine_with_tuple_valued_params():
+    with ensure_clean() as path:
+        db_uri = "postgresql+psycopg2://postgres@localhost:5432/postgres"
+        engine = sqlalchemy.create_engine(db_uri)
+
+        check(
+            assert_type(
+                read_sql_query(
+                    "select * from test where a in %(a)s",
+                    con=engine,
+                    params={"a": (1, 2)},
+                ),
+                DataFrame,
+            ),
+            DataFrame,
+        )
+        check(
+            assert_type(
+                read_sql_query(
+                    "select * from test where a in %s",
+                    con=engine,
+                    params=((1, 2),),
+                ),
+                DataFrame,
+            ),
+            DataFrame,
+        )
+        engine.dispose()
+
+
 def test_read_html():
     check(assert_type(DF.to_html(), str), str)
     with ensure_clean() as path: