yjdata 0.2.2__tar.gz → 0.2.3__tar.gz

{yjdata-0.2.2 → yjdata-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yjdata
-Version: 0.2.2
+Version: 0.2.3
 Summary: yjdata SDK for Python
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -98,6 +98,22 @@ the same query structure with double-quoted identifiers. Expressions retain
 normal SQL NULL semantics: for example, `change_type <> 'conceptNegation'`
 also excludes NULL `change_type` values.
 
+## Batch writes
+
+PostgreSQL inserts, upserts and updates use `psycopg2.extras.execute_batch`.
+MySQL inserts and upserts use PyMySQL's optimized `executemany`. MySQL updates
+combine driver-escaped UPDATE statements into one request, preserving input
+order, repeated keys and ordinary `=` matching (NULL keys do not match).
+MySQL connections enable `CLIENT.MULTI_STATEMENTS` for these update batches.
+
+`batch_size` limits rows per batch. MySQL update batches also respect the
+cursor's `max_stmt_length` in encoded bytes (normally about 1 MB); an individual
+larger statement is sent alone and remains subject to the server's packet limit.
+Every result is consumed so errors in later statements propagate before commit.
+With the default transactional connection settings and transactional tables,
+all batches commit together or roll back on failure. This reduces network
+round trips; the server still executes one UPDATE per input row.
+
 ## Tests
 
 ```sh
@@ -109,3 +125,6 @@ parameter ordering, and result conversion. Generated queries execute against
 SQLite's window engine with DB-API placeholders translated, so tests require
 neither optional database drivers nor running servers. This does not replace
 integration testing against MySQL and PostgreSQL.
+
+With `pymysql` installed, the MySQL batch-update tests also exercise the real
+driver's escaping, encoded request sizes and result handling without a server.

{yjdata-0.2.2 → yjdata-0.2.3}/README.md

@@ -87,6 +87,22 @@ the same query structure with double-quoted identifiers. Expressions retain
 normal SQL NULL semantics: for example, `change_type <> 'conceptNegation'`
 also excludes NULL `change_type` values.
 
+## Batch writes
+
+PostgreSQL inserts, upserts and updates use `psycopg2.extras.execute_batch`.
+MySQL inserts and upserts use PyMySQL's optimized `executemany`. MySQL updates
+combine driver-escaped UPDATE statements into one request, preserving input
+order, repeated keys and ordinary `=` matching (NULL keys do not match).
+MySQL connections enable `CLIENT.MULTI_STATEMENTS` for these update batches.
+
+`batch_size` limits rows per batch. MySQL update batches also respect the
+cursor's `max_stmt_length` in encoded bytes (normally about 1 MB); an individual
+larger statement is sent alone and remains subject to the server's packet limit.
+Every result is consumed so errors in later statements propagate before commit.
+With the default transactional connection settings and transactional tables,
+all batches commit together or roll back on failure. This reduces network
+round trips; the server still executes one UPDATE per input row.
+
 ## Tests
 
 ```sh
@@ -98,3 +114,6 @@ parameter ordering, and result conversion. Generated queries execute against
 SQLite's window engine with DB-API placeholders translated, so tests require
 neither optional database drivers nor running servers. This does not replace
 integration testing against MySQL and PostgreSQL.
+
+With `pymysql` installed, the MySQL batch-update tests also exercise the real
+driver's escaping, encoded request sizes and result handling without a server.

{yjdata-0.2.2 → yjdata-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "yjdata"
-version = "0.2.2"
+version = "0.2.3"
 description = "yjdata SDK for Python"
 readme = "README.md"
 requires-python = ">=3.10"

{yjdata-0.2.2 → yjdata-0.2.3}/src/yjdata/db/mysql.py

@@ -51,7 +51,49 @@ class MysqlAdapter(SqlAdapter):
         return upsert_qry
 
     def _create_connection(self):
-        return self._driver.connect(**self._conn_kwargs)
+        kwargs = self._conn_kwargs.copy()
+        # UPDATE batches send several statements in one request, in row order.
+        kwargs["client_flag"] = (
+            kwargs.get("client_flag", 0) | self._driver.constants.CLIENT.MULTI_STATEMENTS
+        )
+        return self._driver.connect(**kwargs)
+
+    def _execute_batch(
+        self,
+        cursor: Any,
+        qry: Any,
+        data: list,
+        batch_size: int,
+    ) -> None:
+        if not isinstance(qry, str) or not qry.startswith("UPDATE "):
+            # PyMySQL already combines INSERT/UPSERT values into bulk statements.
+            return super()._execute_batch(cursor, qry, data, batch_size)
+
+        statements: list[bytes] = []
+        byte_count = 0
+        for row in data:
+            # Let the driver escape values using the connection's SQL mode.
+            statement = cursor.mogrify(qry, row).encode(
+                cursor.connection.encoding, "surrogateescape"
+            )
+            if statements and (
+                len(statements) >= batch_size
+                or byte_count + 1 + len(statement) > cursor.max_stmt_length
+            ):
+                self._execute_statements(cursor, statements)
+                statements = []
+                byte_count = 0
+            byte_count += len(statement) + bool(statements)
+            statements.append(statement)
+        if statements:
+            self._execute_statements(cursor, statements)
+
+    @staticmethod
+    def _execute_statements(cursor: Any, statements: list[bytes]) -> None:
+        cursor.execute(b";".join(statements))
+        # Errors in later statements surface through nextset(), before commit.
+        while cursor.nextset():
+            pass
 
     def _cursor(self, connection: Any, *, as_dict: bool = False):
         if as_dict:

{yjdata-0.2.2 → yjdata-0.2.3}/src/yjdata/db/sql.py

@@ -294,8 +294,18 @@ class SqlAdapter(DbAdapter):
             return
         with self._connection_scope(write=True) as connection:
             with self._cursor(connection) as cursor:
-                for offset in range(0, len(data), batch_size):
-                    cursor.executemany(qry, data[offset:offset + batch_size])
+                self._execute_batch(cursor, qry, data, batch_size)
+
+    def _execute_batch(
+        self,
+        cursor: Any,
+        qry: Any,
+        data: list,
+        batch_size: int,
+    ) -> None:
+        """Execute batches, allowing adapters to use driver-specific helpers."""
+        for offset in range(0, len(data), batch_size):
+            cursor.executemany(qry, data[offset:offset + batch_size])
 
     @abstractmethod
     def _build_upsert_sql(

{yjdata-0.2.2 → yjdata-0.2.3}/src/yjdata.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yjdata
-Version: 0.2.2
+Version: 0.2.3
 Summary: yjdata SDK for Python
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -98,6 +98,22 @@ the same query structure with double-quoted identifiers. Expressions retain
 normal SQL NULL semantics: for example, `change_type <> 'conceptNegation'`
 also excludes NULL `change_type` values.
 
+## Batch writes
+
+PostgreSQL inserts, upserts and updates use `psycopg2.extras.execute_batch`.
+MySQL inserts and upserts use PyMySQL's optimized `executemany`. MySQL updates
+combine driver-escaped UPDATE statements into one request, preserving input
+order, repeated keys and ordinary `=` matching (NULL keys do not match).
+MySQL connections enable `CLIENT.MULTI_STATEMENTS` for these update batches.
+
+`batch_size` limits rows per batch. MySQL update batches also respect the
+cursor's `max_stmt_length` in encoded bytes (normally about 1 MB); an individual
+larger statement is sent alone and remains subject to the server's packet limit.
+Every result is consumed so errors in later statements propagate before commit.
+With the default transactional connection settings and transactional tables,
+all batches commit together or roll back on failure. This reduces network
+round trips; the server still executes one UPDATE per input row.
+
 ## Tests
 
 ```sh
@@ -109,3 +125,6 @@ parameter ordering, and result conversion. Generated queries execute against
 SQLite's window engine with DB-API placeholders translated, so tests require
 neither optional database drivers nor running servers. This does not replace
 integration testing against MySQL and PostgreSQL.
+
+With `pymysql` installed, the MySQL batch-update tests also exercise the real
+driver's escaping, encoded request sizes and result handling without a server.

{yjdata-0.2.2 → yjdata-0.2.3}/src/yjdata.egg-info/SOURCES.txt

@@ -27,4 +27,6 @@ src/yjdata/web/_input.py
 src/yjdata/web/calendar.py
 src/yjdata/web/client.py
 src/yjdata/web/time_util.py
-tests/test_last_of_group.py
+tests/test_batch_writes.py
+tests/test_last_of_group.py
+tests/test_mysql_batch_updates.py

yjdata-0.2.3/tests/test_batch_writes.py

@@ -0,0 +1,137 @@
+"""Verify batch dispatch and transaction handling without database servers."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, call, sentinel
+
+import polars as pl
+
+from yjdata.db import MysqlAdapter, PostgresAdapter
+from yjdata.db.sql import SqlAdapter
+
+
+def make_adapter(adapter_type):
+    adapter = adapter_type.__new__(adapter_type)
+    SqlAdapter.__init__(adapter)
+    connection = MagicMock()
+    connection.closed = False
+    connection.open = True
+    cursor = connection.cursor.return_value.__enter__.return_value
+    adapter._create_connection = Mock(return_value=connection)
+    adapter._driver = SimpleNamespace(
+        extras=SimpleNamespace(execute_batch=Mock()),
+    )
+    # SQL construction is independent of batch dispatch. Keep the query opaque
+    # to also verify that PostgreSQL composable queries can pass through intact.
+    adapter._build_upsert_sql = Mock(return_value=sentinel.upsert_query)
+    return adapter, connection, cursor
+
+
+class BatchWriteTests(unittest.TestCase):
+    def setUp(self):
+        self.df = pl.DataFrame({"id": [1, 2, 3, 4, 5], "value": ["a", None, "c", "d", "e"]})
+
+    def write(self, adapter, operation, df=None, batch_size=2):
+        df = self.df if df is None else df
+        if operation == "insert_data":
+            return adapter.insert_data("records", df.to_dicts(), batch_size=batch_size)
+        method = getattr(adapter, operation)
+        if operation == "insert_polars":
+            return method("records", df, batch_size=batch_size)
+        return method("records", ["id"], df, batch_size=batch_size)
+
+    def test_postgres_writes_use_driver_batch_helper(self):
+        for operation in ("insert_data", "insert_polars", "upsert_polars", "update_polars"):
+            for batch_size in (1, 2, 10):
+                with self.subTest(operation=operation, batch_size=batch_size):
+                    adapter, connection, cursor = make_adapter(PostgresAdapter)
+                    self.assertIsNone(self.write(adapter, operation, batch_size=batch_size))
+                    execute_batch = adapter._driver.extras.execute_batch
+                    execute_batch.assert_called_once()
+                    args, kwargs = execute_batch.call_args
+                    self.assertIs(args[0], cursor)
+                    expected_rows = self.df.rows()
+                    if operation == "update_polars":
+                        expected_rows = self.df.select("value", "id").rows()
+                    self.assertEqual(args[2], expected_rows)
+                    self.assertEqual(kwargs, {"page_size": batch_size})
+                    if operation == "upsert_polars":
+                        self.assertIs(args[1], sentinel.upsert_query)
+                    cursor.execute.assert_not_called()
+                    cursor.executemany.assert_not_called()
+                    connection.commit.assert_called_once_with()
+                    connection.rollback.assert_not_called()
+                    connection.cursor.return_value.__exit__.assert_called_once()
+
+    def test_postgres_batch_failure_rolls_back_and_propagates(self):
+        for operation in ("insert_polars", "upsert_polars", "update_polars"):
+            with self.subTest(operation=operation):
+                adapter, connection, cursor = make_adapter(PostgresAdapter)
+                error = RuntimeError("batch failed")
+                adapter._driver.extras.execute_batch.side_effect = error
+                with self.assertRaises(RuntimeError) as raised:
+                    self.write(adapter, operation)
+                self.assertIs(raised.exception, error)
+                connection.commit.assert_not_called()
+                connection.rollback.assert_called_once_with()
+                cursor.executemany.assert_not_called()
+                connection.cursor.return_value.__exit__.assert_called_once()
+
+    def test_empty_writes_do_not_connect(self):
+        for adapter_type in (PostgresAdapter, MysqlAdapter):
+            for operation in ("insert_data", "insert_polars", "upsert_polars", "update_polars"):
+                with self.subTest(adapter=adapter_type.__name__, operation=operation):
+                    adapter, _, _ = make_adapter(adapter_type)
+                    self.write(adapter, operation, self.df.head(0))
+                    adapter._create_connection.assert_not_called()
+                    adapter._driver.extras.execute_batch.assert_not_called()
+
+    def test_mysql_keeps_executemany_batches_in_one_transaction(self):
+        for operation in ("insert_data", "insert_polars", "upsert_polars"):
+            with self.subTest(operation=operation):
+                adapter, connection, cursor = make_adapter(MysqlAdapter)
+                self.write(adapter, operation)
+                rows = self.df.rows()
+                qry = cursor.executemany.call_args_list[0].args[0]
+                self.assertEqual(cursor.executemany.call_args_list, [
+                    call(qry, rows[:2]), call(qry, rows[2:4]), call(qry, rows[4:]),
+                ])
+                adapter._driver.extras.execute_batch.assert_not_called()
+                connection.commit.assert_called_once_with()
+                connection.rollback.assert_not_called()
+
+    def test_mysql_updates_batch_requests_and_drain_results(self):
+        adapter, connection, cursor = make_adapter(MysqlAdapter)
+        cursor.connection.encoding = "utf8"
+        cursor.max_stmt_length = 1024000
+        statements = [f"UPDATE records SET value = {i}" for i in range(5)]
+        cursor.mogrify.side_effect = statements
+        cursor.nextset.side_effect = [True, None, True, None, None]
+        self.write(adapter, "update_polars")
+        self.assertEqual(cursor.execute.call_args_list, [
+            call(";".join(statements[:2]).encode()),
+            call(";".join(statements[2:4]).encode()),
+            call(statements[4].encode()),
+        ])
+        self.assertEqual(cursor.nextset.call_count, 5)
+        cursor.executemany.assert_not_called()
+        connection.commit.assert_called_once_with()
+        connection.rollback.assert_not_called()
+
+    def test_mysql_later_statement_error_rolls_back(self):
+        adapter, connection, cursor = make_adapter(MysqlAdapter)
+        cursor.connection.encoding = "utf8"
+        cursor.max_stmt_length = 1024000
+        cursor.mogrify.return_value = "UPDATE records SET value = 1"
+        error = RuntimeError("second update failed")
+        cursor.nextset.side_effect = error
+        with self.assertRaises(RuntimeError) as raised:
+            self.write(adapter, "update_polars")
+        self.assertIs(raised.exception, error)
+        cursor.execute.assert_called_once()
+        connection.commit.assert_not_called()
+        connection.rollback.assert_called_once_with()
+
+
+if __name__ == "__main__":
+    unittest.main()

yjdata-0.2.3/tests/test_mysql_batch_updates.py

@@ -0,0 +1,130 @@
+"""Exercise real PyMySQL escaping and cursors, replacing only server I/O."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import polars as pl
+
+from yjdata.db import MysqlAdapter
+from yjdata.db.sql import SqlAdapter
+
+try:
+    import pymysql
+except ImportError:
+    pymysql = None
+
+
+@unittest.skipIf(pymysql is None, "requires optional pymysql driver")
+class MysqlBatchUpdateDriverTests(unittest.TestCase):
+    def setUp(self):
+        self.adapter = MysqlAdapter(host="localhost")
+        self.connection = pymysql.connect(defer_connect=True)
+        self.connection.server_status = 0
+        self.cursor = self.connection.cursor()
+        self.connection.cursor = Mock(return_value=self.cursor)
+        self.connection.commit = Mock()
+        self.connection.rollback = Mock()
+        self.connection.query = Mock(side_effect=self.query)
+        self.adapter._create_connection = Mock(return_value=self.connection)
+        self.adapter._is_connection_closed = Mock(return_value=False)
+
+    def query(self, sql, **kwargs):
+        self.connection._result = SimpleNamespace(
+            affected_rows=1, warning_count=0, description=None,
+            insert_id=0, rows=None, has_next=False,
+        )
+        return 1
+
+    def requests(self):
+        return [call.args[0] for call in self.connection.query.call_args_list]
+
+    def test_thousand_updates_use_ten_requests(self):
+        df = pl.DataFrame({"id": range(1000), "value": range(1000)})
+        with patch.object(MysqlAdapter, "_execute_batch", SqlAdapter._execute_batch):
+            self.adapter.update_polars("records", ["id"], df, batch_size=100)
+        self.assertEqual(len(self.requests()), 1000)
+        self.connection.query.reset_mock()
+        self.connection.commit.reset_mock()
+        self.cursor = pymysql.cursors.Cursor(self.connection)
+        self.connection.cursor.return_value = self.cursor
+        self.adapter.update_polars("records", ["id"], df, batch_size=100)
+        requests = self.requests()
+        self.assertEqual(len(requests), 10)
+        self.assertTrue(all(request.count(b"UPDATE ") == 100 for request in requests))
+        self.connection.commit.assert_called_once_with()
+
+    def test_composite_keys_nulls_duplicates_and_escaping_preserve_order(self):
+        # Include SQL-looking text, non-ASCII, quotes, backslashes and percent
+        # signs to catch unsafe joining or formatting escaped parameters twice.
+        values = ["中文'; UPDATE records SET v=0; -- \\ %s", None, "last"]
+        df = pl.DataFrame({"id": [1, 2, 1], "region": ["a", None, "a"], "value%": values})
+        qry = "UPDATE `records` SET `value%%` = %s WHERE `id` = %s AND `region` = %s"
+        for sql_mode in (0, pymysql.constants.SERVER_STATUS.SERVER_STATUS_NO_BACKSLASH_ESCAPES):
+            with self.subTest(sql_mode=sql_mode):
+                self.connection.server_status = sql_mode
+                self.connection.query.reset_mock()
+                # A fresh cursor is needed because update_polars closes it.
+                self.cursor = pymysql.cursors.Cursor(self.connection)
+                self.connection.cursor.return_value = self.cursor
+                expected = [self.cursor.mogrify(qry, row).encode("utf8")
+                            for row in df.select("value%", "id", "region").rows()]
+                self.adapter.update_polars("records", ["id", "region"], df, batch_size=2)
+                self.assertEqual(self.requests(), [b";".join(expected[:2]), expected[2]])
+                self.assertIn(b"= NULL", self.requests()[0])
+
+    def test_statement_size_limit_counts_encoded_bytes(self):
+        df = pl.DataFrame({"id": [1, 2, 3], "value": ["中文"] * 3})
+        qry = "UPDATE `records` SET `value` = %s WHERE `id` = %s"
+        statements = [self.cursor.mogrify(qry, ("中文", i)).encode("utf8") for i in (1, 2, 3)]
+        # Two statements fit as characters, but not as encoded bytes.
+        self.cursor.max_stmt_length = len(b";".join(statements[:2])) - 1
+        self.adapter.update_polars("records", ["id"], df)
+        self.assertEqual(self.requests(), statements)
+
+    def test_exact_size_boundary_allows_two_statements(self):
+        df = pl.DataFrame({"id": [1, 2, 3], "value": ["a"] * 3})
+        qry = "UPDATE `records` SET `value` = %s WHERE `id` = %s"
+        statements = [self.cursor.mogrify(qry, ("a", i)).encode("utf8") for i in (1, 2, 3)]
+        self.cursor.max_stmt_length = len(b";".join(statements[:2]))
+        self.adapter.update_polars("records", ["id"], df)
+        self.assertEqual(self.requests(), [b";".join(statements[:2]), statements[2]])
+
+    def test_single_large_statement_is_sent_alone(self):
+        self.cursor.max_stmt_length = 1
+        df = pl.DataFrame({"id": [1, 2], "value": ["a", "b"]})
+        self.adapter.update_polars("records", ["id"], df)
+        self.assertEqual(len(self.requests()), 2)
+        self.assertTrue(all(request.startswith(b"UPDATE ") for request in self.requests()))
+
+    def test_error_in_later_result_prevents_commit(self):
+        error = pymysql.err.IntegrityError(1062, "duplicate entry")
+
+        def query_with_more_results(sql, **kwargs):
+            self.query(sql, **kwargs)
+            self.connection._result.has_next = True
+            return 1
+
+        self.connection.query.side_effect = query_with_more_results
+        self.connection.next_result = Mock(side_effect=error)
+        df = pl.DataFrame({"id": [1, 2, 3], "value": ["a", "b", "c"]})
+        with self.assertRaises(pymysql.err.IntegrityError) as raised:
+            self.adapter.update_polars("records", ["id"], df, batch_size=2)
+        self.assertIs(raised.exception, error)
+        self.assertEqual(len(self.requests()), 1)
+        self.connection.commit.assert_not_called()
+        self.connection.rollback.assert_called_once_with()
+
+    def test_connection_enables_multi_statements_and_preserves_flags(self):
+        for flags in (0, pymysql.constants.CLIENT.FOUND_ROWS):
+            with self.subTest(flags=flags):
+                adapter = MysqlAdapter(host="localhost", client_flag=flags)
+                with patch.object(pymysql, "connect") as connect:
+                    adapter._create_connection()
+                self.assertEqual(connect.call_args.kwargs["client_flag"],
+                                 flags | pymysql.constants.CLIENT.MULTI_STATEMENTS)
+                self.assertEqual(adapter._conn_kwargs["client_flag"], flags)
+
+
+if __name__ == "__main__":
+    unittest.main()