Merge pull request #69 from qaspen-python/feature/support_binary_copy

Added method binary_copy_to_table

Author: chandr-andr

Cargo.lock

@@ -28,15 +28,15 @@ uuid = { version = "1.7.0", features = ["v4"] }
 serde_json = "1.0.113"
 futures-util = "0.3.30"
 macaddr = "1.0.1"
+postgres-types = { git = "https://github.com/chandr-andr/rust-postgres.git", branch = "psqlpy", features = [
+    "derive",
+] }
 tokio-postgres = { git = "https://github.com/chandr-andr/rust-postgres.git", branch = "psqlpy", features = [
     "with-serde_json-1",
     "array-impls",
     "with-chrono-0_4",
     "with-uuid-1",
 ] }
-postgres-types = { git = "https://github.com/chandr-andr/rust-postgres.git", branch = "psqlpy", features = [
-    "derive",
-] }
 postgres-protocol = { git = "https://github.com/chandr-andr/rust-postgres.git", branch = "psqlpy" }
 postgres-openssl = { git = "https://github.com/chandr-andr/rust-postgres.git", branch = "psqlpy" }
 rust_decimal = { git = "https://github.com/chandr-andr/rust-decimal.git", branch = "psqlpy", features = [

Cargo.toml

@@ -1,9 +1,10 @@
 import types
 from enum import Enum
+from io import BytesIO
 from ipaddress import IPv4Address, IPv6Address
 from typing import Any, Callable, List, Optional, Sequence, TypeVar, Union
 
-from typing_extensions import Self
+from typing_extensions import Buffer, Self
 
 _CustomClass = TypeVar(
     "_CustomClass",
@@ -809,6 +810,30 @@ class Transaction:
             await cursor.close()
         ```
         """
+    async def binary_copy_to_table(
+        self: Self,
+        source: Union[bytes, bytearray, Buffer, BytesIO],
+        table_name: str,
+        columns: Optional[Sequence[str]] = None,
+        schema_name: Optional[str] = None,
+    ) -> int:
+        """Perform binary copy to PostgreSQL.
+
+        Execute `COPY table_name (<columns>) FROM STDIN (FORMAT binary)`
+        and start sending bytes to PostgreSQL.
+
+        IMPORTANT! User is responsible for the bytes passed to the database.
+        If bytes are incorrect user will get error from the database.
+
+        ### Parameters:
+        - `source`: source of bytes.
+        - `table_name`: name of the table.
+        - `columns`: sequence of str columns.
+        - `schema_name`: name of the schema.
+
+        ### Returns:
+        number of inserted rows;
+        """
 
 class Connection:
     """Connection from Database Connection Pool.
@@ -1053,6 +1078,30 @@ class Connection:
         It necessary to commit all transactions and close all cursor
         made by this connection. Otherwise, it won't have any practical usage.
         """
+    async def binary_copy_to_table(
+        self: Self,
+        source: Union[bytes, bytearray, Buffer, BytesIO],
+        table_name: str,
+        columns: Optional[Sequence[str]] = None,
+        schema_name: Optional[str] = None,
+    ) -> int:
+        """Perform binary copy to PostgreSQL.
+
+        Execute `COPY table_name (<columns>) FROM STDIN (FORMAT binary)`
+        and start sending bytes to PostgreSQL.
+
+        IMPORTANT! User is responsible for the bytes passed to the database.
+        If bytes are incorrect user will get error from the database.
+
+        ### Parameters:
+        - `source`: source of bytes.
+        - `table_name`: name of the table.
+        - `columns`: sequence of str columns.
+        - `schema_name`: name of the schema.
+
+        ### Returns:
+        number of inserted rows;
+        """
 
 class ConnectionPoolStatus:
     max_size: int

python/psqlpy/_internal/__init__.pyi

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import os
 import typing
+from io import BytesIO
 
 import pytest
+from pgpq import ArrowToPostgresBinaryEncoder
+from pyarrow import parquet
 from tests.helpers import count_rows_in_test_table
 
 from psqlpy import ConnectionPool, Cursor, QueryResult, Transaction
@@ -180,3 +184,55 @@ async def test_closed_connection_error(
 
     with pytest.raises(expected_exception=ConnectionClosedError):
         await connection.execute("SELECT 1")
+
+
+async def test_binary_copy_to_table(
+    psql_pool: ConnectionPool,
+) -> None:
+    """Test binary copy in connection."""
+    table_name: typing.Final = "cars"
+    await psql_pool.execute(f"DROP TABLE IF EXISTS {table_name}")
+    await psql_pool.execute(
+        """
+CREATE TABLE IF NOT EXISTS cars (
+    model VARCHAR,
+    mpg FLOAT8,
+    cyl INTEGER,
+    disp FLOAT8,
+    hp INTEGER,
+    drat FLOAT8,
+    wt FLOAT8,
+    qsec FLOAT8,
+    vs INTEGER,
+    am INTEGER,
+    gear INTEGER,
+    carb INTEGER
+);
+""",
+    )
+
+    arrow_table = parquet.read_table(
+        f"{os.path.dirname(os.path.abspath(__file__))}/test_data/MTcars.parquet",  # noqa: PTH120, PTH100
+    )
+    encoder = ArrowToPostgresBinaryEncoder(arrow_table.schema)
+    buf = BytesIO()
+    buf.write(encoder.write_header())
+    for batch in arrow_table.to_batches():
+        buf.write(encoder.write_batch(batch))
+    buf.write(encoder.finish())
+    buf.seek(0)
+
+    async with psql_pool.acquire() as connection:
+        inserted_rows = await connection.binary_copy_to_table(
+            source=buf,
+            table_name=table_name,
+        )
+
+    expected_inserted_row: typing.Final = 32
+
+    assert inserted_rows == expected_inserted_row
+
+    real_table_rows: typing.Final = await psql_pool.execute(
+        f"SELECT COUNT(*) AS rows_count FROM {table_name}",
+    )
+    assert real_table_rows.result()[0]["rows_count"] == expected_inserted_row

python/tests/test_connection.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import os
 import typing
+from io import BytesIO
 
 import pytest
+from pgpq import ArrowToPostgresBinaryEncoder
+from pyarrow import parquet
 from tests.helpers import count_rows_in_test_table
 
 from psqlpy import ConnectionPool, Cursor, IsolationLevel, ReadVariant
@@ -334,3 +338,55 @@ async def test_transaction_send_underlying_connection_to_pool_manually(
         await transaction.commit()
         assert not psql_pool.status().available
     assert psql_pool.status().available == 1
+
+
+async def test_binary_copy_to_table(
+    psql_pool: ConnectionPool,
+) -> None:
+    """Test binary copy in transaction."""
+    table_name: typing.Final = "cars"
+    await psql_pool.execute(f"DROP TABLE IF EXISTS {table_name}")
+    await psql_pool.execute(
+        """
+CREATE TABLE IF NOT EXISTS cars (
+    model VARCHAR,
+    mpg FLOAT8,
+    cyl INTEGER,
+    disp FLOAT8,
+    hp INTEGER,
+    drat FLOAT8,
+    wt FLOAT8,
+    qsec FLOAT8,
+    vs INTEGER,
+    am INTEGER,
+    gear INTEGER,
+    carb INTEGER
+);
+""",
+    )
+
+    arrow_table = parquet.read_table(
+        f"{os.path.dirname(os.path.abspath(__file__))}/test_data/MTcars.parquet",  # noqa: PTH120, PTH100
+    )
+    encoder = ArrowToPostgresBinaryEncoder(arrow_table.schema)
+    buf = BytesIO()
+    buf.write(encoder.write_header())
+    for batch in arrow_table.to_batches():
+        buf.write(encoder.write_batch(batch))
+    buf.write(encoder.finish())
+    buf.seek(0)
+
+    async with psql_pool.acquire() as connection:
+        inserted_rows = await connection.binary_copy_to_table(
+            source=buf,
+            table_name=table_name,
+        )
+
+    expected_inserted_row: typing.Final = 32
+
+    assert inserted_rows == expected_inserted_row
+
+    real_table_rows: typing.Final = await psql_pool.execute(
+        f"SELECT COUNT(*) AS rows_count FROM {table_name}",
+    )
+    assert real_table_rows.result()[0]["rows_count"] == expected_inserted_row

python/tests/test_data/MTcars.parquet

@@ -119,3 +119,22 @@ impl KeepaliveConfig {
         }
     }
 }
+
+#[pyclass]
+#[derive(Clone, Copy)]
+pub enum CopyCommandFormat {
+    TEXT,
+    CSV,
+    BINARY,
+}
+
+impl CopyCommandFormat {
+    #[must_use]
+    pub fn to_internal(&self) -> String {
+        match self {
+            CopyCommandFormat::TEXT => "text".into(),
+            CopyCommandFormat::CSV => "csv".into(),
+            CopyCommandFormat::BINARY => "binary".into(),
+        }
+    }
+}