binary representation of numpy.complex and postgresql composite #1060

from __future__ import annotations

import asyncpg
import asyncio
import numpy as np

from ast import literal_eval
import struct

SQL_CREATE = """
DO $$ BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'complex') THEN
        CREATE TYPE complex AS (
            r float4,
            i float4
        );
    END IF;
    CREATE TABLE dummy_table (dummy_column complex);
    DROP TABLE dummy_table;
END $$;

DROP TABLE IF EXISTS "poc_asyncpg";

CREATE TABLE "poc_asyncpg" (
    "id" SERIAL PRIMARY KEY,
    "float_value" float4 NULL,
    "complex_value" complex NULL,
    "complex_array" complex[] NULL
)
"""


def _cplx_decode(val) -> np.complex64:
    cplx = complex(*literal_eval(val))
    return np.complex64(cplx)


def _cplx_encode(val: np.complex64 | complex) -> str:
    return str((np.float32(val.real), np.float32(val.imag)))


async def set_type_codec(conn):
    """
    had to use this patch since the conn.set_type_codec does not work for scalar variables
    """
    schema = 'public'
    format = 'text'
    conn._check_open()
    typenames = ('complex',)
    for typename in typenames:
        typeinfo = await conn.fetchrow(
            asyncpg.introspection.TYPE_BY_NAME, typename, schema)
        if not typeinfo:
            raise ValueError('unknown type: {}.{}'.format(schema, typename))

        oid = typeinfo['oid']
        conn._protocol.get_settings().add_python_codec(
            oid, typename, schema, 'scalar',
            lambda a: _cplx_encode(a), lambda a: _cplx_decode(a), format)

# if this part is commented, error message is : 
# asyncpg.exceptions._base.InternalClientError: no binary format encoder for type complex    
        conn._protocol.get_settings().add_python_codec(
            oid, typename, schema, 'scalar',
            encoder=lambda x: struct.pack('!2f', x.real, x.imag),
            decoder=lambda x: np.frombuffer(x, dtype=np.complex64)[0],
            format="binary")

    # Statement cache is no longer valid due to codec changes.
    conn._drop_local_statement_cache()


async def init_connection(conn):
    await set_type_codec(conn)
    await conn.set_type_codec(
        'numeric', encoder=str, decoder=np.float32,
        schema='pg_catalog', format='text'
    )
    await conn.set_type_codec(
        'float4', encoder=str, decoder=np.float32,
        schema='pg_catalog', format='text'
    )
    await conn.set_type_codec(
        'float4', encoder=struct.Struct("!f").pack, decoder=struct.Struct("!f").unpack,
        schema='pg_catalog', format='binary'
    )


async def trunc(pool):
    async with pool.acquire() as conn:
        async with conn.transaction():
            query = "TRUNCATE poc_asyncpg"
            await conn.execute(query)


async def worker_copy(pool, column_name, data):
    await trunc(pool)
    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.copy_records_to_table("poc_asyncpg",
                                             records=[(data,)],
                                             columns=(column_name,)
                                             )


try:
    from common import dbinfo
except ImportError as e:
    class DB:
        user = "user"
        password = "password"
        database = "db"


    dbinfo = DB()


def create_pool(db_info):
    pool = asyncpg.create_pool(
        user=db_info.user,
        password=db_info.password,
        database=dbinfo.database,
        host="127.0.0.1",
        init=init_connection
    )
    return pool


async def main(info):
    pool = await create_pool(info)
    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.execute(SQL_CREATE)

    await worker_copy(pool, "float_value", np.float32(4.2))
    await worker_copy(pool, "complex_value", np.complex64(4.2 + 1j * 4.2))


if __name__ == '__main__':
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(main(dbinfo))