Aiopg result module optimization benchmark setup

benchmark.py

import asyncio
import json
import timeit

import aiopg
import aiopg.sa
import asyncpg
import sqlalchemy as sa
import sqlalchemy.dialects.postgresql


metadata = sa.MetaData()

data_simple = sa.Table(
    'data_simple',
    metadata,
   sa.Column('id', sa.Integer, primary_key=True),
   sa.Column('data', sa.Text)
)

data_json = sa.Table(
    'data_json',
    metadata,
   sa.Column('id', sa.Integer, primary_key=True),
   sa.Column('data', sa.dialects.postgresql.JSON)
)


DATABASE = {
    'username': 'postgres',
    'password': 'postgres',
    'host': 'localhost',
    'port': 5432,
    'database': 'aiopg'
}


def setup_engine():
    return aiopg.sa.create_engine(
        user=DATABASE['username'],
        password=DATABASE['password'],
        host=DATABASE['host'],
        port=DATABASE['port'],
        database=DATABASE['database']
    )


async def _init(connection):
    await connection.set_type_codec(
        'json',
        encoder=json.dumps,
        decoder=json.loads,
        schema='pg_catalog'
    )


def setup_engine_asyncpg():
    return asyncpg.create_pool(
        user=DATABASE['username'],
        password=DATABASE['password'],
        host=DATABASE['host'],
        port=DATABASE['port'],
        database=DATABASE['database'],
        init=_init
    )


async def select_simple(connection):
    data = await connection.execute(
        sa.select([data_simple])
    )
    result = await data.fetchall()
    for item in result:
        assert 'unknown' not in item
        for _ in range(10):
            data = item.data.upper()
            id = item.id ** 2


async def select_simple_asyncpg(connection):
    data = await connection.execute(
        sa.select([data_simple])
    )
    result = await data.fetchall()
    for item in result:
        for _ in range(10):
            data = item['data'].upper()
            id = item['id'] ** 2


async def select_json(connection):
    data = await connection.execute(
        sa.select([data_json])
    )
    for item in await data.fetchall():
        for _ in range(10):
            data = list(item.data.keys())
            id = item.id ** 2


async def select_simple_asyncpg(connection):
    query = str(sa.select([data_simple]))
    data = await connection.fetch(query)
    for item in data:
        for _ in range(10):
            data = item['data'].upper()
            id = item['id'] ** 2


async def select_json_asyncpg(connection):
    query = str(sa.select([data_json]))
    data = await connection.fetch(query)
    for item in data:
        for _ in range(10):
            data = list(item['data'].keys())
            id = item['id'] ** 2


async def benchmark(fn, number=99):
    async with setup_engine() as engine:
        async with engine.acquire() as connection:
            for _ in range(number):
                await fn(connection)


async def benchmark_asyncpg(fn, number=99):
    async with setup_engine_asyncpg() as pool:
        async with pool.acquire() as connection:
            for _ in range(number):
                await fn(connection)


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    results = timeit.repeat(
        lambda: loop.run_until_complete(benchmark(select_simple)),
        repeat=3,
        number=1
    )
    print('simple', results)
    results = timeit.repeat(
        lambda: loop.run_until_complete(benchmark(select_json)),
        repeat=3,
        number=1
    )
    print('json', results)
    results = timeit.repeat(
        lambda: loop.run_until_complete(benchmark_asyncpg(select_simple_asyncpg)),
        repeat=3,
        number=1
    )
    print('simple_asyncpg', results)
    results = timeit.repeat(
        lambda: loop.run_until_complete(benchmark_asyncpg(select_json_asyncpg)),
        repeat=3,
        number=1
    )
    print('json_asyncpg', results)

db.sql

CREATE TABLE data_json (
    id serial PRIMARY KEY,
    data json NULL
);

CREATE TABLE data_simple (
    id serial PRIMARY KEY,
    data text NULL
);


INSERT INTO data_simple (data)
SELECT md5(random()::text) FROM generate_series(1, 500);

INSERT INTO data_json (data)
SELECT '{"foo": 1, "bar": "text", "baz": [1, 2, 3]}'::json FROM generate_series(1, 500);

docker-compose.yml

version: '3.1'

services:
  aiopg-db:
    image: postgres
    restart: always
    volumes:
      - aiopg-db-data:/var/lib/postgresql/data
    environment:
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: postgres
    ports:
      - 5432:5432
volumes:
  aiopg-db-data: