from __future__ import annotations

import random
import typing
from typing import Any
from typing import Iterator

from sqlalchemy import create_engine
from sqlalchemy import select
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.orm import Session

if typing.TYPE_CHECKING:
    from sqlalchemy import FrozenResult
    from sqlalchemy import Result
    from sqlalchemy import Select
    from sqlalchemy import SQLColumnExpression


def windowed_query(
    session: Session,
    stmt: Select[Any],
    column: SQLColumnExpression[Any],
    windowsize: int,
) -> Iterator[Result[Any]]:
    """Given a Session and Select() object, organize and execute the statement
    such that it is invoked for ordered chunks of the total result.   yield
    out individual Result objects for each chunk.

    """

    # add the column we will window / sort on to the statement
    stmt = stmt.add_columns(column).order_by(column)
    last_id = None

    while True:
        subq = stmt

        # filter the statement on the previous "last id" we got, if any
        if last_id is not None:
            subq = subq.filter(column > last_id)

        # execute the query
        result: Result = session.execute(subq.limit(windowsize))

        # turn the Result into a FrozenResult that we can peek at the data
        # first, then spin off new Result objects
        frozen_result: FrozenResult = result.freeze()

        # get the raw data
        chunk = frozen_result().all()

        if not chunk:
            break

        # count how many columns we have and also get the "last id" fetched
        result_width = len(chunk[0])
        last_id = chunk[-1][-1]

        # get a new, unconsumed Result back from the FrozenResult
        yield_result: Result = frozen_result()

        # split off the last column (Result could use a slice method here)
        yield_result = yield_result.columns(*list(range(0, result_width - 1)))

        # yield it out
        yield yield_result


if __name__ == "__main__":

    class Base(DeclarativeBase):
        pass

    class Widget(Base):
        __tablename__ = "widget"
        id: Mapped[int] = mapped_column(primary_key=True)
        data: Mapped[int]

    e = create_engine("mysql://scott:tiger@localhost/test", echo=True)

    Base.metadata.drop_all(e)
    Base.metadata.create_all(e)

    # get some random list of unique values
    data = set([random.randint(1, 1000000) for i in range(10000)])

    s = Session(e)
    s.add_all([Widget(id=i, data=j) for i, j in enumerate(data, 1)])
    s.commit()

    q = select(Widget)

    for result in windowed_query(s, q, Widget.data, 1000):
        for widget in result.scalars():
            print("data:", widget.data)

from __future__ import annotations

import random
import typing
from typing import Any
from typing import Iterator

from sqlalchemy import and_
from sqlalchemy import create_engine
from sqlalchemy import func
from sqlalchemy import select
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.orm import Session

if typing.TYPE_CHECKING:
    from sqlalchemy import Result
    from sqlalchemy import Select
    from sqlalchemy import SQLColumnExpression


def column_windows(
    session: Session,
    stmt: Select[Any],
    column: SQLColumnExpression[Any],
    windowsize: int,
) -> Iterator[SQLColumnExpression[bool]]:
    """Return a series of WHERE clauses against
    a given column that break it into windows.

    Result is an iterable of WHERE clauses that are packaged with
    the individual ranges to select from.

    Requires a database that supports window functions.

    """

    rownum = func.row_number().over(order_by=column).label("rownum")

    subq = stmt.add_columns(rownum).subquery()
    subq_column = list(subq.columns)[-1]

    target_column = subq.corresponding_column(column)  # type: ignore  # will be fixed by #10326
    new_stmt = select(target_column)

    if windowsize > 1:
        new_stmt = new_stmt.filter(subq_column % windowsize == 1)

    """
    # the SQL statement here is intended to give us a list of ranges,
    # and looks like:

    SELECT anon_1.data
    FROM (SELECT widget.id AS id, widget.data AS data, row_number() OVER (ORDER BY widget.data) AS rownum
    FROM widget) AS anon_1
    WHERE anon_1.rownum %% %(rownum_1)s = %(param_1)s

    """

    intervals = list(session.scalars(new_stmt))

    # yield out WHERE clauses for each range
    while intervals:
        start = intervals.pop(0)
        if intervals:
            yield and_(column >= start, column < intervals[0])
        else:
            yield column >= start


def windowed_query(
    session: Session,
    stmt: Select[Any],
    column: SQLColumnExpression[Any],
    windowsize: int,
) -> Iterator[Result[Any]]:
    """Given a Session and Select() object, organize and execute the statement
    such that it is invoked for ordered chunks of the total result.   yield
    out individual Result objects for each chunk.

    """

    for whereclause in column_windows(session, stmt, column, windowsize):
        yield session.execute(stmt.filter(whereclause).order_by(column))


if __name__ == "__main__":

    class Base(DeclarativeBase):
        pass

    class Widget(Base):
        __tablename__ = "widget"
        id: Mapped[int] = mapped_column(primary_key=True)
        data: Mapped[int]

    e = create_engine("postgresql://scott:tiger@localhost/test", echo=True)

    Base.metadata.drop_all(e)
    Base.metadata.create_all(e)

    # get some random list of unique values
    data = set([random.randint(1, 1000000) for i in range(10000)])

    with Session(e) as s:
        s.add_all([Widget(id=i, data=j) for i, j in enumerate(data, 1)])
        s.commit()

    with Session(e) as s:
        q = select(Widget)

        for result in windowed_query(s, q, Widget.data, 1000):
            for widget in result.scalars():
                print("data:", widget.data)