yjdata 0.2.1__tar.gz → 0.2.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {yjdata-0.2.1 → yjdata-0.2.3}/PKG-INFO +20 -1
- {yjdata-0.2.1 → yjdata-0.2.3}/README.md +19 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/pyproject.toml +1 -1
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/db/mysql.py +43 -1
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/db/sql.py +12 -2
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/web/_context.py +3 -4
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata.egg-info/PKG-INFO +20 -1
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata.egg-info/SOURCES.txt +3 -1
- yjdata-0.2.3/tests/test_batch_writes.py +137 -0
- yjdata-0.2.3/tests/test_mysql_batch_updates.py +130 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/setup.cfg +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/__init__.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/_optional.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/dataset/__init__.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/dataset/common.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/dataset/quote.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/dataset/snapshot.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/dataset/trade.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/db/__init__.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/db/adapter.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/db/expression.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/db/postgres.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/dtype.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/utlity/__init__.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/utlity/s3.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/web/__init__.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/web/_input.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/web/calendar.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/web/client.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata/web/time_util.py +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata.egg-info/dependency_links.txt +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata.egg-info/requires.txt +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/src/yjdata.egg-info/top_level.txt +0 -0
- {yjdata-0.2.1 → yjdata-0.2.3}/tests/test_last_of_group.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: yjdata
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.3
|
|
4
4
|
Summary: yjdata SDK for Python
|
|
5
5
|
Requires-Python: >=3.10
|
|
6
6
|
Description-Content-Type: text/markdown
|
|
@@ -98,6 +98,22 @@ the same query structure with double-quoted identifiers. Expressions retain
|
|
|
98
98
|
normal SQL NULL semantics: for example, `change_type <> 'conceptNegation'`
|
|
99
99
|
also excludes NULL `change_type` values.
|
|
100
100
|
|
|
101
|
+
## Batch writes
|
|
102
|
+
|
|
103
|
+
PostgreSQL inserts, upserts and updates use `psycopg2.extras.execute_batch`.
|
|
104
|
+
MySQL inserts and upserts use PyMySQL's optimized `executemany`. MySQL updates
|
|
105
|
+
combine driver-escaped UPDATE statements into one request, preserving input
|
|
106
|
+
order, repeated keys and ordinary `=` matching (NULL keys do not match).
|
|
107
|
+
MySQL connections enable `CLIENT.MULTI_STATEMENTS` for these update batches.
|
|
108
|
+
|
|
109
|
+
`batch_size` limits rows per batch. MySQL update batches also respect the
|
|
110
|
+
cursor's `max_stmt_length` in encoded bytes (normally about 1 MB); an individual
|
|
111
|
+
larger statement is sent alone and remains subject to the server's packet limit.
|
|
112
|
+
Every result is consumed so errors in later statements propagate before commit.
|
|
113
|
+
With the default transactional connection settings and transactional tables,
|
|
114
|
+
all batches commit together or roll back on failure. This reduces network
|
|
115
|
+
round trips; the server still executes one UPDATE per input row.
|
|
116
|
+
|
|
101
117
|
## Tests
|
|
102
118
|
|
|
103
119
|
```sh
|
|
@@ -109,3 +125,6 @@ parameter ordering, and result conversion. Generated queries execute against
|
|
|
109
125
|
SQLite's window engine with DB-API placeholders translated, so tests require
|
|
110
126
|
neither optional database drivers nor running servers. This does not replace
|
|
111
127
|
integration testing against MySQL and PostgreSQL.
|
|
128
|
+
|
|
129
|
+
With `pymysql` installed, the MySQL batch-update tests also exercise the real
|
|
130
|
+
driver's escaping, encoded request sizes and result handling without a server.
|
|
@@ -87,6 +87,22 @@ the same query structure with double-quoted identifiers. Expressions retain
|
|
|
87
87
|
normal SQL NULL semantics: for example, `change_type <> 'conceptNegation'`
|
|
88
88
|
also excludes NULL `change_type` values.
|
|
89
89
|
|
|
90
|
+
## Batch writes
|
|
91
|
+
|
|
92
|
+
PostgreSQL inserts, upserts and updates use `psycopg2.extras.execute_batch`.
|
|
93
|
+
MySQL inserts and upserts use PyMySQL's optimized `executemany`. MySQL updates
|
|
94
|
+
combine driver-escaped UPDATE statements into one request, preserving input
|
|
95
|
+
order, repeated keys and ordinary `=` matching (NULL keys do not match).
|
|
96
|
+
MySQL connections enable `CLIENT.MULTI_STATEMENTS` for these update batches.
|
|
97
|
+
|
|
98
|
+
`batch_size` limits rows per batch. MySQL update batches also respect the
|
|
99
|
+
cursor's `max_stmt_length` in encoded bytes (normally about 1 MB); an individual
|
|
100
|
+
larger statement is sent alone and remains subject to the server's packet limit.
|
|
101
|
+
Every result is consumed so errors in later statements propagate before commit.
|
|
102
|
+
With the default transactional connection settings and transactional tables,
|
|
103
|
+
all batches commit together or roll back on failure. This reduces network
|
|
104
|
+
round trips; the server still executes one UPDATE per input row.
|
|
105
|
+
|
|
90
106
|
## Tests
|
|
91
107
|
|
|
92
108
|
```sh
|
|
@@ -98,3 +114,6 @@ parameter ordering, and result conversion. Generated queries execute against
|
|
|
98
114
|
SQLite's window engine with DB-API placeholders translated, so tests require
|
|
99
115
|
neither optional database drivers nor running servers. This does not replace
|
|
100
116
|
integration testing against MySQL and PostgreSQL.
|
|
117
|
+
|
|
118
|
+
With `pymysql` installed, the MySQL batch-update tests also exercise the real
|
|
119
|
+
driver's escaping, encoded request sizes and result handling without a server.
|
|
@@ -51,7 +51,49 @@ class MysqlAdapter(SqlAdapter):
|
|
|
51
51
|
return upsert_qry
|
|
52
52
|
|
|
53
53
|
def _create_connection(self):
|
|
54
|
-
|
|
54
|
+
kwargs = self._conn_kwargs.copy()
|
|
55
|
+
# UPDATE batches send several statements in one request, in row order.
|
|
56
|
+
kwargs["client_flag"] = (
|
|
57
|
+
kwargs.get("client_flag", 0) | self._driver.constants.CLIENT.MULTI_STATEMENTS
|
|
58
|
+
)
|
|
59
|
+
return self._driver.connect(**kwargs)
|
|
60
|
+
|
|
61
|
+
def _execute_batch(
|
|
62
|
+
self,
|
|
63
|
+
cursor: Any,
|
|
64
|
+
qry: Any,
|
|
65
|
+
data: list,
|
|
66
|
+
batch_size: int,
|
|
67
|
+
) -> None:
|
|
68
|
+
if not isinstance(qry, str) or not qry.startswith("UPDATE "):
|
|
69
|
+
# PyMySQL already combines INSERT/UPSERT values into bulk statements.
|
|
70
|
+
return super()._execute_batch(cursor, qry, data, batch_size)
|
|
71
|
+
|
|
72
|
+
statements: list[bytes] = []
|
|
73
|
+
byte_count = 0
|
|
74
|
+
for row in data:
|
|
75
|
+
# Let the driver escape values using the connection's SQL mode.
|
|
76
|
+
statement = cursor.mogrify(qry, row).encode(
|
|
77
|
+
cursor.connection.encoding, "surrogateescape"
|
|
78
|
+
)
|
|
79
|
+
if statements and (
|
|
80
|
+
len(statements) >= batch_size
|
|
81
|
+
or byte_count + 1 + len(statement) > cursor.max_stmt_length
|
|
82
|
+
):
|
|
83
|
+
self._execute_statements(cursor, statements)
|
|
84
|
+
statements = []
|
|
85
|
+
byte_count = 0
|
|
86
|
+
byte_count += len(statement) + bool(statements)
|
|
87
|
+
statements.append(statement)
|
|
88
|
+
if statements:
|
|
89
|
+
self._execute_statements(cursor, statements)
|
|
90
|
+
|
|
91
|
+
@staticmethod
|
|
92
|
+
def _execute_statements(cursor: Any, statements: list[bytes]) -> None:
|
|
93
|
+
cursor.execute(b";".join(statements))
|
|
94
|
+
# Errors in later statements surface through nextset(), before commit.
|
|
95
|
+
while cursor.nextset():
|
|
96
|
+
pass
|
|
55
97
|
|
|
56
98
|
def _cursor(self, connection: Any, *, as_dict: bool = False):
|
|
57
99
|
if as_dict:
|
|
@@ -294,8 +294,18 @@ class SqlAdapter(DbAdapter):
|
|
|
294
294
|
return
|
|
295
295
|
with self._connection_scope(write=True) as connection:
|
|
296
296
|
with self._cursor(connection) as cursor:
|
|
297
|
-
|
|
298
|
-
|
|
297
|
+
self._execute_batch(cursor, qry, data, batch_size)
|
|
298
|
+
|
|
299
|
+
def _execute_batch(
|
|
300
|
+
self,
|
|
301
|
+
cursor: Any,
|
|
302
|
+
qry: Any,
|
|
303
|
+
data: list,
|
|
304
|
+
batch_size: int,
|
|
305
|
+
) -> None:
|
|
306
|
+
"""Execute batches, allowing adapters to use driver-specific helpers."""
|
|
307
|
+
for offset in range(0, len(data), batch_size):
|
|
308
|
+
cursor.executemany(qry, data[offset:offset + batch_size])
|
|
299
309
|
|
|
300
310
|
@abstractmethod
|
|
301
311
|
def _build_upsert_sql(
|
|
@@ -44,14 +44,13 @@ class RequestContext(abc.ABC):
|
|
|
44
44
|
def _proc_args(kwargs: dict):
|
|
45
45
|
return {
|
|
46
46
|
**kwargs,
|
|
47
|
-
'follow_redirects': True
|
|
48
|
-
'verify': False
|
|
47
|
+
'follow_redirects': True
|
|
49
48
|
}
|
|
50
49
|
|
|
51
50
|
|
|
52
51
|
class SyncRequestContext(RequestContext):
|
|
53
52
|
def __init__(self):
|
|
54
|
-
self._client = httpx.Client()
|
|
53
|
+
self._client = httpx.Client(verify=False)
|
|
55
54
|
|
|
56
55
|
def get(self, url: str, **kwargs):
|
|
57
56
|
res = self._client.get(url, **self._proc_args(kwargs))
|
|
@@ -82,7 +81,7 @@ class SyncRequestContext(RequestContext):
|
|
|
82
81
|
|
|
83
82
|
class AsyncRequestContext(RequestContext):
|
|
84
83
|
def __init__(self):
|
|
85
|
-
self._client = httpx.AsyncClient()
|
|
84
|
+
self._client = httpx.AsyncClient(verify=False)
|
|
86
85
|
|
|
87
86
|
async def get(self, url: str, **kwargs):
|
|
88
87
|
res = await self._client.get(url, **self._proc_args(kwargs))
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: yjdata
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.3
|
|
4
4
|
Summary: yjdata SDK for Python
|
|
5
5
|
Requires-Python: >=3.10
|
|
6
6
|
Description-Content-Type: text/markdown
|
|
@@ -98,6 +98,22 @@ the same query structure with double-quoted identifiers. Expressions retain
|
|
|
98
98
|
normal SQL NULL semantics: for example, `change_type <> 'conceptNegation'`
|
|
99
99
|
also excludes NULL `change_type` values.
|
|
100
100
|
|
|
101
|
+
## Batch writes
|
|
102
|
+
|
|
103
|
+
PostgreSQL inserts, upserts and updates use `psycopg2.extras.execute_batch`.
|
|
104
|
+
MySQL inserts and upserts use PyMySQL's optimized `executemany`. MySQL updates
|
|
105
|
+
combine driver-escaped UPDATE statements into one request, preserving input
|
|
106
|
+
order, repeated keys and ordinary `=` matching (NULL keys do not match).
|
|
107
|
+
MySQL connections enable `CLIENT.MULTI_STATEMENTS` for these update batches.
|
|
108
|
+
|
|
109
|
+
`batch_size` limits rows per batch. MySQL update batches also respect the
|
|
110
|
+
cursor's `max_stmt_length` in encoded bytes (normally about 1 MB); an individual
|
|
111
|
+
larger statement is sent alone and remains subject to the server's packet limit.
|
|
112
|
+
Every result is consumed so errors in later statements propagate before commit.
|
|
113
|
+
With the default transactional connection settings and transactional tables,
|
|
114
|
+
all batches commit together or roll back on failure. This reduces network
|
|
115
|
+
round trips; the server still executes one UPDATE per input row.
|
|
116
|
+
|
|
101
117
|
## Tests
|
|
102
118
|
|
|
103
119
|
```sh
|
|
@@ -109,3 +125,6 @@ parameter ordering, and result conversion. Generated queries execute against
|
|
|
109
125
|
SQLite's window engine with DB-API placeholders translated, so tests require
|
|
110
126
|
neither optional database drivers nor running servers. This does not replace
|
|
111
127
|
integration testing against MySQL and PostgreSQL.
|
|
128
|
+
|
|
129
|
+
With `pymysql` installed, the MySQL batch-update tests also exercise the real
|
|
130
|
+
driver's escaping, encoded request sizes and result handling without a server.
|
|
@@ -0,0 +1,137 @@
|
|
|
1
|
+
"""Verify batch dispatch and transaction handling without database servers."""
|
|
2
|
+
|
|
3
|
+
import unittest
|
|
4
|
+
from types import SimpleNamespace
|
|
5
|
+
from unittest.mock import MagicMock, Mock, call, sentinel
|
|
6
|
+
|
|
7
|
+
import polars as pl
|
|
8
|
+
|
|
9
|
+
from yjdata.db import MysqlAdapter, PostgresAdapter
|
|
10
|
+
from yjdata.db.sql import SqlAdapter
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def make_adapter(adapter_type):
|
|
14
|
+
adapter = adapter_type.__new__(adapter_type)
|
|
15
|
+
SqlAdapter.__init__(adapter)
|
|
16
|
+
connection = MagicMock()
|
|
17
|
+
connection.closed = False
|
|
18
|
+
connection.open = True
|
|
19
|
+
cursor = connection.cursor.return_value.__enter__.return_value
|
|
20
|
+
adapter._create_connection = Mock(return_value=connection)
|
|
21
|
+
adapter._driver = SimpleNamespace(
|
|
22
|
+
extras=SimpleNamespace(execute_batch=Mock()),
|
|
23
|
+
)
|
|
24
|
+
# SQL construction is independent of batch dispatch. Keep the query opaque
|
|
25
|
+
# to also verify that PostgreSQL composable queries can pass through intact.
|
|
26
|
+
adapter._build_upsert_sql = Mock(return_value=sentinel.upsert_query)
|
|
27
|
+
return adapter, connection, cursor
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class BatchWriteTests(unittest.TestCase):
|
|
31
|
+
def setUp(self):
|
|
32
|
+
self.df = pl.DataFrame({"id": [1, 2, 3, 4, 5], "value": ["a", None, "c", "d", "e"]})
|
|
33
|
+
|
|
34
|
+
def write(self, adapter, operation, df=None, batch_size=2):
|
|
35
|
+
df = self.df if df is None else df
|
|
36
|
+
if operation == "insert_data":
|
|
37
|
+
return adapter.insert_data("records", df.to_dicts(), batch_size=batch_size)
|
|
38
|
+
method = getattr(adapter, operation)
|
|
39
|
+
if operation == "insert_polars":
|
|
40
|
+
return method("records", df, batch_size=batch_size)
|
|
41
|
+
return method("records", ["id"], df, batch_size=batch_size)
|
|
42
|
+
|
|
43
|
+
def test_postgres_writes_use_driver_batch_helper(self):
|
|
44
|
+
for operation in ("insert_data", "insert_polars", "upsert_polars", "update_polars"):
|
|
45
|
+
for batch_size in (1, 2, 10):
|
|
46
|
+
with self.subTest(operation=operation, batch_size=batch_size):
|
|
47
|
+
adapter, connection, cursor = make_adapter(PostgresAdapter)
|
|
48
|
+
self.assertIsNone(self.write(adapter, operation, batch_size=batch_size))
|
|
49
|
+
execute_batch = adapter._driver.extras.execute_batch
|
|
50
|
+
execute_batch.assert_called_once()
|
|
51
|
+
args, kwargs = execute_batch.call_args
|
|
52
|
+
self.assertIs(args[0], cursor)
|
|
53
|
+
expected_rows = self.df.rows()
|
|
54
|
+
if operation == "update_polars":
|
|
55
|
+
expected_rows = self.df.select("value", "id").rows()
|
|
56
|
+
self.assertEqual(args[2], expected_rows)
|
|
57
|
+
self.assertEqual(kwargs, {"page_size": batch_size})
|
|
58
|
+
if operation == "upsert_polars":
|
|
59
|
+
self.assertIs(args[1], sentinel.upsert_query)
|
|
60
|
+
cursor.execute.assert_not_called()
|
|
61
|
+
cursor.executemany.assert_not_called()
|
|
62
|
+
connection.commit.assert_called_once_with()
|
|
63
|
+
connection.rollback.assert_not_called()
|
|
64
|
+
connection.cursor.return_value.__exit__.assert_called_once()
|
|
65
|
+
|
|
66
|
+
def test_postgres_batch_failure_rolls_back_and_propagates(self):
|
|
67
|
+
for operation in ("insert_polars", "upsert_polars", "update_polars"):
|
|
68
|
+
with self.subTest(operation=operation):
|
|
69
|
+
adapter, connection, cursor = make_adapter(PostgresAdapter)
|
|
70
|
+
error = RuntimeError("batch failed")
|
|
71
|
+
adapter._driver.extras.execute_batch.side_effect = error
|
|
72
|
+
with self.assertRaises(RuntimeError) as raised:
|
|
73
|
+
self.write(adapter, operation)
|
|
74
|
+
self.assertIs(raised.exception, error)
|
|
75
|
+
connection.commit.assert_not_called()
|
|
76
|
+
connection.rollback.assert_called_once_with()
|
|
77
|
+
cursor.executemany.assert_not_called()
|
|
78
|
+
connection.cursor.return_value.__exit__.assert_called_once()
|
|
79
|
+
|
|
80
|
+
def test_empty_writes_do_not_connect(self):
|
|
81
|
+
for adapter_type in (PostgresAdapter, MysqlAdapter):
|
|
82
|
+
for operation in ("insert_data", "insert_polars", "upsert_polars", "update_polars"):
|
|
83
|
+
with self.subTest(adapter=adapter_type.__name__, operation=operation):
|
|
84
|
+
adapter, _, _ = make_adapter(adapter_type)
|
|
85
|
+
self.write(adapter, operation, self.df.head(0))
|
|
86
|
+
adapter._create_connection.assert_not_called()
|
|
87
|
+
adapter._driver.extras.execute_batch.assert_not_called()
|
|
88
|
+
|
|
89
|
+
def test_mysql_keeps_executemany_batches_in_one_transaction(self):
|
|
90
|
+
for operation in ("insert_data", "insert_polars", "upsert_polars"):
|
|
91
|
+
with self.subTest(operation=operation):
|
|
92
|
+
adapter, connection, cursor = make_adapter(MysqlAdapter)
|
|
93
|
+
self.write(adapter, operation)
|
|
94
|
+
rows = self.df.rows()
|
|
95
|
+
qry = cursor.executemany.call_args_list[0].args[0]
|
|
96
|
+
self.assertEqual(cursor.executemany.call_args_list, [
|
|
97
|
+
call(qry, rows[:2]), call(qry, rows[2:4]), call(qry, rows[4:]),
|
|
98
|
+
])
|
|
99
|
+
adapter._driver.extras.execute_batch.assert_not_called()
|
|
100
|
+
connection.commit.assert_called_once_with()
|
|
101
|
+
connection.rollback.assert_not_called()
|
|
102
|
+
|
|
103
|
+
def test_mysql_updates_batch_requests_and_drain_results(self):
|
|
104
|
+
adapter, connection, cursor = make_adapter(MysqlAdapter)
|
|
105
|
+
cursor.connection.encoding = "utf8"
|
|
106
|
+
cursor.max_stmt_length = 1024000
|
|
107
|
+
statements = [f"UPDATE records SET value = {i}" for i in range(5)]
|
|
108
|
+
cursor.mogrify.side_effect = statements
|
|
109
|
+
cursor.nextset.side_effect = [True, None, True, None, None]
|
|
110
|
+
self.write(adapter, "update_polars")
|
|
111
|
+
self.assertEqual(cursor.execute.call_args_list, [
|
|
112
|
+
call(";".join(statements[:2]).encode()),
|
|
113
|
+
call(";".join(statements[2:4]).encode()),
|
|
114
|
+
call(statements[4].encode()),
|
|
115
|
+
])
|
|
116
|
+
self.assertEqual(cursor.nextset.call_count, 5)
|
|
117
|
+
cursor.executemany.assert_not_called()
|
|
118
|
+
connection.commit.assert_called_once_with()
|
|
119
|
+
connection.rollback.assert_not_called()
|
|
120
|
+
|
|
121
|
+
def test_mysql_later_statement_error_rolls_back(self):
|
|
122
|
+
adapter, connection, cursor = make_adapter(MysqlAdapter)
|
|
123
|
+
cursor.connection.encoding = "utf8"
|
|
124
|
+
cursor.max_stmt_length = 1024000
|
|
125
|
+
cursor.mogrify.return_value = "UPDATE records SET value = 1"
|
|
126
|
+
error = RuntimeError("second update failed")
|
|
127
|
+
cursor.nextset.side_effect = error
|
|
128
|
+
with self.assertRaises(RuntimeError) as raised:
|
|
129
|
+
self.write(adapter, "update_polars")
|
|
130
|
+
self.assertIs(raised.exception, error)
|
|
131
|
+
cursor.execute.assert_called_once()
|
|
132
|
+
connection.commit.assert_not_called()
|
|
133
|
+
connection.rollback.assert_called_once_with()
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
if __name__ == "__main__":
|
|
137
|
+
unittest.main()
|
|
@@ -0,0 +1,130 @@
|
|
|
1
|
+
"""Exercise real PyMySQL escaping and cursors, replacing only server I/O."""
|
|
2
|
+
|
|
3
|
+
import unittest
|
|
4
|
+
from types import SimpleNamespace
|
|
5
|
+
from unittest.mock import Mock, patch
|
|
6
|
+
|
|
7
|
+
import polars as pl
|
|
8
|
+
|
|
9
|
+
from yjdata.db import MysqlAdapter
|
|
10
|
+
from yjdata.db.sql import SqlAdapter
|
|
11
|
+
|
|
12
|
+
try:
|
|
13
|
+
import pymysql
|
|
14
|
+
except ImportError:
|
|
15
|
+
pymysql = None
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@unittest.skipIf(pymysql is None, "requires optional pymysql driver")
|
|
19
|
+
class MysqlBatchUpdateDriverTests(unittest.TestCase):
|
|
20
|
+
def setUp(self):
|
|
21
|
+
self.adapter = MysqlAdapter(host="localhost")
|
|
22
|
+
self.connection = pymysql.connect(defer_connect=True)
|
|
23
|
+
self.connection.server_status = 0
|
|
24
|
+
self.cursor = self.connection.cursor()
|
|
25
|
+
self.connection.cursor = Mock(return_value=self.cursor)
|
|
26
|
+
self.connection.commit = Mock()
|
|
27
|
+
self.connection.rollback = Mock()
|
|
28
|
+
self.connection.query = Mock(side_effect=self.query)
|
|
29
|
+
self.adapter._create_connection = Mock(return_value=self.connection)
|
|
30
|
+
self.adapter._is_connection_closed = Mock(return_value=False)
|
|
31
|
+
|
|
32
|
+
def query(self, sql, **kwargs):
|
|
33
|
+
self.connection._result = SimpleNamespace(
|
|
34
|
+
affected_rows=1, warning_count=0, description=None,
|
|
35
|
+
insert_id=0, rows=None, has_next=False,
|
|
36
|
+
)
|
|
37
|
+
return 1
|
|
38
|
+
|
|
39
|
+
def requests(self):
|
|
40
|
+
return [call.args[0] for call in self.connection.query.call_args_list]
|
|
41
|
+
|
|
42
|
+
def test_thousand_updates_use_ten_requests(self):
|
|
43
|
+
df = pl.DataFrame({"id": range(1000), "value": range(1000)})
|
|
44
|
+
with patch.object(MysqlAdapter, "_execute_batch", SqlAdapter._execute_batch):
|
|
45
|
+
self.adapter.update_polars("records", ["id"], df, batch_size=100)
|
|
46
|
+
self.assertEqual(len(self.requests()), 1000)
|
|
47
|
+
self.connection.query.reset_mock()
|
|
48
|
+
self.connection.commit.reset_mock()
|
|
49
|
+
self.cursor = pymysql.cursors.Cursor(self.connection)
|
|
50
|
+
self.connection.cursor.return_value = self.cursor
|
|
51
|
+
self.adapter.update_polars("records", ["id"], df, batch_size=100)
|
|
52
|
+
requests = self.requests()
|
|
53
|
+
self.assertEqual(len(requests), 10)
|
|
54
|
+
self.assertTrue(all(request.count(b"UPDATE ") == 100 for request in requests))
|
|
55
|
+
self.connection.commit.assert_called_once_with()
|
|
56
|
+
|
|
57
|
+
def test_composite_keys_nulls_duplicates_and_escaping_preserve_order(self):
|
|
58
|
+
# Include SQL-looking text, non-ASCII, quotes, backslashes and percent
|
|
59
|
+
# signs to catch unsafe joining or formatting escaped parameters twice.
|
|
60
|
+
values = ["中文'; UPDATE records SET v=0; -- \\ %s", None, "last"]
|
|
61
|
+
df = pl.DataFrame({"id": [1, 2, 1], "region": ["a", None, "a"], "value%": values})
|
|
62
|
+
qry = "UPDATE `records` SET `value%%` = %s WHERE `id` = %s AND `region` = %s"
|
|
63
|
+
for sql_mode in (0, pymysql.constants.SERVER_STATUS.SERVER_STATUS_NO_BACKSLASH_ESCAPES):
|
|
64
|
+
with self.subTest(sql_mode=sql_mode):
|
|
65
|
+
self.connection.server_status = sql_mode
|
|
66
|
+
self.connection.query.reset_mock()
|
|
67
|
+
# A fresh cursor is needed because update_polars closes it.
|
|
68
|
+
self.cursor = pymysql.cursors.Cursor(self.connection)
|
|
69
|
+
self.connection.cursor.return_value = self.cursor
|
|
70
|
+
expected = [self.cursor.mogrify(qry, row).encode("utf8")
|
|
71
|
+
for row in df.select("value%", "id", "region").rows()]
|
|
72
|
+
self.adapter.update_polars("records", ["id", "region"], df, batch_size=2)
|
|
73
|
+
self.assertEqual(self.requests(), [b";".join(expected[:2]), expected[2]])
|
|
74
|
+
self.assertIn(b"= NULL", self.requests()[0])
|
|
75
|
+
|
|
76
|
+
def test_statement_size_limit_counts_encoded_bytes(self):
|
|
77
|
+
df = pl.DataFrame({"id": [1, 2, 3], "value": ["中文"] * 3})
|
|
78
|
+
qry = "UPDATE `records` SET `value` = %s WHERE `id` = %s"
|
|
79
|
+
statements = [self.cursor.mogrify(qry, ("中文", i)).encode("utf8") for i in (1, 2, 3)]
|
|
80
|
+
# Two statements fit as characters, but not as encoded bytes.
|
|
81
|
+
self.cursor.max_stmt_length = len(b";".join(statements[:2])) - 1
|
|
82
|
+
self.adapter.update_polars("records", ["id"], df)
|
|
83
|
+
self.assertEqual(self.requests(), statements)
|
|
84
|
+
|
|
85
|
+
def test_exact_size_boundary_allows_two_statements(self):
|
|
86
|
+
df = pl.DataFrame({"id": [1, 2, 3], "value": ["a"] * 3})
|
|
87
|
+
qry = "UPDATE `records` SET `value` = %s WHERE `id` = %s"
|
|
88
|
+
statements = [self.cursor.mogrify(qry, ("a", i)).encode("utf8") for i in (1, 2, 3)]
|
|
89
|
+
self.cursor.max_stmt_length = len(b";".join(statements[:2]))
|
|
90
|
+
self.adapter.update_polars("records", ["id"], df)
|
|
91
|
+
self.assertEqual(self.requests(), [b";".join(statements[:2]), statements[2]])
|
|
92
|
+
|
|
93
|
+
def test_single_large_statement_is_sent_alone(self):
|
|
94
|
+
self.cursor.max_stmt_length = 1
|
|
95
|
+
df = pl.DataFrame({"id": [1, 2], "value": ["a", "b"]})
|
|
96
|
+
self.adapter.update_polars("records", ["id"], df)
|
|
97
|
+
self.assertEqual(len(self.requests()), 2)
|
|
98
|
+
self.assertTrue(all(request.startswith(b"UPDATE ") for request in self.requests()))
|
|
99
|
+
|
|
100
|
+
def test_error_in_later_result_prevents_commit(self):
|
|
101
|
+
error = pymysql.err.IntegrityError(1062, "duplicate entry")
|
|
102
|
+
|
|
103
|
+
def query_with_more_results(sql, **kwargs):
|
|
104
|
+
self.query(sql, **kwargs)
|
|
105
|
+
self.connection._result.has_next = True
|
|
106
|
+
return 1
|
|
107
|
+
|
|
108
|
+
self.connection.query.side_effect = query_with_more_results
|
|
109
|
+
self.connection.next_result = Mock(side_effect=error)
|
|
110
|
+
df = pl.DataFrame({"id": [1, 2, 3], "value": ["a", "b", "c"]})
|
|
111
|
+
with self.assertRaises(pymysql.err.IntegrityError) as raised:
|
|
112
|
+
self.adapter.update_polars("records", ["id"], df, batch_size=2)
|
|
113
|
+
self.assertIs(raised.exception, error)
|
|
114
|
+
self.assertEqual(len(self.requests()), 1)
|
|
115
|
+
self.connection.commit.assert_not_called()
|
|
116
|
+
self.connection.rollback.assert_called_once_with()
|
|
117
|
+
|
|
118
|
+
def test_connection_enables_multi_statements_and_preserves_flags(self):
|
|
119
|
+
for flags in (0, pymysql.constants.CLIENT.FOUND_ROWS):
|
|
120
|
+
with self.subTest(flags=flags):
|
|
121
|
+
adapter = MysqlAdapter(host="localhost", client_flag=flags)
|
|
122
|
+
with patch.object(pymysql, "connect") as connect:
|
|
123
|
+
adapter._create_connection()
|
|
124
|
+
self.assertEqual(connect.call_args.kwargs["client_flag"],
|
|
125
|
+
flags | pymysql.constants.CLIENT.MULTI_STATEMENTS)
|
|
126
|
+
self.assertEqual(adapter._conn_kwargs["client_flag"], flags)
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
if __name__ == "__main__":
|
|
130
|
+
unittest.main()
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|