Advanced Topics

Accessing the cursor object

The cursor is a temporary object created in memory that allows you to perform row-by-row operations on your data and use handy methods such as .description, .fetchall() and .fetchone(). As long as you are running a SQL SELECT query, you can access the cursor object by appending _cursor to the end of the queries name.

For example, say you have the following query named get-all-greetings in a sql file:

-- name: get_all_greetings
-- Get all the greetings in the database
select greeting_id, greeting
  from greetings
 order by 1;

With this query, you can get all greeting_id’s and greeting’s, access the cursor object, and print the column names with the following code:

import asyncio
import aiosql
import aiosqlite
from typing import List

queries = aiosql.from_path("greetings.sql", "aiosqlite")

async def access_cursor():
    async with aiosqlite.connect("greetings.db") as conn:
        # append _cursor after query name
        async with queries.get_all_greetings_cursor(conn) as cur:
            print([col_info[0] for col_info in cur.description])
            first_row = await cur.fetchone()
            all_data = await cur.fetchall()
            print(f"FIRST ROW: {first_row}")  # tuple of first row data
            print(f"OTHER DATA: {all_data}")    # remaining rows

asyncio.run(access_cursor())

# ['greeting_id', 'greeting']
# FIRST ROW: (1, 'Hi')
# OTHER DATA: [(2, 'Aloha'), (3, 'Hola'), (4, 'Bonjour'), (5, '你好')]

Accessing prepared SQL as a string

When you need to do something not directly supported by aiosql, this is your escape hatch. You can still define your SQL in a file and load it with aiosql, but then you may choose to use it without calling your aiosql method. The prepared SQL string of a method is available as an attribute of each method queries.<method_name>.sql. Here’s an example of how you might use it with a unique feature of psycopg2 like execute_values.

import aiosql
import psycopg2
from psycopg2.extras import execute_values

SQL = """
-- name: create_schema#
create table if not exists test (id int primary key, v1 int, v2 int);

-- name: insert!
INSERT INTO test (id, v1, v2) VALUES %s;

-- name: update!
UPDATE test SET v1 = data.v1 FROM (VALUES %s) AS data (id, v1)
WHERE test.id = data.id;

-- name: getem
select * from test order by id;
"""

queries = aiosql.from_str(SQL, "psycopg2")
conn = psycopg2.connect("dbname=test")
queries.create_schema(conn)
with conn.cursor() as cur:
    execute_values(cur, queries.insert.sql, [(1, 2, 3), (4, 5, 6), (7, 8, 9)])
    execute_values(cur, queries.update.sql, [(1, 20), (4, 50)])

print(list(queries.getem(conn)))
#  [(1, 20, 3), (4, 50, 6), (7, 8, 9)]

Accessing the SQL Operation Type

Query functions also provide access to the SQL operation type you define in your library. This can be useful for observability (such as metrics, tracing, or logging), or customizing how you manage different operations within your codebase. Extending from the above example:

import time
import logging
import contextlib
import pg_execute_values as pev

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("metrics")

def report_metrics(op, sql, op_time):
    log.info(f"Operation: {op.name!r}\nSQL: {sql!r} \nTime (ms): {op_time}")

@contextlib.contextmanager
def observe_query(func):
    op = func.operation
    sql = func.sql
    start = time.time()
    yield
    end = time.time()
    op_time = end - start
    report_metrics(op, sql, op_time)

with observe_query(pev.queries.getem):
    pev.queries.getem(pev.conn)

# INFO:metrics:Operation: 'SELECT'
# SQL: 'select * from test order by id;' 
# Time (ms): 2.6226043701171875e-06