Repository: pola-rs/polars
File: py-polars/tests/unit/io/database/test_read.py (Python)

Code before:

    t = Table("test_data", MetaData(), autoload_with=alchemy_engine)

    # establish sqlalchemy "selectable" and validate usage
    selectable_query = select(
        alchemy_cast(func.strftime("%Y", t.c.date), Integer).label("year"),
        t.c.name,
        t.c.value,
    ).where(t.c.value < 0)

    expected = pl.DataFrame({"year": [2021], "name": ["other"], "value": [-99.5]})

    for conn in (alchemy_session, alchemy_engine, alchemy_conn):
        assert_frame_equal(
            pl.read_database(selectable_query, connection=conn),
            expected,
        )

        batches = list(
            pl.read_database(
                selectable_query,
                connection=conn,
                iter_batches=True,
                batch_size=1,
            )
        )
        assert len(batches) == 1
        assert_frame_equal(batches[0], expected)


def test_read_database_alchemy_textclause(tmp_sqlite_db: Path) -> None:
    # various flavours of alchemy connection
    alchemy_engine = create_engine(f"sqlite:///{tmp_sqlite_db}")
    alchemy_session: ConnectionOrCursor = sessionmaker(bind=alchemy_engine)()
    alchemy_conn: ConnectionOrCursor = alchemy_engine.connect()

    # establish sqlalchemy "textclause" and validate usage
    textclause_query = text(
        """
                SELECT CAST(STRFTIME('%Y',"date") AS INT) as "year", name, value
                FROM test_data
                WHERE value < 0
            """
    )

    expected = pl.DataFrame({"year": [2021], "name": ["other"], "value": [-99.5]})

    for conn in (alchemy_session, alchemy_engine, alchemy_conn):
        assert_frame_equal(
            pl.read_database(textclause_query, connection=conn),
            expected,

Review comment (suggestion):
drive by - should be indented as iterator variable `conn` is referenced

Code after:

    # establish sqlalchemy "selectable" and validate usage
    selectable_query = select(
        alchemy_cast(func.strftime("%Y", t.c.date), Integer).label("year"),
        t.c.name,
        t.c.value,
    ).where(t.c.value < 0)

    expected = pl.DataFrame({"year": [2021], "name": ["other"], "value": [-99.5]})

    for conn in (alchemy_session, alchemy_engine, alchemy_conn):
        assert_frame_equal(
            pl.read_database(selectable_query, connection=conn),
            expected,
        )

        batches = list(
            pl.read_database(
                selectable_query,
                connection=conn,
                iter_batches=True,
                batch_size=1,
            )
        )
        assert len(batches) == 1
        assert_frame_equal(batches[0], expected)


def test_read_database_alchemy_textclause(tmp_sqlite_db: Path) -> None:
    # various flavours of alchemy connection
    alchemy_engine = create_engine(f"sqlite:///{tmp_sqlite_db}")
    alchemy_session: ConnectionOrCursor = sessionmaker(bind=alchemy_engine)()
    alchemy_conn: ConnectionOrCursor = alchemy_engine.connect()

    # establish sqlalchemy "textclause" and validate usage
    textclause_query = text(
        """
                SELECT CAST(STRFTIME('%Y',"date") AS INT) as "year", name, value
                FROM test_data
                WHERE value < 0
            """
    )

    expected = pl.DataFrame({"year": [2021], "name": ["other"], "value": [-99.5]})

    for conn in (alchemy_session, alchemy_engine, alchemy_conn):
        assert_frame_equal(
            pl.read_database(textclause_query, connection=conn),
            expected,
        )