diff --git a/pandas/io/sql.py b/pandas/io/sql.py index 9024961688c6b..5a5702bf36587 100644 --- a/pandas/io/sql.py +++ b/pandas/io/sql.py @@ -1636,6 +1636,8 @@ def run_transaction(self): def execute(self, sql: str | Select | TextClause, params=None): """Simple passthrough to SQLAlchemy connectable""" args = [] if params is None else [params] + if self.returns_generator: + self.con.execution_options(stream_results=True) if isinstance(sql, str): return self.con.exec_driver_sql(sql, *args) return self.con.execute(sql, *args) @@ -1814,11 +1816,11 @@ def read_query( read_sql """ + self.returns_generator = chunksize is not None result = self.execute(sql, params) columns = result.keys() if chunksize is not None: - self.returns_generator = True return self._query_iterator( result, self.exit_stack, diff --git a/pandas/tests/io/test_sql.py b/pandas/tests/io/test_sql.py index c8f4d68230e5b..fb900a033f6f5 100644 --- a/pandas/tests/io/test_sql.py +++ b/pandas/tests/io/test_sql.py @@ -1189,6 +1189,24 @@ def test_read_iris_table_chunksize(conn, request): check_iris_frame(iris_frame) +@pytest.mark.parametrize("conn", sqlalchemy_connectable_iris) +def test_sqlalchemy_read_table_chunksize_stream_results(conn, request): + conn = request.getfixturevalue(conn) + with sql.SQLDatabase(conn) as pandasSQL: + assert list(pandasSQL.read_table("iris", chunksize=10000)) + execution_options = pandasSQL.con.get_execution_options() + assert execution_options["stream_results"] is True + + +@pytest.mark.parametrize("conn", sqlalchemy_connectable_iris) +def test_sqlalchemy_read_sql_chunksize_stream_results(conn, request): + conn = request.getfixturevalue(conn) + with sql.SQLDatabase(conn) as pandasSQL: + assert list(pandasSQL.read_sql("SELECT * FROM iris", chunksize=10000)) + execution_options = pandasSQL.con.get_execution_options() + assert execution_options["stream_results"] is True + + @pytest.mark.parametrize("conn", sqlalchemy_connectable) def test_to_sql_callable(conn, test_frame1, request): conn = request.getfixturevalue(conn)