yjdata 0.2.2__tar.gz → 0.2.4__tar.gz

yjdata-0.2.4/PKG-INFO

@@ -0,0 +1,66 @@
+Metadata-Version: 2.4
+Name: yjdata
+Version: 0.2.4
+Summary: yjdata SDK for Python
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: httpx>=0.25
+Requires-Dist: polars>=1.30
+Requires-Dist: pyarrow>=18
+Requires-Dist: pydantic>=2.11
+
+# yjdata
+
+Python SDK for the YJData web API, MySQL/PostgreSQL adapters, and Parquet utilities.
+
+## Web API
+
+```python
+from yjdata import WebDataClient
+
+client = WebDataClient()
+prices = client.cn_stock.stk_eod_md("2026-11-06", ids="000001.XSHE")
+for price in prices:
+    print(price.id, price.close)
+
+trading_days = client.calendar.range("2026-11-01", "2026-11-06")
+```
+
+| Method | Data |
+| --- | --- |
+| `cn_stock.stk_eod_md(date, ids=None, stable_id=True)` | Stock end-of-day prices |
+| `cn_stock.stk_sod_ref(date, ids=None)` | Stock pre-market reference data |
+| `cn_future.fut_eod_md(date, ids=None)` | Futures end-of-day prices |
+| `cn_future.fut_sod_ref(date, ids=None)` | Futures pre-market reference data |
+| `calendar.next_n(start, n)`, `last_n(end, n)`, `range(start, end)`, `is_trading_day(date)` | Trading calendar; optional `calendar="XSHG"` |
+
+- Dates accept `YYYY-MM-DD`, `YYYYMMDD`, `datetime.date`, or `datetime.datetime`.
+- IDs accept lists or comma-separated strings. `None` returns all records;
+  `[]` or a blank string returns none.
+- Stock prices use first-listing IDs by default. `stable_id=False` uses the
+  ticker on the requested date.
+- Market data returns Pydantic models with friendly fields such as `close`,
+  `volume`, and `display_id`; `time` is a `YYYY-MM-DD` string. Models accept
+  friendly names and API aliases. `model_dump(by_alias=True)` exports API names.
+- Use `AsyncWebDataClient` with `await` for the same methods.
+
+See the [OpenAPI schema](https://yjdata-api.yaojingriver.internal/openapi.json)
+for endpoint and field details.
+
+## Database adapters
+
+`MysqlAdapter` (MySQL 8.0+) and `PostgresAdapter` support
+`select_last_of_group(table, group_by=..., sort_by=...)` and
+`select_last_of_group_polars(...)`, returning dictionaries or a Polars DataFrame.
+
+These methods select the last complete row per group in SQL. Use `ascending`
+to control sort direction, `filter` before selection, `post_filter` afterward,
+and `columns` for projection. NULLs sort below non-NULL values; include a unique
+tie-breaker in `sort_by` for deterministic selection. Group output order is
+unspecified.
+
+## Tests
+
+```sh
+PYTHONPATH=src python -m unittest discover -s tests -v
+```

yjdata-0.2.4/README.md

@@ -0,0 +1,55 @@
+# yjdata
+
+Python SDK for the YJData web API, MySQL/PostgreSQL adapters, and Parquet utilities.
+
+## Web API
+
+```python
+from yjdata import WebDataClient
+
+client = WebDataClient()
+prices = client.cn_stock.stk_eod_md("2026-11-06", ids="000001.XSHE")
+for price in prices:
+    print(price.id, price.close)
+
+trading_days = client.calendar.range("2026-11-01", "2026-11-06")
+```
+
+| Method | Data |
+| --- | --- |
+| `cn_stock.stk_eod_md(date, ids=None, stable_id=True)` | Stock end-of-day prices |
+| `cn_stock.stk_sod_ref(date, ids=None)` | Stock pre-market reference data |
+| `cn_future.fut_eod_md(date, ids=None)` | Futures end-of-day prices |
+| `cn_future.fut_sod_ref(date, ids=None)` | Futures pre-market reference data |
+| `calendar.next_n(start, n)`, `last_n(end, n)`, `range(start, end)`, `is_trading_day(date)` | Trading calendar; optional `calendar="XSHG"` |
+
+- Dates accept `YYYY-MM-DD`, `YYYYMMDD`, `datetime.date`, or `datetime.datetime`.
+- IDs accept lists or comma-separated strings. `None` returns all records;
+  `[]` or a blank string returns none.
+- Stock prices use first-listing IDs by default. `stable_id=False` uses the
+  ticker on the requested date.
+- Market data returns Pydantic models with friendly fields such as `close`,
+  `volume`, and `display_id`; `time` is a `YYYY-MM-DD` string. Models accept
+  friendly names and API aliases. `model_dump(by_alias=True)` exports API names.
+- Use `AsyncWebDataClient` with `await` for the same methods.
+
+See the [OpenAPI schema](https://yjdata-api.yaojingriver.internal/openapi.json)
+for endpoint and field details.
+
+## Database adapters
+
+`MysqlAdapter` (MySQL 8.0+) and `PostgresAdapter` support
+`select_last_of_group(table, group_by=..., sort_by=...)` and
+`select_last_of_group_polars(...)`, returning dictionaries or a Polars DataFrame.
+
+These methods select the last complete row per group in SQL. Use `ascending`
+to control sort direction, `filter` before selection, `post_filter` afterward,
+and `columns` for projection. NULLs sort below non-NULL values; include a unique
+tie-breaker in `sort_by` for deterministic selection. Group output order is
+unspecified.
+
+## Tests
+
+```sh
+PYTHONPATH=src python -m unittest discover -s tests -v
+```

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

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

{yjdata-0.2.2 → yjdata-0.2.4}/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.4}/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.4}/src/yjdata/web/calendar.py

@@ -49,7 +49,7 @@ class CalendarWebClient:
         return self._context.post(f'{self._api_root}/next-n', payload={
             'start': parse_date_str_arg(start, False),
             'n': n,
-            'calendar': calendar
+            'exchange': calendar
         })
 
     def last_n(
@@ -83,7 +83,7 @@ class CalendarWebClient:
         return self._context.post(f'{self._api_root}/last-n', payload={
             'end': parse_date_str_arg(end, False),
             'n': n,
-            'calendar': calendar
+            'exchange': calendar
         })
 
     def range(
@@ -118,7 +118,7 @@ class CalendarWebClient:
         return self._context.post(f'{self._api_root}/range', payload={
             'start': parse_date_str_arg(start, False),
             'end': parse_date_str_arg(end, False),
-            'calendar': calendar
+            'exchange': calendar
         })
 
     @abstractmethod

{yjdata-0.2.2 → yjdata-0.2.4}/src/yjdata/web/client.py

@@ -1,6 +1,8 @@
 import os
 
 from .calendar import SyncCalenderWebClient, AsyncCalendarWebClient
+from .cn_future import CnFutureWebClient
+from .cn_stock import CnStockWebClient
 from ._context import RequestContext, SyncRequestContext, AsyncRequestContext
 
 
@@ -23,6 +25,16 @@ class _WebClient:
             self._host = host
         self._context = context
 
+    @property
+    def cn_stock(self) -> CnStockWebClient:
+        """API set for domestic stock market data."""
+        return CnStockWebClient(self._host, self._context)
+
+    @property
+    def cn_future(self) -> CnFutureWebClient:
+        """API set for domestic futures market data."""
+        return CnFutureWebClient(self._host, self._context)
+
 
 class WebDataClient(_WebClient):
     """Client for accessing YJData web API."""

yjdata-0.2.4/src/yjdata/web/cn_future.py

@@ -0,0 +1,110 @@
+from typing import Awaitable, List, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ._context import RequestContext
+from ._input import DateStrArgType, StrListArgType, parse_date_str_arg, parse_str_list_arg
+
+
+__all__ = ['CnFutureWebClient', 'FutEodMd', 'FutSodRef']
+
+
+class FutEodMd(BaseModel):
+    """Futures end-of-day prices, volume, turnover and open interest."""
+
+    model_config = ConfigDict(validate_by_name=True)
+
+    id: str
+    time: str          = Field(alias='date')
+    pre_settle: float  = Field(alias='ps')
+    pre_close: float   = Field(alias='pc')
+    open: float        = Field(alias='o')
+    high: float        = Field(alias='h')
+    low: float         = Field(alias='l')
+    close: float       = Field(alias='c')
+    settle: float      = Field(alias='s')
+    volume: int        = Field(alias='vol')
+    amount: float      = Field(alias='amt')
+    open_interest: int = Field(alias='oi')
+
+
+class FutSodRef(BaseModel):
+    """Futures pre-market reference data, including contract specifications."""
+
+    model_config = ConfigDict(validate_by_name=True)
+
+    id: str
+    time: str                       = Field(alias='date')
+    name: str
+    display_id: str                 = Field(alias='dispId')
+    ticker: str
+    trading_status: str             = Field(alias='status')
+    underlying_id: Union[str, None] = Field(alias='ulId')
+    product: str                    = Field(alias='prod')
+    multiplier: float               = Field(alias='mult')
+    tick: float
+    pre_close: float                = Field(alias='pc')
+    pre_settle: float               = Field(alias='ps')
+    auct_hours: str                 = Field(alias='auct')
+    cont_hours: str                 = Field(alias='cont')
+    upper_limit: float              = Field(alias='ul')
+    lower_limit: float              = Field(alias='ll')
+
+
+class CnFutureWebClient:
+
+    def __init__(self, host: str, context: RequestContext):
+        self._api_root = f'http://{host}/v1/cn_future'
+        self._context = context
+
+    def fut_eod_md(
+            self,
+            date: DateStrArgType,
+            ids: StrListArgType = None
+    ) -> Union[List[FutEodMd], Awaitable[List[FutEodMd]]]:
+        """Get domestic futures end-of-day market data for a trading date.
+
+        Parameters
+        ----------
+        date : str, datetime.date or datetime.datetime
+            Trading date; strings accept YYYY-MM-DD or YYYYMMDD.
+        ids : str or list[str], optional
+            IDs as a list or comma-separated string. None returns all records;
+            an empty list or blank string returns no records.
+
+        Returns
+        -------
+        list[FutEodMd] or Awaitable[list[FutEodMd]]
+            Validated records; await the result with AsyncWebDataClient.
+        """
+        return self._context.post_model_list(FutEodMd, f'{self._api_root}/fut-eod-md', payload={
+            'date': parse_date_str_arg(date, False),
+            'ids': None if ids is None else parse_str_list_arg(ids)
+        })
+
+    def fut_sod_ref(
+            self,
+            date: DateStrArgType,
+            ids: StrListArgType = None
+    ) -> Union[List[FutSodRef], Awaitable[List[FutSodRef]]]:
+        """Get domestic futures pre-market reference data for a trading date.
+
+        Parameters
+        ----------
+        date : str, datetime.date or datetime.datetime
+            Trading date; strings accept YYYY-MM-DD or YYYYMMDD.
+        ids : str or list[str], optional
+            IDs as a list or comma-separated string. None returns all records;
+            an empty list or blank string returns no records.
+
+        Returns
+        -------
+        list[FutSodRef] or Awaitable[list[FutSodRef]]
+            Validated records; await the result with AsyncWebDataClient.
+            API fields dispId and ulId use display_id and underlying_id.
+            underlying_id is None for contracts without an underlying index.
+        """
+        return self._context.post_model_list(FutSodRef, f'{self._api_root}/fut-sod-ref', payload={
+            'date': parse_date_str_arg(date, False),
+            'ids': None if ids is None else parse_str_list_arg(ids)
+        })

yjdata-0.2.4/src/yjdata/web/cn_stock.py

@@ -0,0 +1,108 @@
+from typing import Awaitable, List, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ._context import RequestContext
+from ._input import DateStrArgType, StrListArgType, parse_date_str_arg, parse_str_list_arg
+
+
+__all__ = ['CnStockWebClient', 'StkEodMd', 'StkSodRef']
+
+
+class StkEodMd(BaseModel):
+    """Stock end-of-day prices, volume and turnover."""
+
+    model_config = ConfigDict(validate_by_name=True)
+
+    id: str
+    time: str        = Field(alias='date')
+    pre_close: float = Field(alias='pc')
+    open: float      = Field(alias='o')
+    high: float      = Field(alias='h')
+    low: float       = Field(alias='l')
+    close: float     = Field(alias='c')
+    volume: int      = Field(alias='vol')
+    amount: float    = Field(alias='amt')
+
+
+class StkSodRef(BaseModel):
+    """Stock pre-market reference data. IDs use the first listing ticker."""
+
+    model_config = ConfigDict(validate_by_name=True)
+
+    id: str
+    time: str           = Field(alias='date')
+    name: str
+    display_id: str     = Field(alias='dispId')
+    ticker: str
+    trading_status: str = Field(alias='status')
+    st_status: str      = Field(alias='st')
+    pre_close: float    = Field(alias='pc')
+    prior_close: float  = Field(alias='prc')
+    upper_limit: float  = Field(alias='ul')
+    lower_limit: float  = Field(alias='ll')
+    auct_hours: str     = Field(alias='auct')
+    cont_hours: str     = Field(alias='cont')
+
+
+class CnStockWebClient:
+
+    def __init__(self, host: str, context: RequestContext):
+        self._api_root = f'http://{host}/v1/cn_stock'
+        self._context = context
+
+    def stk_eod_md(
+            self,
+            date: DateStrArgType,
+            ids: StrListArgType = None,
+            stable_id: bool = True
+    ) -> Union[List[StkEodMd], Awaitable[List[StkEodMd]]]:
+        """Get stock end-of-day market data for a trading date.
+
+        Parameters
+        ----------
+        date : str, datetime.date or datetime.datetime
+            Trading date; strings accept YYYY-MM-DD or YYYYMMDD.
+        ids : str or list[str], optional
+            IDs as a list or comma-separated string. None returns all records;
+            an empty list or blank string returns no records.
+        stable_id : bool, default: True
+            Use the first listing ticker plus exchange for both the ID filter
+            and returned IDs. False uses the ticker on the requested date.
+
+        Returns
+        -------
+        list[StkEodMd] or Awaitable[list[StkEodMd]]
+            Validated records; await the result with AsyncWebDataClient.
+        """
+        return self._context.post_model_list(StkEodMd, f'{self._api_root}/stk-eod-md', payload={
+            'date': parse_date_str_arg(date, False),
+            'ids': None if ids is None else parse_str_list_arg(ids),
+            'stableId': stable_id
+        })
+
+    def stk_sod_ref(
+            self,
+            date: DateStrArgType,
+            ids: StrListArgType = None
+    ) -> Union[List[StkSodRef], Awaitable[List[StkSodRef]]]:
+        """Get stock pre-market reference data for a trading date.
+
+        Parameters
+        ----------
+        date : str, datetime.date or datetime.datetime
+            Trading date; strings accept YYYY-MM-DD or YYYYMMDD.
+        ids : str or list[str], optional
+            IDs as a list or comma-separated string. None returns all records;
+            an empty list or blank string returns no records.
+
+        Returns
+        -------
+        list[StkSodRef] or Awaitable[list[StkSodRef]]
+            Validated records; await the result with AsyncWebDataClient.
+            Fields use backend Python names, such as display_id for dispId.
+        """
+        return self._context.post_model_list(StkSodRef, f'{self._api_root}/stk-sod-ref', payload={
+            'date': parse_date_str_arg(date, False),
+            'ids': None if ids is None else parse_str_list_arg(ids)
+        })

yjdata-0.2.4/src/yjdata.egg-info/PKG-INFO

@@ -0,0 +1,66 @@
+Metadata-Version: 2.4
+Name: yjdata
+Version: 0.2.4
+Summary: yjdata SDK for Python
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: httpx>=0.25
+Requires-Dist: polars>=1.30
+Requires-Dist: pyarrow>=18
+Requires-Dist: pydantic>=2.11
+
+# yjdata
+
+Python SDK for the YJData web API, MySQL/PostgreSQL adapters, and Parquet utilities.
+
+## Web API
+
+```python
+from yjdata import WebDataClient
+
+client = WebDataClient()
+prices = client.cn_stock.stk_eod_md("2026-11-06", ids="000001.XSHE")
+for price in prices:
+    print(price.id, price.close)
+
+trading_days = client.calendar.range("2026-11-01", "2026-11-06")
+```
+
+| Method | Data |
+| --- | --- |
+| `cn_stock.stk_eod_md(date, ids=None, stable_id=True)` | Stock end-of-day prices |
+| `cn_stock.stk_sod_ref(date, ids=None)` | Stock pre-market reference data |
+| `cn_future.fut_eod_md(date, ids=None)` | Futures end-of-day prices |
+| `cn_future.fut_sod_ref(date, ids=None)` | Futures pre-market reference data |
+| `calendar.next_n(start, n)`, `last_n(end, n)`, `range(start, end)`, `is_trading_day(date)` | Trading calendar; optional `calendar="XSHG"` |
+
+- Dates accept `YYYY-MM-DD`, `YYYYMMDD`, `datetime.date`, or `datetime.datetime`.
+- IDs accept lists or comma-separated strings. `None` returns all records;
+  `[]` or a blank string returns none.
+- Stock prices use first-listing IDs by default. `stable_id=False` uses the
+  ticker on the requested date.
+- Market data returns Pydantic models with friendly fields such as `close`,
+  `volume`, and `display_id`; `time` is a `YYYY-MM-DD` string. Models accept
+  friendly names and API aliases. `model_dump(by_alias=True)` exports API names.
+- Use `AsyncWebDataClient` with `await` for the same methods.
+
+See the [OpenAPI schema](https://yjdata-api.yaojingriver.internal/openapi.json)
+for endpoint and field details.
+
+## Database adapters
+
+`MysqlAdapter` (MySQL 8.0+) and `PostgresAdapter` support
+`select_last_of_group(table, group_by=..., sort_by=...)` and
+`select_last_of_group_polars(...)`, returning dictionaries or a Polars DataFrame.
+
+These methods select the last complete row per group in SQL. Use `ascending`
+to control sort direction, `filter` before selection, `post_filter` afterward,
+and `columns` for projection. NULLs sort below non-NULL values; include a unique
+tie-breaker in `sort_by` for deterministic selection. Group output order is
+unspecified.
+
+## Tests
+
+```sh
+PYTHONPATH=src python -m unittest discover -s tests -v
+```

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

@@ -26,5 +26,9 @@ src/yjdata/web/_context.py
 src/yjdata/web/_input.py
 src/yjdata/web/calendar.py
 src/yjdata/web/client.py
+src/yjdata/web/cn_future.py
+src/yjdata/web/cn_stock.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.4/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.4/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()

yjdata-0.2.2/PKG-INFO

@@ -1,111 +0,0 @@
-Metadata-Version: 2.4
-Name: yjdata
-Version: 0.2.2
-Summary: yjdata SDK for Python
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-Requires-Dist: httpx>=0.25
-Requires-Dist: polars>=1.30
-Requires-Dist: pyarrow>=18
-Requires-Dist: pydantic>=2.11
-
-# yjdata
-
-Python SDK for yjdata, including MySQL and PostgreSQL database adapters.
-
-## Select the last row of each group
-
-`MysqlAdapter` (MySQL 8.0+) and `PostgresAdapter` provide:
-
-- `select_last_of_group(...) -> list[dict[str, Any]]`
-- `select_last_of_group_polars(...) -> polars.DataFrame`
-
-Both methods execute grouping, row selection, filtering, and projection in SQL
-using `ROW_NUMBER()`. They accept the table name followed by keyword arguments:
-
-| Parameter | Meaning |
-| --- | --- |
-| `table` | Table name, optionally qualified by schema or database. |
-| `group_by: Sequence[str]` | Required nonempty sequence of grouping columns. |
-| `sort_by: Sequence[str]` | Required nonempty sequence of sort columns, in priority order. |
-| `ascending: bool \| Sequence[bool] = True` | One direction for every sort column, or one boolean per column. |
-| `columns: Sequence[str] \| None = None` | Nonempty sequence of output columns; `None` returns all original columns. |
-| `filter: Expression \| None = None` | Filter source rows before choosing a row per group. |
-| `post_filter: Expression \| None = None` | Filter the chosen rows, before projecting output columns. |
-
-The operation first applies `filter`, orders each group, selects its last complete
-source row, applies `post_filter`, and finally projects `columns`. Values come
-from the same row, including NULL values; this is not a separate last-non-NULL
-aggregation for each column. A post-filter rejection removes the group entirely,
-without falling back to an older row.
-
-NULLs sort below all non-NULL values on both backends: first for ascending order
-and last for descending order. Consequently, ascending selection prefers the
-largest non-NULL value, while descending selection prefers NULL when present.
-This rule applies to each sort column. NULL grouping keys group together.
-For deterministic selection among ties, include a unique column in `sort_by`,
-such as `sort_by=["change_time", "record_id"]`. Output order across groups is
-unspecified.
-
-Group, sort, and filter columns need not appear in `columns`. Internal ranking
-columns are excluded by the SQL projection, even if the table has similarly
-named columns. With `columns=None`, an additional zero-row query reads column
-metadata so the adapter can explicitly project all original columns.
-An empty result is `[]` for dictionaries or a DataFrame with no columns for
-Polars, matching the existing selection methods. Polars infers types from all
-returned rows.
-
-### THS concept snapshots
-
-For an existing adapter `client` and a cutoff value `cutoff`:
-
-```python
-import yjdata.db.expression as ye
-
-data = client.select_last_of_group_polars(
-    table="thsetl.concept_change_record",
-    group_by=["concept_id"],
-    sort_by=["change_time"],
-    ascending=True,
-    columns=["concept_id", "concept_name"],
-    filter=(
-        (ye.col("change_time") <= cutoff)
-        & (ye.col("isvalid") == "1")
-    ),
-    post_filter=ye.col("change_type") != "conceptNegation",
-)
-```
-
-This selects the latest valid change at or before the cutoff, then excludes
-concepts whose selected change is a negation. The MySQL query has this form:
-
-```sql
-SELECT `concept_id`, `concept_name`
-FROM (
-    SELECT `concept_id`, `concept_name`, `change_type`,
-           ROW_NUMBER() OVER (
-               PARTITION BY `concept_id`
-               ORDER BY (`change_time` IS NULL) ASC, `change_time` DESC
-           ) AS `_yjdata_row_number`
-    FROM `thsetl`.`concept_change_record`
-    WHERE (`change_time` <= %s) AND (`isvalid` = %s)
-) AS `_yjdata_latest`
-WHERE `_yjdata_row_number` = 1 AND (`change_type` <> %s)
-```
-
-Values are bound in order: `[cutoff, "1", "conceptNegation"]`. PostgreSQL uses
-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.
-
-## Tests
-
-```sh
-python -m unittest discover -s tests -v
-```
-
-The grouped-selection tests exercise both adapters' validation, quoting,
-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.

yjdata-0.2.2/README.md

@@ -1,100 +0,0 @@
-# yjdata
-
-Python SDK for yjdata, including MySQL and PostgreSQL database adapters.
-
-## Select the last row of each group
-
-`MysqlAdapter` (MySQL 8.0+) and `PostgresAdapter` provide:
-
-- `select_last_of_group(...) -> list[dict[str, Any]]`
-- `select_last_of_group_polars(...) -> polars.DataFrame`
-
-Both methods execute grouping, row selection, filtering, and projection in SQL
-using `ROW_NUMBER()`. They accept the table name followed by keyword arguments:
-
-| Parameter | Meaning |
-| --- | --- |
-| `table` | Table name, optionally qualified by schema or database. |
-| `group_by: Sequence[str]` | Required nonempty sequence of grouping columns. |
-| `sort_by: Sequence[str]` | Required nonempty sequence of sort columns, in priority order. |
-| `ascending: bool \| Sequence[bool] = True` | One direction for every sort column, or one boolean per column. |
-| `columns: Sequence[str] \| None = None` | Nonempty sequence of output columns; `None` returns all original columns. |
-| `filter: Expression \| None = None` | Filter source rows before choosing a row per group. |
-| `post_filter: Expression \| None = None` | Filter the chosen rows, before projecting output columns. |
-
-The operation first applies `filter`, orders each group, selects its last complete
-source row, applies `post_filter`, and finally projects `columns`. Values come
-from the same row, including NULL values; this is not a separate last-non-NULL
-aggregation for each column. A post-filter rejection removes the group entirely,
-without falling back to an older row.
-
-NULLs sort below all non-NULL values on both backends: first for ascending order
-and last for descending order. Consequently, ascending selection prefers the
-largest non-NULL value, while descending selection prefers NULL when present.
-This rule applies to each sort column. NULL grouping keys group together.
-For deterministic selection among ties, include a unique column in `sort_by`,
-such as `sort_by=["change_time", "record_id"]`. Output order across groups is
-unspecified.
-
-Group, sort, and filter columns need not appear in `columns`. Internal ranking
-columns are excluded by the SQL projection, even if the table has similarly
-named columns. With `columns=None`, an additional zero-row query reads column
-metadata so the adapter can explicitly project all original columns.
-An empty result is `[]` for dictionaries or a DataFrame with no columns for
-Polars, matching the existing selection methods. Polars infers types from all
-returned rows.
-
-### THS concept snapshots
-
-For an existing adapter `client` and a cutoff value `cutoff`:
-
-```python
-import yjdata.db.expression as ye
-
-data = client.select_last_of_group_polars(
-    table="thsetl.concept_change_record",
-    group_by=["concept_id"],
-    sort_by=["change_time"],
-    ascending=True,
-    columns=["concept_id", "concept_name"],
-    filter=(
-        (ye.col("change_time") <= cutoff)
-        & (ye.col("isvalid") == "1")
-    ),
-    post_filter=ye.col("change_type") != "conceptNegation",
-)
-```
-
-This selects the latest valid change at or before the cutoff, then excludes
-concepts whose selected change is a negation. The MySQL query has this form:
-
-```sql
-SELECT `concept_id`, `concept_name`
-FROM (
-    SELECT `concept_id`, `concept_name`, `change_type`,
-           ROW_NUMBER() OVER (
-               PARTITION BY `concept_id`
-               ORDER BY (`change_time` IS NULL) ASC, `change_time` DESC
-           ) AS `_yjdata_row_number`
-    FROM `thsetl`.`concept_change_record`
-    WHERE (`change_time` <= %s) AND (`isvalid` = %s)
-) AS `_yjdata_latest`
-WHERE `_yjdata_row_number` = 1 AND (`change_type` <> %s)
-```
-
-Values are bound in order: `[cutoff, "1", "conceptNegation"]`. PostgreSQL uses
-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.
-
-## Tests
-
-```sh
-python -m unittest discover -s tests -v
-```
-
-The grouped-selection tests exercise both adapters' validation, quoting,
-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.

yjdata-0.2.2/src/yjdata.egg-info/PKG-INFO

@@ -1,111 +0,0 @@
-Metadata-Version: 2.4
-Name: yjdata
-Version: 0.2.2
-Summary: yjdata SDK for Python
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-Requires-Dist: httpx>=0.25
-Requires-Dist: polars>=1.30
-Requires-Dist: pyarrow>=18
-Requires-Dist: pydantic>=2.11
-
-# yjdata
-
-Python SDK for yjdata, including MySQL and PostgreSQL database adapters.
-
-## Select the last row of each group
-
-`MysqlAdapter` (MySQL 8.0+) and `PostgresAdapter` provide:
-
-- `select_last_of_group(...) -> list[dict[str, Any]]`
-- `select_last_of_group_polars(...) -> polars.DataFrame`
-
-Both methods execute grouping, row selection, filtering, and projection in SQL
-using `ROW_NUMBER()`. They accept the table name followed by keyword arguments:
-
-| Parameter | Meaning |
-| --- | --- |
-| `table` | Table name, optionally qualified by schema or database. |
-| `group_by: Sequence[str]` | Required nonempty sequence of grouping columns. |
-| `sort_by: Sequence[str]` | Required nonempty sequence of sort columns, in priority order. |
-| `ascending: bool \| Sequence[bool] = True` | One direction for every sort column, or one boolean per column. |
-| `columns: Sequence[str] \| None = None` | Nonempty sequence of output columns; `None` returns all original columns. |
-| `filter: Expression \| None = None` | Filter source rows before choosing a row per group. |
-| `post_filter: Expression \| None = None` | Filter the chosen rows, before projecting output columns. |
-
-The operation first applies `filter`, orders each group, selects its last complete
-source row, applies `post_filter`, and finally projects `columns`. Values come
-from the same row, including NULL values; this is not a separate last-non-NULL
-aggregation for each column. A post-filter rejection removes the group entirely,
-without falling back to an older row.
-
-NULLs sort below all non-NULL values on both backends: first for ascending order
-and last for descending order. Consequently, ascending selection prefers the
-largest non-NULL value, while descending selection prefers NULL when present.
-This rule applies to each sort column. NULL grouping keys group together.
-For deterministic selection among ties, include a unique column in `sort_by`,
-such as `sort_by=["change_time", "record_id"]`. Output order across groups is
-unspecified.
-
-Group, sort, and filter columns need not appear in `columns`. Internal ranking
-columns are excluded by the SQL projection, even if the table has similarly
-named columns. With `columns=None`, an additional zero-row query reads column
-metadata so the adapter can explicitly project all original columns.
-An empty result is `[]` for dictionaries or a DataFrame with no columns for
-Polars, matching the existing selection methods. Polars infers types from all
-returned rows.
-
-### THS concept snapshots
-
-For an existing adapter `client` and a cutoff value `cutoff`:
-
-```python
-import yjdata.db.expression as ye
-
-data = client.select_last_of_group_polars(
-    table="thsetl.concept_change_record",
-    group_by=["concept_id"],
-    sort_by=["change_time"],
-    ascending=True,
-    columns=["concept_id", "concept_name"],
-    filter=(
-        (ye.col("change_time") <= cutoff)
-        & (ye.col("isvalid") == "1")
-    ),
-    post_filter=ye.col("change_type") != "conceptNegation",
-)
-```
-
-This selects the latest valid change at or before the cutoff, then excludes
-concepts whose selected change is a negation. The MySQL query has this form:
-
-```sql
-SELECT `concept_id`, `concept_name`
-FROM (
-    SELECT `concept_id`, `concept_name`, `change_type`,
-           ROW_NUMBER() OVER (
-               PARTITION BY `concept_id`
-               ORDER BY (`change_time` IS NULL) ASC, `change_time` DESC
-           ) AS `_yjdata_row_number`
-    FROM `thsetl`.`concept_change_record`
-    WHERE (`change_time` <= %s) AND (`isvalid` = %s)
-) AS `_yjdata_latest`
-WHERE `_yjdata_row_number` = 1 AND (`change_type` <> %s)
-```
-
-Values are bound in order: `[cutoff, "1", "conceptNegation"]`. PostgreSQL uses
-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.
-
-## Tests
-
-```sh
-python -m unittest discover -s tests -v
-```
-
-The grouped-selection tests exercise both adapters' validation, quoting,
-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.