BUG: Allow read_sql to work with chunksize. by cdcadman · Pull Request #49967 · pandas-dev/pandas (original) (raw)

To keep the connection open, I need to make sure that the context manager at line 864 stays open. My solution is to create generating functions _read_sql_table, _read_sql_query, and _read_sql. If chunksize=None, these generating functions return a single DataFrame. Then I have the opposite problem: if the caller wants a single DataFrame, I have to exit out of the context manager at line 864 so that the connection will close. If instead of

for frame in result:
    pass
return frame

I use

then the line 864 context manager would not exit. I agree it is roundabout, but I think having SQLDatabase only accept a Connection, not an Engine, and having SQLDatabase.read_query only return an Iterator[DataFrame] also simplifies things a bit.