Make statements classes implementing a contract

.gitignore

@@ -21,3 +21,4 @@
 
 # Ignore noise:
 .DS_Store
+*.pyc

src/music_catalogue/__init__.py

@@ -11,65 +11,68 @@
 from ._generated._core import (
     NoRowError as NoRowError,
 )
+from ._generated._core import (
+    Statement as Statement,
+)
 from ._generated._register import (
     register_types as register_types,
 )
 from ._generated.statements.insert_album import (
-    InsertAlbumRow as InsertAlbumRow,
+    InsertAlbum as InsertAlbum,
 )
 from ._generated.statements.insert_album import (
-    insert_album as insert_album,
+    InsertAlbumRow as InsertAlbumRow,
 )
 from ._generated.statements.insert_multiple_albums import (
-    InsertMultipleAlbumsRow as InsertMultipleAlbumsRow,
+    InsertMultipleAlbums as InsertMultipleAlbums,
 )
 from ._generated.statements.insert_multiple_albums import (
-    insert_multiple_albums as insert_multiple_albums,
+    InsertMultipleAlbumsRow as InsertMultipleAlbumsRow,
 )
 from ._generated.statements.select_album_by_format import (
-    SelectAlbumByFormatRow as SelectAlbumByFormatRow,
+    SelectAlbumByFormat as SelectAlbumByFormat,
 )
 from ._generated.statements.select_album_by_format import (
-    select_album_by_format as select_album_by_format,
+    SelectAlbumByFormatRow as SelectAlbumByFormatRow,
 )
 from ._generated.statements.select_album_by_id import (
-    SelectAlbumByIdRow as SelectAlbumByIdRow,
+    SelectAlbumById as SelectAlbumById,
 )
 from ._generated.statements.select_album_by_id import (
-    select_album_by_id as select_album_by_id,
+    SelectAlbumByIdRow as SelectAlbumByIdRow,
 )
 from ._generated.statements.select_album_by_name import (
-    SelectAlbumByNameRow as SelectAlbumByNameRow,
+    SelectAlbumByName as SelectAlbumByName,
 )
 from ._generated.statements.select_album_by_name import (
-    select_album_by_name as select_album_by_name,
+    SelectAlbumByNameRow as SelectAlbumByNameRow,
 )
 from ._generated.statements.select_album_with_filters import (
-    SelectAlbumWithFiltersRow as SelectAlbumWithFiltersRow,
+    SelectAlbumWithFilters as SelectAlbumWithFilters,
 )
 from ._generated.statements.select_album_with_filters import (
-    select_album_with_filters as select_album_with_filters,
+    SelectAlbumWithFiltersRow as SelectAlbumWithFiltersRow,
 )
 from ._generated.statements.select_album_with_tracks import (
-    SelectAlbumWithTracksRow as SelectAlbumWithTracksRow,
+    SelectAlbumWithTracks as SelectAlbumWithTracks,
 )
 from ._generated.statements.select_album_with_tracks import (
-    select_album_with_tracks as select_album_with_tracks,
+    SelectAlbumWithTracksRow as SelectAlbumWithTracksRow,
 )
 from ._generated.statements.select_genre_by_artist import (
-    SelectGenreByArtistRow as SelectGenreByArtistRow,
+    SelectGenreByArtist as SelectGenreByArtist,
 )
 from ._generated.statements.select_genre_by_artist import (
-    select_genre_by_artist as select_genre_by_artist,
+    SelectGenreByArtistRow as SelectGenreByArtistRow,
 )
 from ._generated.statements.update_album_recording_returning import (
-    UpdateAlbumRecordingReturningRow as UpdateAlbumRecordingReturningRow,
+    UpdateAlbumRecordingReturning as UpdateAlbumRecordingReturning,
 )
 from ._generated.statements.update_album_recording_returning import (
-    update_album_recording_returning as update_album_recording_returning,
+    UpdateAlbumRecordingReturningRow as UpdateAlbumRecordingReturningRow,
 )
 from ._generated.statements.update_album_released import (
-    update_album_released as update_album_released,
+    UpdateAlbumReleased as UpdateAlbumReleased,
 )
 from ._generated.types.album_format import (
     AlbumFormat as AlbumFormat,
@@ -87,6 +90,7 @@
 __all__ = [
     "JsonValue",
     "NoRowError",
+    "Statement",
 ]
 __all__ += [
     "AlbumFormat",
@@ -106,16 +110,16 @@
     "UpdateAlbumRecordingReturningRow",
 ]
 __all__ += [
-    "insert_album",
-    "insert_multiple_albums",
-    "select_album_by_format",
-    "select_album_by_id",
-    "select_album_by_name",
-    "select_album_with_filters",
-    "select_album_with_tracks",
-    "select_genre_by_artist",
-    "update_album_recording_returning",
-    "update_album_released",
+    "InsertAlbum",
+    "InsertMultipleAlbums",
+    "SelectAlbumByFormat",
+    "SelectAlbumById",
+    "SelectAlbumByName",
+    "SelectAlbumWithFilters",
+    "SelectAlbumWithTracks",
+    "SelectGenreByArtist",
+    "UpdateAlbumRecordingReturning",
+    "UpdateAlbumReleased",
 ]
 __all__ += [
     "register_types",

src/music_catalogue/_generated/_core.py

@@ -6,6 +6,10 @@
 
 from __future__ import annotations
 
+from typing import ClassVar, Protocol
+
+from psycopg import AsyncConnection, Connection
+
 type JsonValue = None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"]
 
 
@@ -15,3 +19,13 @@ class NoRowError(RuntimeError):
     def __init__(self, sql: str) -> None:
         self.sql = sql
         super().__init__(f"single-row query returned no rows: {sql}")
+
+
+class Statement[R](Protocol):
+    # SQL is exposed on the class itself so generic concerns (logging, metrics, tracing)
+    # can read it off any statement without executing it.
+    SQL: ClassVar[str]
+
+    async def execute(self, conn: AsyncConnection[object]) -> R: ...
+
+    def execute_sync(self, conn: Connection[object]) -> R: ...

src/music_catalogue/_generated/statements/insert_album.py

@@ -6,55 +6,49 @@
 
 from dataclasses import dataclass
 from datetime import date
+from typing import ClassVar
 
 from psycopg import AsyncConnection, Connection
 from psycopg.rows import args_row as _args_row
 
 from .. import types as _db_types
+from .._core import Statement
 from .._runtime import fetch_single as _fetch_single
 from ..sync._runtime import fetch_single as _fetch_single_sync
 
-SQL = """\
-insert into album (name, released, format, recording)
-values (%(name)s, %(released)s, %(format)s, %(recording)s)
-returning id
-"""
-
 
 @dataclass(frozen=True, slots=True)
 class InsertAlbumRow:
     id: int
 
 
-async def insert_album(
-    conn: AsyncConnection[object],
-    *,
-    name: str,
-    released: date,
-    format: _db_types.AlbumFormat,
-    recording: _db_types.RecordingInfo,
-) -> InsertAlbumRow:
-    params: dict[str, object] = {
-        "name": name,
-        "released": released,
-        "format": format,
-        "recording": recording,
-    }
-    return await _fetch_single(conn, SQL, params, _args_row(InsertAlbumRow))
-
-
-def insert_album_sync(
-    conn: Connection[object],
-    *,
-    name: str,
-    released: date,
-    format: _db_types.AlbumFormat,
-    recording: _db_types.RecordingInfo,
-) -> InsertAlbumRow:
-    params: dict[str, object] = {
-        "name": name,
-        "released": released,
-        "format": format,
-        "recording": recording,
-    }
-    return _fetch_single_sync(conn, SQL, params, _args_row(InsertAlbumRow))
+@dataclass(frozen=True, slots=True)
+class InsertAlbum(Statement[InsertAlbumRow]):
+    name: str
+    released: date
+    format: _db_types.AlbumFormat
+    recording: _db_types.RecordingInfo
+
+    SQL: ClassVar[str] = """\
+insert into album (name, released, format, recording)
+values (%(name)s, %(released)s, %(format)s, %(recording)s)
+returning id
+"""
+
+    async def execute(self, conn: AsyncConnection[object]) -> InsertAlbumRow:
+        params: dict[str, object] = {
+            "name": self.name,
+            "released": self.released,
+            "format": self.format,
+            "recording": self.recording,
+        }
+        return await _fetch_single(conn, self.SQL, params, _args_row(InsertAlbumRow))
+
+    def execute_sync(self, conn: Connection[object]) -> InsertAlbumRow:
+        params: dict[str, object] = {
+            "name": self.name,
+            "released": self.released,
+            "format": self.format,
+            "recording": self.recording,
+        }
+        return _fetch_single_sync(conn, self.SQL, params, _args_row(InsertAlbumRow))

src/music_catalogue/_generated/statements/insert_multiple_albums.py

@@ -6,15 +6,29 @@
 
 from dataclasses import dataclass
 from datetime import date
+from typing import ClassVar
 
 from psycopg import AsyncConnection, Connection
 from psycopg.rows import args_row as _args_row
 
 from .. import types as _db_types
+from .._core import Statement
 from .._runtime import fetch_many as _fetch_many
 from ..sync._runtime import fetch_many as _fetch_many_sync
 
-SQL = """\
+
+@dataclass(frozen=True, slots=True)
+class InsertMultipleAlbumsRow:
+    id: int
+
+
+@dataclass(frozen=True, slots=True)
+class InsertMultipleAlbums(Statement[list[InsertMultipleAlbumsRow]]):
+    name: list[str]
+    released: list[date]
+    format: list[_db_types.AlbumFormat]
+
+    SQL: ClassVar[str] = """\
 -- This is an example of a bulk-insert (batch-insert) technique.
 -- We pass in all fields as arrays of the same size, and we unnest it to insert multiple rows at once.
 insert into album (name, released, format)
@@ -27,37 +41,18 @@
 returning id
 """
 
-
-@dataclass(frozen=True, slots=True)
-class InsertMultipleAlbumsRow:
-    id: int
-
-
-async def insert_multiple_albums(
-    conn: AsyncConnection[object],
-    *,
-    name: list[str],
-    released: list[date],
-    format: list[_db_types.AlbumFormat],
-) -> list[InsertMultipleAlbumsRow]:
-    params: dict[str, object] = {
-        "name": name,
-        "released": released,
-        "format": format,
-    }
-    return await _fetch_many(conn, SQL, params, _args_row(InsertMultipleAlbumsRow))
-
-
-def insert_multiple_albums_sync(
-    conn: Connection[object],
-    *,
-    name: list[str],
-    released: list[date],
-    format: list[_db_types.AlbumFormat],
-) -> list[InsertMultipleAlbumsRow]:
-    params: dict[str, object] = {
-        "name": name,
-        "released": released,
-        "format": format,
-    }
-    return _fetch_many_sync(conn, SQL, params, _args_row(InsertMultipleAlbumsRow))
+    async def execute(self, conn: AsyncConnection[object]) -> list[InsertMultipleAlbumsRow]:
+        params: dict[str, object] = {
+            "name": self.name,
+            "released": self.released,
+            "format": self.format,
+        }
+        return await _fetch_many(conn, self.SQL, params, _args_row(InsertMultipleAlbumsRow))
+
+    def execute_sync(self, conn: Connection[object]) -> list[InsertMultipleAlbumsRow]:
+        params: dict[str, object] = {
+            "name": self.name,
+            "released": self.released,
+            "format": self.format,
+        }
+        return _fetch_many_sync(conn, self.SQL, params, _args_row(InsertMultipleAlbumsRow))