Compare different ways to have unique columns with nulls under SQLite and PostgreSQL

unique_null.py

"""Compare ways to have unique columns with NULLs."""

import os
import subprocess
import time
import uuid

import sqlalchemy as sa
import sqlalchemy.orm

DB = 'sqlite:///unique_null.sqlite3'
#DB = 'postgresql://postgres@/unique_null'

ENGINE = sa.create_engine(DB, echo=True)

REGISTRY = sa.orm.registry()


class Model:

    @sa.orm.declared_attr
    def __tablename__(cls) -> str:
        return cls.__name__.lower()

    id = sa.Column(sa.Integer, primary_key=True)

    spam = sa.Column(sa.Text, nullable=False)

    eggs = None

    @staticmethod
    def iterinsert(eggs, needle: str = 'spam', nhay: int = 500_000):
        quarter, rest = divmod(nhay, 4)
        for _ in range(quarter * 3):
            yield uuid.uuid4().hex, eggs
        yield needle, needle
        for _ in range(quarter + rest):
            yield uuid.uuid4().hex, eggs

    _eggs = None

    @classmethod
    def insert(cls, *, bind=None) -> None:
        params = cls.iterinsert(eggs=cls._eggs)
        insert = sa.insert(cls)
        if bind.dialect.name == 'sqlite':
            insert = insert.compile(bind=bind, column_keys=['spam', 'eggs'])
            assert insert.positional
            bind.connection.executemany(insert.string, params)
        else:
            bind.execute(insert, [{'spam': s, 'eggs': e} for s, e in params])

    @classmethod
    def query(cls, *, needle: str = 'spam') -> sa.sql.Select:
        return (sa.select(cls.id)
                .where(cls.spam == sa.bindparam('spam', needle)))


@REGISTRY.mapped
class Noindex(Model):

    eggs = sa.Column(sa.Text, nullable=False)

    _eggs = ''


@REGISTRY.mapped
class Index(Model):

    eggs = sa.Column(sa.Text, nullable=False)

    __table_args__ = (
        sa.UniqueConstraint('spam', eggs),
    )

    _eggs = ''


@REGISTRY.mapped
class Nullable(Model):

    eggs = sa.Column(sa.Text)

    __table_args__ = (
        sa.UniqueConstraint('spam', eggs),
    )


@REGISTRY.mapped
class Partial(Model):

    eggs = sa.Column(sa.Text)

    __table_args__ = (
        sa.Index('i1', 'spam', eggs, sqlite_where=(eggs != None),
                 postgresql_where=(eggs != None), unique=True),
        sa.Index('i2', 'spam', sqlite_where=(eggs != None),
                 postgresql_where=(eggs == None), unique=True),
    )


@REGISTRY.mapped
class Coalesce(Model):

    eggs = sa.Column(sa.Text)

    __table_args__ = (
        sa.Index('i3', 'spam', sa.func.coalesce(eggs, ''), unique=True),
    )


models = Model.__subclasses__()


def create(engine=ENGINE) -> None:
    if engine.dialect.name == 'sqlite':
        if os.path.exists(engine.url.database):
            os.remove(engine.url.database)
    elif engine.dialect.name == 'postgresql':
        user, db = engine.url.username, engine.url.database
        subprocess.run(['dropdb', '-U', user, '--if-exists', db], check=True)
        subprocess.run(['createdb', '-U', user, db], check=True)
    else:
        raise NotImplementedError

    REGISTRY.metadata.create_all(engine)

    with engine.begin() as conn:
        for m in models:
            m.insert(bind=conn)


def query(engine=ENGINE) -> None:
    for m in models:
        print(m)
        query = m.query()
        start = time.perf_counter_ns()
        with engine.connect() as conn:
            print(conn.scalar(query))
        duration = (time.perf_counter_ns() - start) / 1_000_000_000
        print(duration, '', sep='\n')


if __name__ == '__main__':
    create()
    query()