Skip to content

Added method execute_batch for Connection and Transaction #84

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 19, 2024
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
22 changes: 22 additions & 0 deletions docs/components/connection.md
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,28 @@ async def main() -> None:
dict_results: list[dict[str, Any]] = results.result()
```

### Execute Batch

#### Parameters:

- `querystring`: querystrings separated by semicolons.

Executes a sequence of SQL statements using the simple query protocol.

Statements should be separated by semicolons.
If an error occurs, execution of the sequence will stop at that point.
This is intended for use when, for example,
initializing a database schema.

```python
async def main() -> None:
...
connection = await db_pool.connection()
await connection.execute_batch(
"CREATE TABLE psqlpy (name VARCHAR); CREATE TABLE psqlpy2 (name VARCHAR);",
)
```

### Fetch

#### Parameters:
Expand Down
23 changes: 23 additions & 0 deletions docs/components/transaction.md
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,29 @@ async def main() -> None:
dict_results: list[dict[str, Any]] = results.result()
```

### Execute Batch

#### Parameters:

- `querystring`: querystrings separated by semicolons.

Executes a sequence of SQL statements using the simple query protocol.

Statements should be separated by semicolons.
If an error occurs, execution of the sequence will stop at that point.
This is intended for use when, for example,
initializing a database schema.

```python
async def main() -> None:
...
connection = await db_pool.connection()
async with connection.transaction() as transaction:
await transaction.execute_batch(
"CREATE TABLE psqlpy (name VARCHAR); CREATE TABLE psqlpy2 (name VARCHAR);",
)
```

### Fetch

#### Parameters:
Expand Down
30 changes: 30 additions & 0 deletions python/psqlpy/_internal/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -455,6 +455,21 @@ class Transaction:
await transaction.commit()
```
"""
async def execute_batch(
self: Self,
querystring: str,
) -> None:
"""
Executes a sequence of SQL statements using the simple query protocol.

Statements should be separated by semicolons.
If an error occurs, execution of the sequence will stop at that point.
This is intended for use when, for example,
initializing a database schema.

### Parameters:
- `querystring`: querystrings separated by semicolons.
"""
async def execute_many(
self: Self,
querystring: str,
Expand Down Expand Up @@ -885,6 +900,21 @@ class Connection:
dict_result: List[Dict[Any, Any]] = query_result.result()
```
"""
async def execute_batch(
self: Self,
querystring: str,
) -> None:
"""
Executes a sequence of SQL statements using the simple query protocol.

Statements should be separated by semicolons.
If an error occurs, execution of the sequence will stop at that point.
This is intended for use when, for example,
initializing a database schema.

### Parameters:
- `querystring`: querystrings separated by semicolons.
"""
async def execute_many(
self: Self,
querystring: str,
Expand Down
11 changes: 11 additions & 0 deletions python/tests/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -236,3 +236,14 @@ async def test_binary_copy_to_table(
f"SELECT COUNT(*) AS rows_count FROM {table_name}",
)
assert real_table_rows.result()[0]["rows_count"] == expected_inserted_row


async def test_execute_batch_method(psql_pool: ConnectionPool) -> None:
"""Test `execute_batch` method."""
await psql_pool.execute(querystring="DROP TABLE IF EXISTS execute_batch")
await psql_pool.execute(querystring="DROP TABLE IF EXISTS execute_batch2")
query = "CREATE TABLE execute_batch (name VARCHAR);CREATE TABLE execute_batch2 (name VARCHAR);"
async with psql_pool.acquire() as conn:
await conn.execute_batch(querystring=query)
await conn.execute(querystring="SELECT * FROM execute_batch")
await conn.execute(querystring="SELECT * FROM execute_batch2")
11 changes: 11 additions & 0 deletions python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,3 +390,14 @@ async def test_binary_copy_to_table(
f"SELECT COUNT(*) AS rows_count FROM {table_name}",
)
assert real_table_rows.result()[0]["rows_count"] == expected_inserted_row


async def test_execute_batch_method(psql_pool: ConnectionPool) -> None:
"""Test `execute_batch` method."""
await psql_pool.execute(querystring="DROP TABLE IF EXISTS execute_batch")
await psql_pool.execute(querystring="DROP TABLE IF EXISTS execute_batch2")
query = "CREATE TABLE execute_batch (name VARCHAR);CREATE TABLE execute_batch2 (name VARCHAR);"
async with psql_pool.acquire() as conn, conn.transaction() as transaction:
await transaction.execute_batch(querystring=query)
await transaction.execute(querystring="SELECT * FROM execute_batch")
await transaction.execute(querystring="SELECT * FROM execute_batch2")
25 changes: 25 additions & 0 deletions src/driver/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,31 @@ impl Connection {
Err(RustPSQLDriverError::ConnectionClosedError)
}

/// Executes a sequence of SQL statements using the simple query protocol.
///
/// Statements should be separated by semicolons.
/// If an error occurs, execution of the sequence will stop at that point.
/// This is intended for use when, for example,
/// initializing a database schema.
///
/// # Errors
///
/// May return Err Result if:
/// 1) Connection is closed.
/// 2) Cannot execute querystring.
pub async fn execute_batch(
self_: pyo3::Py<Self>,
querystring: String,
) -> RustPSQLDriverPyResult<()> {
let db_client = pyo3::Python::with_gil(|gil| self_.borrow(gil).db_client.clone());

if let Some(db_client) = db_client {
return Ok(db_client.batch_execute(&querystring).await?);
}

Err(RustPSQLDriverError::ConnectionClosedError)
}

/// Execute querystring with parameters.
///
/// It converts incoming parameters to rust readable
Expand Down
25 changes: 25 additions & 0 deletions src/driver/transaction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,31 @@ impl Transaction {
Err(RustPSQLDriverError::TransactionClosedError)
}

/// Executes a sequence of SQL statements using the simple query protocol.
///
/// Statements should be separated by semicolons.
/// If an error occurs, execution of the sequence will stop at that point.
/// This is intended for use when, for example,
/// initializing a database schema.
///
/// # Errors
///
/// May return Err Result if:
/// 1) Transaction is closed.
/// 2) Cannot execute querystring.
pub async fn execute_batch(self_: Py<Self>, querystring: String) -> RustPSQLDriverPyResult<()> {
let (is_transaction_ready, db_client) = pyo3::Python::with_gil(|gil| {
let self_ = self_.borrow(gil);
(self_.check_is_transaction_ready(), self_.db_client.clone())
});
is_transaction_ready?;
if let Some(db_client) = db_client {
return Ok(db_client.batch_execute(&querystring).await?);
}

Err(RustPSQLDriverError::TransactionClosedError)
}

/// Fetch result from the database.
///
/// It converts incoming parameters to rust readable
Expand Down
Loading