query.py

# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.30.0
# source: query.sql
import numpy
from numpy.typing import NDArray
from typing import Optional

import sqlalchemy
import sqlalchemy.ext.asyncio

from python import models


CREATE_ITEM = """-- name: create_item \\:one
INSERT INTO items (embedding) VALUES (:p1) RETURNING id, embedding
"""


GET_ITEM = """-- name: get_item \\:one
SELECT id, embedding FROM items WHERE id = :p1
"""


class Querier:
    def __init__(self, conn: sqlalchemy.engine.Connection):
        self._conn = conn

    def create_item(self, *, embedding: Optional[NDArray[numpy.float32]]) -> Optional[models.Item]:
        row = self._conn.execute(sqlalchemy.text(CREATE_ITEM), {"p1": embedding}).first()
        if row is None:
            return None
        return models.Item(
            id=row[0],
            embedding=row[1],
        )

    def get_item(self, *, id: int) -> Optional[models.Item]:
        row = self._conn.execute(sqlalchemy.text(GET_ITEM), {"p1": id}).first()
        if row is None:
            return None
        return models.Item(
            id=row[0],
            embedding=row[1],
        )


class AsyncQuerier:
    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
        self._conn = conn

    async def create_item(self, *, embedding: Optional[NDArray[numpy.float32]]) -> Optional[models.Item]:
        row = (await self._conn.execute(sqlalchemy.text(CREATE_ITEM), {"p1": embedding})).first()
        if row is None:
            return None
        return models.Item(
            id=row[0],
            embedding=row[1],
        )

    async def get_item(self, *, id: int) -> Optional[models.Item]:
        row = (await self._conn.execute(sqlalchemy.text(GET_ITEM), {"p1": id})).first()
        if row is None:
            return None
        return models.Item(
            id=row[0],
            embedding=row[1],
        )