Merge pull request #47 from nackjicholson/select-value-op

add select_value operation

aiosql/adapters/aiosqlite.py

@@ -37,6 +37,12 @@ async def select_one(conn, _query_name, sql, parameters, record_class=None):
                 result = record_class(**dict(zip(column_names, result)))
         return result
 
+    @staticmethod
+    async def select_value(conn, _query_name, sql, parameters):
+        async with conn.execute(sql, parameters) as cur:
+            result = await cur.fetchone()
+        return result[0] if result else None
+
     @staticmethod
     @aiocontextmanager
     async def select_cursor(conn, _query_name, sql, parameters):

aiosql/adapters/asyncpg.py

@@ -86,6 +86,11 @@ async def select_one(self, conn, query_name, sql, parameters, record_class=None)
                 result = record_class(**dict(result))
         return result
 
+    async def select_value(self, conn, query_name, sql, parameters):
+        parameters = self.maybe_order_params(query_name, parameters)
+        async with MaybeAcquire(conn) as connection:
+            return await connection.fetchval(sql, *parameters)
+
     @aiocontextmanager
     async def select_cursor(self, conn, query_name, sql, parameters):
         parameters = self.maybe_order_params(query_name, parameters)

aiosql/adapters/psycopg2.py

@@ -38,6 +38,13 @@ def select_one(conn, _query_name, sql, parameters, record_class=None):
                 result = record_class(**dict(zip(column_names, result)))
         return result
 
+    @staticmethod
+    def select_value(conn, _query_name, sql, parameters):
+        with conn.cursor() as cur:
+            cur.execute(sql, parameters)
+            result = cur.fetchone()
+        return result[0] if result else None
+
     @staticmethod
     @contextmanager
     def select_cursor(conn, _query_name, sql, parameters):

aiosql/adapters/sqlite3.py

@@ -40,6 +40,14 @@ def select_one(conn, _query_name, sql, parameters, record_class=None):
         cur.close()
         return result
 
+    @staticmethod
+    def select_value(conn, _query_name, sql, parameters):
+        cur = conn.cursor()
+        cur.execute(sql, parameters)
+        result = cur.fetchone()
+        cur.close()
+        return result[0] if result else None
+
     @staticmethod
     @contextmanager
     def select_cursor(conn, _query_name, sql, parameters):

aiosql/queries.py

@@ -53,6 +53,11 @@ def fn(self: Queries, conn, *args, **kwargs):
                 conn, query_name, sql, _params(args, kwargs), record_class
             )
 
+    elif operation_type == SQLOperationType.SELECT_VALUE:
+
+        def fn(self: Queries, conn, *args, **kwargs):
+            return self.driver_adapter.select_value(conn, query_name, sql, _params(args, kwargs))
+
     else:
         raise ValueError(f"Unknown operation_type: {operation_type}")
 

aiosql/query_loader.py

@@ -36,6 +36,9 @@ def _make_query_datum(self, query_str: str):
         elif query_name.endswith("^"):
             operation_type = SQLOperationType.SELECT_ONE
             query_name = query_name[:-1]
+        elif query_name.endswith("$"):
+            operation_type = SQLOperationType.SELECT_VALUE
+            query_name = query_name[:-1]
         else:
             operation_type = SQLOperationType.SELECT
 

aiosql/types.py

@@ -24,6 +24,7 @@ class SQLOperationType(Enum):
     SCRIPT = 3
     SELECT = 4
     SELECT_ONE = 5
+    SELECT_VALUE = 6
 
 
 class QueryDatum(NamedTuple):
@@ -63,6 +64,11 @@ def select_one(
     ) -> Optional[Any]:
         ...
 
+    def select_value(
+        self, conn: Any, query_name: str, sql: str, parameters: Union[List, Dict],
+    ) -> Optional[Any]:
+        ...
+
     def select_cursor(
         self, conn: Any, query_name: str, sql: str, parameters: Union[List, Dict]
     ) -> ContextManager[Any]:
@@ -113,6 +119,11 @@ async def select_one(
     ) -> Optional[Any]:
         ...
 
+    async def select_value(
+        self, conn: Any, query_name: str, sql: str, parameters: Union[List, Dict],
+    ) -> Optional[Any]:
+        ...
+
     async def select_cursor(
         self, conn: Any, query_name: str, sql: str, parameters: Union[List, Dict]
     ) -> AsyncContextManager[Any]:

docs/defining-sql-queries.md

@@ -72,6 +72,22 @@ queries.get_user_by_username(conn, username="willvaughn")
 # => (1, "willvaughn", "William Vaughn") or None
 ```
 
+### `$` Select Value
+
+The `$` operator will execute the query, and only return the first value of the first row of a result set. If there are no rows in the result set it returns `None`. This is implemented by returing the first element of the tuple returned by `cur.fetchone()` of the underlying driver. This is mostly useful for queries returning IDs, COUNTs or other aggregates.
+
+```sql
+-- name: get-count$
+select count(*) from users
+```
+
+When used from Python:
+
+```python
+queries.get_count(conn)
+# => 3
+```
+
 ### `!` Insert/Update/Delete
 
 The `!` operator executes SQL without returning any results. It is meant for statements that use `insert`, `update`, and `delete` to make modifications to database rows without a necessary return value.

tests/blogdb/sql/users/users.sql

@@ -21,6 +21,12 @@ select userid,
    where lastname = :lastname
 order by username asc;
 
+
 -- name: get-all-sorted
 -- Get all user records sorted by username
 select * from users order by username asc;
+
+
+-- name: get-count$
+-- Get number of users
+select count(*) from users;

tests/test_aiosqlite.py

@@ -100,6 +100,15 @@ async def test_select_one(sqlite3_db_path, queries):
     assert actual == expected
 
 
+@pytest.mark.asyncio
+async def test_select_value(sqlite3_db_path, queries):
+    async with aiosqlite.connect(sqlite3_db_path) as conn:
+        actual = await queries.users.get_count(conn)
+
+    expected = 3
+    assert actual == expected
+
+
 @pytest.mark.asyncio
 async def test_insert_returning(sqlite3_db_path, queries):
     async with aiosqlite.connect(sqlite3_db_path) as conn:

tests/test_asyncpg.py

@@ -98,6 +98,14 @@ async def test_select_one(pg_dsn, queries):
     assert actual == expected
 
 
+@pytest.mark.asyncio
+async def test_select_value(pg_dsn, queries):
+    conn = await asyncpg.connect(pg_dsn)
+    actual = await queries.users.get_count(conn)
+    expected = 3
+    assert actual == expected
+
+
 @pytest.mark.asyncio
 async def test_insert_returning(pg_dsn, queries):
     async with asyncpg.create_pool(pg_dsn) as pool:

tests/test_psycopg2.py

@@ -80,6 +80,12 @@ def test_select_one(pg_conn, queries):
     assert actual == expected
 
 
+def test_select_value(pg_conn, queries):
+    actual = queries.users.get_count(pg_conn)
+    expected = 3
+    assert actual == expected
+
+
 def test_insert_returning(pg_conn, queries):
     with pg_conn:
         blogid, title = queries.blogs.pg_publish_blog(

tests/test_sqlite3.py

@@ -81,6 +81,12 @@ def test_select_one(sqlite3_conn, queries):
     assert actual == expected
 
 
+def test_select_value(sqlite3_conn, queries):
+    actual = queries.users.get_count(sqlite3_conn)
+    expected = 3
+    assert actual == expected
+
+
 def test_insert_returning(sqlite3_conn, queries):
     with sqlite3_conn:
         blogid = queries.blogs.publish_blog(