Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions python/psqlpy/_internal/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -942,6 +942,50 @@ class Connection:
- `read_variant`: configure read variant of the transaction.
- `deferrable`: configure deferrable of the transaction.
"""
def cursor(
self: Self,
querystring: str,
parameters: Sequence[Any] | None = None,
fetch_number: int | None = None,
scroll: bool | None = None,
prepared: bool = True,
) -> Cursor:
"""Create new cursor object.

Cursor can be used as an asynchronous iterator.

### Parameters:
- `querystring`: querystring to execute.
- `parameters`: list of parameters to pass in the query.
- `fetch_number`: how many rows need to fetch.
- `scroll`: SCROLL or NO SCROLL cursor.
- `prepared`: should the querystring be prepared before the request.
By default any querystring will be prepared.

### Returns:
new initialized cursor.

### Example:
```python
import asyncio

from psqlpy import PSQLPool, QueryResult


async def main() -> None:
db_pool = PSQLPool()
connection = await db_pool.connection()
async with connection.transaction():
async with connection.cursor(
querystring="SELECT * FROM users WHERE username = $1",
parameters=["Some_Username"],
fetch_number=5,
) as cursor:
async for fetched_result in cursor:
dict_result: List[Dict[Any, Any]] = fetched_result.result()
... # do something with this result.
```
"""

class ConnectionPool:
"""Connection pool for executing queries.
Expand Down
21 changes: 20 additions & 1 deletion python/tests/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import pytest
from tests.helpers import count_rows_in_test_table

from psqlpy import ConnectionPool, QueryResult, Transaction
from psqlpy import ConnectionPool, Cursor, QueryResult, Transaction
from psqlpy.exceptions import ConnectionExecuteError, TransactionExecuteError

pytestmark = pytest.mark.anyio
Expand Down Expand Up @@ -128,3 +128,22 @@ async def test_connection_fetch_val_more_than_one_row(
f"SELECT * FROM {table_name}",
[],
)


async def test_connection_cursor(
psql_pool: ConnectionPool,
table_name: str,
number_database_records: int,
) -> None:
"""Test cursor from Connection."""
connection = await psql_pool.connection()
cursor: Cursor
all_results: list[dict[typing.Any, typing.Any]] = []

async with connection.transaction(), connection.cursor(
querystring=f"SELECT * FROM {table_name}",
) as cursor:
async for cur_res in cursor:
all_results.extend(cur_res.result())

assert len(all_results) == number_database_records
21 changes: 21 additions & 0 deletions src/driver/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use crate::{
};

use super::{
cursor::Cursor,
transaction::Transaction,
transaction_options::{IsolationLevel, ReadVariant},
};
Expand Down Expand Up @@ -396,4 +397,24 @@ impl Connection {
HashSet::new(),
)
}

#[must_use]
pub fn cursor(
&self,
querystring: String,
parameters: Option<Py<PyAny>>,
fetch_number: Option<usize>,
scroll: Option<bool>,
prepared: Option<bool>,
) -> Cursor {
Cursor::new(
self.db_client.clone(),
querystring,
parameters,
"cur_name".into(),
fetch_number.unwrap_or(10),
scroll,
prepared,
)
}
}