from debian.debian_support import version_compare
import sqlalchemy.event

def fetchiter(cursor):
    rows = cursor.fetchmany(1024)
    while rows:
        for row in rows:
            yield row
        rows = cursor.fetchmany(1024)

def enable_sqlite_foreign_keys(engine):
    if engine.dialect.name != "sqlite":
        return
    @sqlalchemy.event.listens_for(engine, "connect")
    def pragma_foreign_keys(connection, _):
        connection.execute("PRAGMA foreign_keys=ON;")

def sqlite_add_version_compare(engine):
    if engine.dialect.name != "sqlite":
        return
    @sqlalchemy.event.listens_for(engine, "connect")
    def add_version_compare(connection, _):
        connection.create_collation("debian_version", version_compare)
        connection.create_function("debian_version_compare", 2, version_compare)

def postgres_enable_result_streaming(engine):
    if engine.dialect.driver != "psycopg2":
        return
    engine.server_side_cursors = True

def configure_database_engine(engine):
    enable_sqlite_foreign_keys(engine)
    sqlite_add_version_compare(engine)
    postgres_enable_result_streaming(engine)