# stdlib imports
import asyncio
import enum

# venv imports
from sqlalchemy import (
   Column,
   exc,
   insert,
   Integer,
   MetaData,
   select,
   Table,
   text,
   __version__ as sqla_version,
)
from sqlalchemy.dialects.postgresql import (
   ENUM
)
from sqlalchemy.ext.asyncio import create_async_engine

# assert the version for this test
assert sqla_version == '2.0.9'

meta = MetaData()

class AppGroup(enum.Enum):
   """Application groups"""
   ADMIN = 'admin'
   STAFF = 'staff'
   CUST = 'customer'


Group = Table(
   'mygroup',
   meta,
   Column('id', Integer, primary_key=True),
   Column('name', ENUM(AppGroup, name=AppGroup.__name__.lower()),
          nullable=False,
          unique=True),
)

DBURI = 'postgresql+asyncpg://tcp:pass4TCP!@postgres/foo'

async def main():
   engine = create_async_engine(DBURI, future=True, echo=True)

   async with engine.begin() as conn:
      await conn.run_sync(meta.create_all)

      await conn.execute(insert(Group).values([dict(name=grp) for grp in AppGroup]))

      # using the enum works with both 1.4 and 2.0
      stmt = select(Group).where(Group.c.name==AppGroup.ADMIN)
      result = await conn.execute(stmt)
      print('>>>>>>>>>>>>>> Found group ADMIN:', result.fetchall())

      # strings of course work in raw SQL
      stmt = text("select id, name from mygroup where name = 'ADMIN'")
      result = await conn.execute(stmt)
      print('>>>>>>>>>>>>>> Found group ADMIN:', result.fetchall())

      # searching by string works in 1.4, but raises a
      # ProgrammingError in 2.0, "operator does not exist: appgroup = character varying"
      try:
         stmt = select(Group).where(Group.c.name=='ADMIN')
         result = await conn.execute(stmt)
         print('>>>>>>>>>>>>>> Found group ADMIN:', result.fetchall())
      except exc.ProgrammingError as error:
         print('!!!!!!!!!!!!!!!!!!!!', error)


      await conn.rollback()

asyncio.run(main())