Initial commit

Robert Holt · Robert Holt · commit c39ca38a6ad5 · 2021-03-03T12:12:01.000-05:00
diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml
@@ -0,0 +1,32 @@
+name: Linting
+on: [push, pull_request]
+jobs:
+  test:
+    name: linting
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.9"
+      - name: Install poetry
+        uses: snok/install-poetry@v1.1.1
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+      - name: Load cached venv
+        id: cached-poetry-dependencies
+        uses: actions/cache@v2
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ hashFiles('**/poetry.lock') }}
+      - name: Install dependencies
+        if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
+        run: poetry install --no-interaction
+      - name: Run lint checks
+        run: |
+          source .venv/bin/activate
+          black --check --diff .
+          mypy
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -0,0 +1,27 @@
+name: Release
+on:
+  release:
+    types:
+      - created
+jobs:
+  release:
+    name: release
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.9"
+      - name: Install poetry
+        uses: snok/install-poetry@v1.1.1
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+      - name: Publish release
+        env:
+          PYPI_USERNAME: ${{ secrets.PYPI_USERNAME }}
+          PYPI_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+        run: |
+          poetry publish --build --username "${PYPI_USERNAME}" --password "${PYPI_PASSWORD}"
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,7 @@
+dist/
+.idea/
+__pycache__
+*.egg-info
+.mypy_cache
+.python-version
+*.venv
diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Datto Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
@@ -0,0 +1,26 @@
+# sqlc-python-runtime
+
+A small runtime library for python code generated by [sqlc](https://github.com/kyleconroy/sqlc)
+
+## Installation
+
+Add the `sqlc-python-runtime` package to your project's dependencies.
+
+## Usage
+
+All query functions generated by sqlc take a database connection object as their first argument.
+This connection object must conform to the `Connection`/`AsyncConnection` interface defined in this package.
+This package also provides wrappers for `psycopg2` and `asyncpg` that conform to this interface.
+
+`psycopg2` example:
+```python
+import psycopg2
+from sqlc_runtime.psycopg2 import build_psycopg2_connection
+
+
+conn = build_psycopg2_connection(psycopg2.connect("postgresql://localhost:5432/mydatabase"))
+```
+
+Also note that the generated query functions can use either a `Connection` or an `AsyncConnection`.
+If an `AsyncConnection` is passed the result must be awaited. The function type hints communicate this so any type
+checkers (like mypy) will raise errors if any are missed or incorrectly used.
diff --git a/mypy.ini b/mypy.ini
@@ -0,0 +1,4 @@
+[mypy]
+ignore_missing_imports = True
+disallow_untyped_defs = True
+files = sqlc_runtime
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -0,0 +1,28 @@
+[tool.poetry]
+name = "sqlc-python-runtime"
+version = "1.0.1"
+description = "Runtime components for code generated by sqlc"
+authors = ["Robert Holt <rholt@datto.com>"]
+license = "MIT"
+homepage = "https://github.com/robholt/sqlc-python-runtime"
+readme = "README.md"
+packages = [{include = "sqlc_runtime"}]
+include = ["README.md"]
+
+[tool.poetry.dependencies]
+python = "^3.7"
+psycopg2-binary = {version = "^2.8.6", optional = true}
+asyncpg = {version = "^0.22.0", optional = true}
+pydantic = "^1.4"
+
+[tool.poetry.dev-dependencies]
+mypy = "^0.812"
+black = "^20.8b1"
+
+[tool.poetry.extras]
+psycopg2 = ["psycopg2-binary"]
+asyncpg = ["asyncpg"]
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
diff --git a/sqlc_runtime/__init__.py b/sqlc_runtime/__init__.py
@@ -0,0 +1,10 @@
+from .interface import (
+    GenericConnection,
+    Connection,
+    AsyncConnection,
+    GenericCursor,
+    Cursor,
+    AsyncCursor,
+    ReturnType,
+    IteratorReturn,
+)
diff --git a/sqlc_runtime/asyncpg.py b/sqlc_runtime/asyncpg.py
@@ -0,0 +1,51 @@
+from typing import Any, Type, Optional, AsyncIterator, TypeVar
+
+import asyncpg
+import pydantic
+
+from sqlc_runtime import AsyncCursor, AsyncConnection
+
+T = TypeVar("T", bound=pydantic.BaseModel)
+
+
+def build_asyncpg_connection(conn: asyncpg.Connection) -> AsyncConnection:
+    return AsyncpgConnection(conn)
+
+
+class AsyncpgConnection:
+    def __init__(self, conn: asyncpg.Connection):
+        self._conn = conn
+
+    async def execute(self, query: str, *params: Any) -> AsyncCursor:
+        return await self._conn.cursor(query, *params)
+
+    async def execute_none(self, query: str, *params: Any) -> None:
+        await self._conn.execute(query, *params)
+
+    async def execute_rowcount(self, query: str, *params: Any) -> int:
+        status = await self._conn.execute(query, *params)
+        return int(status.split(" ")[-1])
+
+    async def execute_one(self, query: str, *params: Any) -> Any:
+        row = await self._conn.fetchrow(query, *params)
+        return row[0] if row is not None else None
+
+    async def execute_one_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> Optional[T]:
+        row = await self._conn.fetchrow(query, *params)
+        if row is None:
+            return None
+        return model.parse_obj(row)
+
+    async def execute_many(self, query: str, *params: Any) -> AsyncIterator[Any]:
+        async with self._conn.transaction():
+            async for row in self._conn.cursor(query, *params):
+                yield row[0]
+
+    async def execute_many_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> AsyncIterator[T]:
+        async with self._conn.transaction():
+            async for row in self._conn.cursor(query, *params):
+                yield model.parse_obj(row)
diff --git a/sqlc_runtime/interface.py b/sqlc_runtime/interface.py
@@ -0,0 +1,113 @@
+from typing import (
+    Protocol,
+    Any,
+    Union,
+    Sequence,
+    Tuple,
+    Optional,
+    Mapping,
+    Iterator,
+    Type,
+    TypeVar,
+    AsyncIterator,
+    Awaitable,
+)
+
+import pydantic
+
+
+class Cursor(Protocol):
+    rowcount: int
+
+    def fetchone(self) -> Mapping[str, Any]:
+        pass
+
+    def fetchmany(self, size: int = None) -> Sequence[Mapping[str, Any]]:
+        pass
+
+    def fetchall(self) -> Sequence[Mapping[str, Any]]:
+        pass
+
+    def __iter__(self) -> Iterator[Mapping[str, Any]]:
+        pass
+
+    def __len__(self) -> int:
+        pass
+
+
+T = TypeVar("T", bound=pydantic.BaseModel)
+
+
+class Connection(Protocol):
+    def execute(self, query: str, *params: Any) -> Cursor:
+        pass
+
+    def execute_none(self, query: str, *params: Any) -> None:
+        pass
+
+    def execute_rowcount(self, query: str, *params: Any) -> int:
+        pass
+
+    def execute_one(self, query: str, *params: Any) -> Any:
+        pass
+
+    def execute_one_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> Optional[T]:
+        pass
+
+    def execute_many(self, query: str, *params: Any) -> Iterator[Any]:
+        pass
+
+    def execute_many_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> Iterator[T]:
+        pass
+
+
+class AsyncCursor(Protocol):
+    async def fetch(
+        self, n: int, *, timeout: float = None
+    ) -> Sequence[Mapping[str, Any]]:
+        pass
+
+    async def fetchrow(self, *, timeout: float = None) -> Mapping[str, Any]:
+        pass
+
+    async def forward(self, n: int, *, timeout: float = None) -> int:
+        pass
+
+
+class AsyncConnection(Protocol):
+    async def execute(self, query: str, *params: Any) -> AsyncCursor:
+        pass
+
+    async def execute_none(self, query: str, *params: Any) -> None:
+        pass
+
+    async def execute_rowcount(self, query: str, *params: Any) -> int:
+        pass
+
+    async def execute_one(self, query: str, *params: Any) -> Any:
+        pass
+
+    async def execute_one_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> Optional[T]:
+        pass
+
+    def execute_many(self, query: str, *params: Any) -> AsyncIterator[Any]:
+        pass
+
+    def execute_many_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> AsyncIterator[T]:
+        pass
+
+
+GenericConnection = Union[Connection, AsyncConnection]
+GenericCursor = Union[Cursor, AsyncCursor]
+
+RT = TypeVar("RT")
+ReturnType = Union[RT, Awaitable[RT]]
+IteratorReturn = Union[Iterator[RT], AsyncIterator[RT]]
diff --git a/sqlc_runtime/psycopg2.py b/sqlc_runtime/psycopg2.py
@@ -0,0 +1,60 @@
+import re
+from typing import Any, Type, Iterator, TypeVar, Optional, TYPE_CHECKING
+
+import pydantic
+from psycopg2.extensions import connection
+from psycopg2.extras import DictCursor
+
+from sqlc_runtime import Connection
+
+T = TypeVar("T", bound=pydantic.BaseModel)
+PSYCOPG2_PLACEHOLDER_REGEXP = re.compile(r"\B\$\d+\b")
+
+
+def build_psycopg2_connection(conn: connection) -> Connection:
+    return Psycopg2Connection(conn)
+
+
+class Psycopg2Connection:
+    def __init__(self, conn: connection):
+        self._conn = conn
+
+    def execute(self, query: str, *params: Any) -> DictCursor:
+        query = PSYCOPG2_PLACEHOLDER_REGEXP.sub("%s", query)
+        cur = self._conn.cursor(cursor_factory=DictCursor)
+        cur.execute(query, params)
+        return cur
+
+    def execute_none(self, query: str, *params: Any) -> None:
+        with self.execute(query, *params):
+            return
+
+    def execute_rowcount(self, query: str, *params: Any) -> int:
+        with self.execute(query, *params) as cur:
+            return cur.rowcount
+
+    def execute_one(self, query: str, *params: Any) -> Any:
+        with self.execute(query, *params) as cur:
+            row = cur.fetchone()
+            return row[0] if row is not None else None
+
+    def execute_one_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> Optional[T]:
+        with self.execute(query, *params) as cur:
+            row = cur.fetchone()
+            if row is None:
+                return None
+            return model.parse_obj(row)
+
+    def execute_many(self, query: str, *params: Any) -> Iterator[Any]:
+        with self.execute(query, *params) as cur:
+            for row in cur:
+                yield row[0]
+
+    def execute_many_model(
+        self, model: Type[T], query: str, *params: Any
+    ) -> Iterator[T]:
+        with self.execute(query, *params) as cur:
+            for row in cur:
+                yield model.parse_obj(row)
diff --git a/sqlc_runtime/py.typed b/sqlc_runtime/py.typed
