From 7c38832632b553bff02ce04dac4d9d2ef1f84338 Mon Sep 17 00:00:00 2001 From: David Zhang Date: Wed, 24 Sep 2025 16:52:04 +0800 Subject: [PATCH] optimize: rename all MySQL-related package names, file names, and variable names to Oceanbase --- .../README.md | 8 ++-- .../README_CN.md | 8 ++-- .../tests/conftest_checkpointer.py | 18 ++++----- .../langgraph-tests/tests/conftest_store.py | 18 ++++----- .../{mysql => oceanbase}/__init__.py | 12 +++--- .../{mysql => oceanbase}/_ainternal.py | 0 .../{mysql => oceanbase}/_internal.py | 0 .../checkpoint/{mysql => oceanbase}/aio.py | 6 +-- .../{mysql => oceanbase}/aio_base.py | 6 +-- .../{mysql => oceanbase}/asyncmy.py | 4 +- .../checkpoint/{mysql => oceanbase}/base.py | 2 +- .../checkpoint/{mysql => oceanbase}/py.typed | 0 .../pymysql.py => oceanbase/pyoceanbase.py} | 14 +++---- .../{mysql => oceanbase}/shallow.py | 8 ++-- .../checkpoint/{mysql => oceanbase}/utils.py | 0 .../langgraph/store/mysql/__init__.py | 5 --- .../langgraph/store/oceanbase/__init__.py | 5 +++ .../store/{mysql => oceanbase}/aio.py | 2 +- .../store/{mysql => oceanbase}/aio_base.py | 4 +- .../store/{mysql => oceanbase}/asyncmy.py | 2 +- .../store/{mysql => oceanbase}/base.py | 6 +-- .../store/{mysql => oceanbase}/py.typed | 0 .../pymysql.py => oceanbase/pyoceanbase.py} | 8 ++-- .../pyproject.toml | 2 +- .../tests/test_async.py | 8 ++-- .../tests/test_async_store.py | 6 +-- .../tests/test_store.py | 38 +++++++++---------- .../tests/test_sync.py | 22 +++++------ 28 files changed, 106 insertions(+), 106 deletions(-) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/__init__.py (97%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/_ainternal.py (100%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/_internal.py (100%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/aio.py (94%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/aio_base.py (99%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/asyncmy.py (96%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/base.py (99%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/py.typed (100%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql/pymysql.py => oceanbase/pyoceanbase.py} (86%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/shallow.py (99%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/{mysql => oceanbase}/utils.py (100%) delete mode 100644 langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/__init__.py create mode 100644 langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/__init__.py rename langgraph-checkpoint-oceanbase-plugin/langgraph/store/{mysql => oceanbase}/aio.py (96%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/store/{mysql => oceanbase}/aio_base.py (98%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/store/{mysql => oceanbase}/asyncmy.py (96%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/store/{mysql => oceanbase}/base.py (98%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/store/{mysql => oceanbase}/py.typed (100%) rename langgraph-checkpoint-oceanbase-plugin/langgraph/store/{mysql/pymysql.py => oceanbase/pyoceanbase.py} (84%) diff --git a/langgraph-checkpoint-oceanbase-plugin/README.md b/langgraph-checkpoint-oceanbase-plugin/README.md index 1e1ba45e..d7111b9b 100644 --- a/langgraph-checkpoint-oceanbase-plugin/README.md +++ b/langgraph-checkpoint-oceanbase-plugin/README.md @@ -24,13 +24,13 @@ In LangGraph's definition, Checkpointer refers to short-term memory (i.e. memory ```python from langchain.chat_models import init_chat_model from langgraph.graph import StateGraph, MessagesState, START -from langgraph.checkpoint.mysql.pymysql import PyMySQLSaver +from langgraph.checkpoint.oceanbase.pyoceanbase import PyOceanBaseSaver from langchain_core.runnables.config import RunnableConfig from langchain_core.messages import HumanMessage model = init_chat_model(model="qwen-max-latest", api_key="xxx", base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", model_provider="openai",temperature=0) DB_URI = "mysql://username:password@ip:port/database" -with PyMySQLSaver.from_conn_string(DB_URI) as checkpointer: +with PyOceanBaseSaver.from_conn_string(DB_URI) as checkpointer: checkpointer.setup() def call_model(state: MessagesState): @@ -63,7 +63,7 @@ with PyMySQLSaver.from_conn_string(DB_URI) as checkpointer: from langchain_core.runnables import RunnableConfig from langgraph.config import get_store from langgraph.prebuilt import create_react_agent -from langgraph.store.mysql import PyMySQLStore +from langgraph.store.oceanbase import PyOceanBaseStore from langchain.chat_models import init_chat_model from langchain_core.messages import HumanMessage from typing_extensions import TypedDict @@ -86,7 +86,7 @@ def save_user_info(user_info: UserInfo, config: RunnableConfig) -> str: return "Successfully saved user info." -with PyMySQLStore.from_conn_string(DB_URI) as store: +with PyOceanBaseStore.from_conn_string(DB_URI) as store: store.setup() agent=create_react_agent( model=model, diff --git a/langgraph-checkpoint-oceanbase-plugin/README_CN.md b/langgraph-checkpoint-oceanbase-plugin/README_CN.md index 2ae4468f..39c34872 100644 --- a/langgraph-checkpoint-oceanbase-plugin/README_CN.md +++ b/langgraph-checkpoint-oceanbase-plugin/README_CN.md @@ -24,13 +24,13 @@ langgraph-checkpoint-oceanbase 已经上传到 PyPI。可以使用下面的命 ```python from langchain.chat_models import init_chat_model from langgraph.graph import StateGraph, MessagesState, START -from langgraph.checkpoint.mysql.pymysql import PyMySQLSaver +from langgraph.checkpoint.oceanbase.pyoceanbase import PyOceanBaseSaver from langchain_core.runnables.config import RunnableConfig from langchain_core.messages import HumanMessage model = init_chat_model(model="qwen-max-latest", api_key="xxx", base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", model_provider="openai",temperature=0) DB_URI = "mysql://username:password@ip:port/database" -with PyMySQLSaver.from_conn_string(DB_URI) as checkpointer: +with PyOceanBaseSaver.from_conn_string(DB_URI) as checkpointer: checkpointer.setup() def call_model(state: MessagesState): @@ -63,7 +63,7 @@ with PyMySQLSaver.from_conn_string(DB_URI) as checkpointer: from langchain_core.runnables import RunnableConfig from langgraph.config import get_store from langgraph.prebuilt import create_react_agent -from langgraph.store.mysql import PyMySQLStore +from langgraph.store.oceanbase import PyOceanBaseStore from langchain.chat_models import init_chat_model from langchain_core.messages import HumanMessage from typing_extensions import TypedDict @@ -86,7 +86,7 @@ def save_user_info(user_info: UserInfo, config: RunnableConfig) -> str: return "Successfully saved user info." -with PyMySQLStore.from_conn_string(DB_URI) as store: +with PyOceanBaseStore.from_conn_string(DB_URI) as store: store.setup() agent=create_react_agent( model=model, diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_checkpointer.py b/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_checkpointer.py index 0bf29aaf..9c28e2ad 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_checkpointer.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_checkpointer.py @@ -6,9 +6,9 @@ import pymysql from sqlalchemy import Engine, create_engine -from langgraph.checkpoint.mysql.aio import AIOMySQLSaver -from langgraph.checkpoint.mysql.asyncmy import AsyncMySaver -from langgraph.checkpoint.mysql.pymysql import PyMySQLSaver +from langgraph.checkpoint.oceanbase.aio import AIOMySQLSaver +from langgraph.checkpoint.oceanbase.asyncmy import AsyncMySaver +from langgraph.checkpoint.oceanbase.pyoceanbase import PyOceanBaseSaver DEFAULT_MYSQL_URI = "mysql://mysql:mysql@localhost:5441/" @@ -23,19 +23,19 @@ def _checkpointer_pymysql(): database = f"test_{uuid4().hex[:16]}" # create unique db - with pymysql.connect(**PyMySQLSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE {database}") try: # yield checkpointer - with PyMySQLSaver.from_conn_string( + with PyOceanBaseSaver.from_conn_string( DEFAULT_MYSQL_URI + database ) as checkpointer: checkpointer.setup() yield checkpointer finally: # drop unique db - with pymysql.connect(**PyMySQLSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"DROP DATABASE {database}") @@ -45,17 +45,17 @@ def _checkpointer_pymysql_pool(): database = f"test_{uuid4().hex[:16]}" # create unique db - with pymysql.connect(**PyMySQLSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE {database}") try: pool = get_pymysql_sqlalchemy_engine(DEFAULT_MYSQL_URI + database) - checkpointer = PyMySQLSaver(pool.raw_connection) + checkpointer = PyOceanBaseSaver(pool.raw_connection) checkpointer.setup() yield checkpointer finally: # drop unique db - with pymysql.connect(**PyMySQLSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"DROP DATABASE {database}") diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_store.py b/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_store.py index 6927a599..5dc24941 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_store.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph-tests/tests/conftest_store.py @@ -6,9 +6,9 @@ import pymysql from sqlalchemy import Engine, create_engine -from langgraph.store.mysql.aio import AIOMySQLStore -from langgraph.store.mysql.asyncmy import AsyncMyStore -from langgraph.store.mysql.pymysql import PyMySQLStore +from langgraph.store.oceanbase.aio import AIOMySQLStore +from langgraph.store.oceanbase.asyncmy import AsyncMyStore +from langgraph.store.oceanbase.pyoceanbase import PyOceanBaseStore DEFAULT_MYSQL_URI = "mysql://mysql:mysql@localhost:5441/" @@ -23,17 +23,17 @@ def _store_pymysql(): database = f"test_{uuid4().hex[:16]}" # create unique db - with pymysql.connect(**PyMySQLStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE {database}") try: # yield store - with PyMySQLStore.from_conn_string(DEFAULT_MYSQL_URI + database) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_MYSQL_URI + database) as store: store.setup() yield store finally: # drop unique db - with pymysql.connect(**PyMySQLStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"DROP DATABASE {database}") @@ -43,18 +43,18 @@ def _store_pymysql_pool(): database = f"test_{uuid4().hex[:16]}" # create unique db - with pymysql.connect(**PyMySQLStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE {database}") try: # yield store engine = get_pymysql_sqlalchemy_engine(DEFAULT_MYSQL_URI + database) - store = PyMySQLStore(engine.raw_connection) + store = PyOceanBaseStore(engine.raw_connection) store.setup() yield store finally: # drop unique db - with pymysql.connect(**PyMySQLStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: + with pymysql.connect(**PyOceanBaseStore.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn: with conn.cursor() as cursor: cursor.execute(f"DROP DATABASE {database}") diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/__init__.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/__init__.py similarity index 97% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/__init__.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/__init__.py index c8050a66..93686964 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/__init__.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/__init__.py @@ -18,9 +18,9 @@ get_checkpoint_id, get_checkpoint_metadata, ) -from langgraph.checkpoint.mysql import _internal -from langgraph.checkpoint.mysql.base import BaseMySQLSaver -from langgraph.checkpoint.mysql.utils import ( +from langgraph.checkpoint.oceanbase import _internal +from langgraph.checkpoint.oceanbase.base import BaseMySQLSaver +from langgraph.checkpoint.oceanbase.utils import ( deserialize_channel_values, deserialize_pending_sends, deserialize_pending_writes, @@ -115,7 +115,7 @@ def list( Iterator[CheckpointTuple]: An iterator of checkpoint tuples. Examples: - >>> from langgraph.checkpoint.mysql import PyMySQLSaver + >>> from langgraph.checkpoint.oceanbase import PyOceanBaseSaver >>> DB_URI = "mysql://mysql:mysql@localhost:5432/mysql" >>> with PyMySQLSaver.from_conn_string(DB_URI) as memory: ... # Run a graph, then list the checkpoints @@ -284,9 +284,9 @@ def put( Examples: - >>> from langgraph.checkpoint.mysql import PyMySQLSaver + >>> from langgraph.checkpoint.oceanbase import PyOceanBaseSaver >>> DB_URI = "mysql://mysql:mysql@localhost:5432/mysql" - >>> with PyMySQLSaver.from_conn_string(DB_URI) as memory: + >>> with PyOceanBaseSaver.from_conn_string(DB_URI) as memory: >>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} >>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}} >>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {}) diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/_ainternal.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/_ainternal.py similarity index 100% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/_ainternal.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/_ainternal.py diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/_internal.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/_internal.py similarity index 100% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/_internal.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/_internal.py diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/aio.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/aio.py similarity index 94% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/aio.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/aio.py index 1e1fb3a9..76b87526 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/aio.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/aio.py @@ -9,9 +9,9 @@ import aiomysql # type: ignore from typing_extensions import Self, override -from langgraph.checkpoint.mysql import _ainternal -from langgraph.checkpoint.mysql.aio_base import BaseAsyncMySQLSaver -from langgraph.checkpoint.mysql.shallow import BaseShallowAsyncMySQLSaver +from langgraph.checkpoint.oceanbase import _ainternal +from langgraph.checkpoint.oceanbase.aio_base import BaseAsyncMySQLSaver +from langgraph.checkpoint.oceanbase.shallow import BaseShallowAsyncMySQLSaver from langgraph.checkpoint.serde.base import SerializerProtocol Conn = _ainternal.Conn[aiomysql.Connection] # For backward compatibility diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/aio_base.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/aio_base.py similarity index 99% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/aio_base.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/aio_base.py index ca3add33..6a279e42 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/aio_base.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/aio_base.py @@ -18,9 +18,9 @@ get_checkpoint_id, get_checkpoint_metadata, ) -from langgraph.checkpoint.mysql import _ainternal -from langgraph.checkpoint.mysql.base import BaseMySQLSaver -from langgraph.checkpoint.mysql.utils import ( +from langgraph.checkpoint.oceanbase import _ainternal +from langgraph.checkpoint.oceanbase.base import BaseMySQLSaver +from langgraph.checkpoint.oceanbase.utils import ( deserialize_channel_values, deserialize_pending_sends, deserialize_pending_writes, diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/asyncmy.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/asyncmy.py similarity index 96% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/asyncmy.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/asyncmy.py index fc413519..af503b02 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/asyncmy.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/asyncmy.py @@ -10,8 +10,8 @@ from asyncmy.cursors import DictCursor # type: ignore from typing_extensions import Self, override -from langgraph.checkpoint.mysql.aio_base import BaseAsyncMySQLSaver -from langgraph.checkpoint.mysql.shallow import BaseShallowAsyncMySQLSaver +from langgraph.checkpoint.oceanbase.aio_base import BaseAsyncMySQLSaver +from langgraph.checkpoint.oceanbase.shallow import BaseShallowAsyncMySQLSaver from langgraph.checkpoint.serde.base import SerializerProtocol diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/base.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/base.py similarity index 99% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/base.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/base.py index 15d40775..c5dd8ca5 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/base.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/base.py @@ -14,7 +14,7 @@ CheckpointMetadata, get_checkpoint_id, ) -from langgraph.checkpoint.mysql.utils import mysql_mariadb_branch +from langgraph.checkpoint.oceanbase.utils import mysql_mariadb_branch from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import TASKS diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/py.typed b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/py.typed similarity index 100% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/py.typed rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/py.typed diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/pymysql.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/pyoceanbase.py similarity index 86% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/pymysql.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/pyoceanbase.py index c03fa21f..2e7e895e 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/pymysql.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/pyoceanbase.py @@ -10,16 +10,16 @@ from pymysql.cursors import DictCursor from typing_extensions import Self, override -from langgraph.checkpoint.mysql import BaseSyncMySQLSaver, _internal -from langgraph.checkpoint.mysql import Conn as BaseConn -from langgraph.checkpoint.mysql.shallow import BaseShallowSyncMySQLSaver +from langgraph.checkpoint.oceanbase import BaseSyncMySQLSaver, _internal +from langgraph.checkpoint.oceanbase import Conn as BaseConn +from langgraph.checkpoint.oceanbase.shallow import BaseShallowSyncMySQLSaver from langgraph.checkpoint.serde.base import SerializerProtocol Conn = BaseConn[pymysql.Connection] # type: ignore -class PyMySQLSaver(BaseSyncMySQLSaver[pymysql.Connection, DictCursor]): - """Checkpointer that stores checkpoints in a MySQL database.""" +class PyOceanBaseSaver(BaseSyncMySQLSaver[pymysql.Connection, DictCursor]): + """Checkpointer that stores checkpoints in an OceanBase database.""" @staticmethod def parse_conn_string(conn_string: str) -> dict[str, Any]: @@ -100,7 +100,7 @@ def from_conn_string( conn_string=mysql://user:password@localhost/db?unix_socket=/path/to/socket """ with pymysql.connect( - **PyMySQLSaver.parse_conn_string(conn_string), + **PyOceanBaseSaver.parse_conn_string(conn_string), autocommit=True, ) as conn: yield cls(conn) @@ -111,4 +111,4 @@ def _get_cursor_from_connection(conn: pymysql.Connection) -> DictCursor: return conn.cursor(DictCursor) -__all__ = ["PyMySQLSaver", "ShallowPyMySQLSaver", "Conn"] +__all__ = ["PyOceanBaseSaver", "ShallowPyMySQLSaver", "Conn"] diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/shallow.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/shallow.py similarity index 99% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/shallow.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/shallow.py index f468ced7..16091797 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/shallow.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/shallow.py @@ -17,9 +17,9 @@ CheckpointTuple, get_checkpoint_metadata, ) -from langgraph.checkpoint.mysql import _ainternal, _internal -from langgraph.checkpoint.mysql.base import BaseMySQLSaver -from langgraph.checkpoint.mysql.utils import ( +from langgraph.checkpoint.oceanbase import _ainternal, _internal +from langgraph.checkpoint.oceanbase.base import BaseMySQLSaver +from langgraph.checkpoint.oceanbase.utils import ( deserialize_channel_values, deserialize_pending_sends, deserialize_pending_writes, @@ -432,7 +432,7 @@ def put( Examples: - >>> from langgraph.checkpoint.mysql import PyMySQLSaver + >>> from langgraph.checkpoint.oceanbase import PyOceanBaseSaver >>> DB_URI = "mysql://mysql:mysql@localhost:5432/mysql" >>> with ShallowPyMySQLSaver.from_conn_string(DB_URI) as memory: >>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/utils.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/utils.py similarity index 100% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/mysql/utils.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/checkpoint/oceanbase/utils.py diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/__init__.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/__init__.py deleted file mode 100644 index e5418195..00000000 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from langgraph.store.mysql.aio import AIOMySQLStore -from langgraph.store.mysql.asyncmy import AsyncMyStore -from langgraph.store.mysql.pymysql import PyMySQLStore - -__all__ = ["AIOMySQLStore", "AsyncMyStore", "PyMySQLStore"] diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/__init__.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/__init__.py new file mode 100644 index 00000000..d1aff4ef --- /dev/null +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/__init__.py @@ -0,0 +1,5 @@ +from langgraph.store.oceanbase.aio import AIOMySQLStore +from langgraph.store.oceanbase.asyncmy import AsyncMyStore +from langgraph.store.oceanbase.pyoceanbase import PyOceanBaseStore + +__all__ = ["AIOMySQLStore", "AsyncMyStore", "PyOceanBaseStore"] diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/aio.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/aio.py similarity index 96% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/aio.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/aio.py index 1959322b..79534b58 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/aio.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/aio.py @@ -7,7 +7,7 @@ import aiomysql # type: ignore from typing_extensions import Self, override -from langgraph.store.mysql.aio_base import BaseAsyncMySQLStore +from langgraph.store.oceanbase.aio_base import BaseAsyncMySQLStore logger = logging.getLogger(__name__) diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/aio_base.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/aio_base.py similarity index 98% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/aio_base.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/aio_base.py index 0b08f6c9..47060518 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/aio_base.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/aio_base.py @@ -8,7 +8,7 @@ import orjson -from langgraph.checkpoint.mysql import _ainternal +from langgraph.checkpoint.oceanbase import _ainternal from langgraph.store.base import ( GetOp, ListNamespacesOp, @@ -18,7 +18,7 @@ SearchOp, ) from langgraph.store.base.batch import AsyncBatchedBaseStore -from langgraph.store.mysql.base import ( +from langgraph.store.oceanbase.base import ( BaseMySQLStore, Row, _decode_ns_bytes, diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/asyncmy.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/asyncmy.py similarity index 96% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/asyncmy.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/asyncmy.py index 8ff1e3fe..287ec81a 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/asyncmy.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/asyncmy.py @@ -8,7 +8,7 @@ from asyncmy.cursors import DictCursor # type: ignore from typing_extensions import Self, override -from langgraph.store.mysql.aio_base import BaseAsyncMySQLStore +from langgraph.store.oceanbase.aio_base import BaseAsyncMySQLStore logger = logging.getLogger(__name__) diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/base.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/base.py similarity index 98% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/base.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/base.py index 813d7d68..027f70ea 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/base.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/base.py @@ -20,9 +20,9 @@ import orjson from typing_extensions import TypedDict -from langgraph.checkpoint.mysql import _ainternal as _ainternal -from langgraph.checkpoint.mysql import _internal as _internal -from langgraph.checkpoint.mysql.utils import mysql_mariadb_branch +from langgraph.checkpoint.oceanbase import _ainternal as _ainternal +from langgraph.checkpoint.oceanbase import _internal as _internal +from langgraph.checkpoint.oceanbase.utils import mysql_mariadb_branch from langgraph.store.base import ( BaseStore, GetOp, diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/py.typed b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/py.typed similarity index 100% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/py.typed rename to langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/py.typed diff --git a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/pymysql.py b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/pyoceanbase.py similarity index 84% rename from langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/pymysql.py rename to langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/pyoceanbase.py index 3dab6f07..34253b7f 100644 --- a/langgraph-checkpoint-oceanbase-plugin/langgraph/store/mysql/pymysql.py +++ b/langgraph-checkpoint-oceanbase-plugin/langgraph/store/oceanbase/pyoceanbase.py @@ -7,10 +7,10 @@ from pymysql.cursors import DictCursor from typing_extensions import Self, override -from langgraph.store.mysql.base import BaseSyncMySQLStore +from langgraph.store.oceanbase.base import BaseSyncMySQLStore -class PyMySQLStore(BaseSyncMySQLStore[pymysql.Connection, DictCursor]): +class PyOceanBaseStore(BaseSyncMySQLStore[pymysql.Connection, DictCursor]): @staticmethod def parse_conn_string(conn_string: str) -> dict[str, Any]: parsed = urllib.parse.urlparse(conn_string) @@ -35,13 +35,13 @@ def from_conn_string( cls, conn_string: str, ) -> Iterator[Self]: - """Create a new PyMySQLStore instance from a connection string. + """Create a new PyOceanBaseStore instance from a connection string. Args: conn_string: The MySQL connection info string. Returns: - PyMySQLStore: A new PyMySQLStore instance. + PyOceanBaseStore: A new PyOceanBaseStore instance. """ with pymysql.connect( **cls.parse_conn_string(conn_string), diff --git a/langgraph-checkpoint-oceanbase-plugin/pyproject.toml b/langgraph-checkpoint-oceanbase-plugin/pyproject.toml index f5ff4378..29b739cf 100644 --- a/langgraph-checkpoint-oceanbase-plugin/pyproject.toml +++ b/langgraph-checkpoint-oceanbase-plugin/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "langgraph-checkpoint-oceanbase" -version = "2.0.17" +version = "2.0.18" description = "Library with a OceanBase MySQL Mode implementation of LangGraph checkpoint saver." authors = [{ name = "David Zhang" }] requires-python = ">=3.9" diff --git a/langgraph-checkpoint-oceanbase-plugin/tests/test_async.py b/langgraph-checkpoint-oceanbase-plugin/tests/test_async.py index 34f0f1a3..23a4558d 100644 --- a/langgraph-checkpoint-oceanbase-plugin/tests/test_async.py +++ b/langgraph-checkpoint-oceanbase-plugin/tests/test_async.py @@ -20,10 +20,10 @@ create_checkpoint, empty_checkpoint, ) -from langgraph.checkpoint.mysql.aio import AIOMySQLSaver, ShallowAIOMySQLSaver -from langgraph.checkpoint.mysql.aio_base import BaseAsyncMySQLSaver -from langgraph.checkpoint.mysql.asyncmy import AsyncMySaver, ShallowAsyncMySaver -from langgraph.checkpoint.mysql.shallow import BaseShallowAsyncMySQLSaver +from langgraph.checkpoint.oceanbase.aio import AIOMySQLSaver, ShallowAIOMySQLSaver +from langgraph.checkpoint.oceanbase.aio_base import BaseAsyncMySQLSaver +from langgraph.checkpoint.oceanbase.asyncmy import AsyncMySaver, ShallowAsyncMySaver +from langgraph.checkpoint.oceanbase.shallow import BaseShallowAsyncMySQLSaver from langgraph.checkpoint.serde.types import TASKS from langgraph.graph import END, START, MessagesState, StateGraph from tests.conftest import DEFAULT_BASE_URI diff --git a/langgraph-checkpoint-oceanbase-plugin/tests/test_async_store.py b/langgraph-checkpoint-oceanbase-plugin/tests/test_async_store.py index 13d326d0..95397b3a 100644 --- a/langgraph-checkpoint-oceanbase-plugin/tests/test_async_store.py +++ b/langgraph-checkpoint-oceanbase-plugin/tests/test_async_store.py @@ -10,9 +10,9 @@ import pytest from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp -from langgraph.store.mysql.aio import AIOMySQLStore -from langgraph.store.mysql.aio_base import BaseAsyncMySQLStore -from langgraph.store.mysql.asyncmy import AsyncMyStore +from langgraph.store.oceanbase.aio import AIOMySQLStore +from langgraph.store.oceanbase.aio_base import BaseAsyncMySQLStore +from langgraph.store.oceanbase.asyncmy import AsyncMyStore from tests.conftest import DEFAULT_BASE_URI pytestmark = pytest.mark.anyio diff --git a/langgraph-checkpoint-oceanbase-plugin/tests/test_store.py b/langgraph-checkpoint-oceanbase-plugin/tests/test_store.py index ffd11a9d..c930e76a 100644 --- a/langgraph-checkpoint-oceanbase-plugin/tests/test_store.py +++ b/langgraph-checkpoint-oceanbase-plugin/tests/test_store.py @@ -15,7 +15,7 @@ PutOp, SearchOp, ) -from langgraph.store.mysql import PyMySQLStore +from langgraph.store.oceanbase import PyOceanBasetore from tests.conftest import ( DEFAULT_BASE_URI, DEFAULT_URI, @@ -33,38 +33,38 @@ @pytest.fixture(scope="function", params=STORES) -def store(request) -> PyMySQLStore: +def store(request) -> PyOceanBaseStore: database = f"test_{uuid4().hex[:16]}" with pymysql.connect( - **PyMySQLStore.parse_conn_string(DEFAULT_BASE_URI), + **PyOceanBaseStore.parse_conn_string(DEFAULT_BASE_URI), autocommit=True, ) as conn: with conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE {database}") try: - with PyMySQLStore.from_conn_string(DEFAULT_BASE_URI + database) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_BASE_URI + database) as store: store.setup() if request.param == "sqlalchemy_engine": engine = get_pymysql_sqlalchemy_engine(DEFAULT_BASE_URI + database) - yield PyMySQLStore(engine.raw_connection) + yield PyOceanBaseStore(engine.raw_connection) elif request.param == "sqlalchemy_pool": pool = get_pymysql_sqlalchemy_pool(DEFAULT_BASE_URI + database) - yield PyMySQLStore(pool.connect) + yield PyOceanBaseStore(pool.connect) else: # default - with PyMySQLStore.from_conn_string(DEFAULT_BASE_URI + database) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_BASE_URI + database) as store: yield store finally: with pymysql.connect( - **PyMySQLStore.parse_conn_string(DEFAULT_BASE_URI), + **PyOceanBaseStore.parse_conn_string(DEFAULT_BASE_URI), autocommit=True, ) as conn: with conn.cursor() as cursor: cursor.execute(f"DROP DATABASE {database}") -def test_batch_order(store: PyMySQLStore) -> None: +def test_batch_order(store: PyOceanBaseStore) -> None: # Setup test data store.put(("test", "foo"), "key1", {"data": "value1"}) store.put(("test", "bar"), "key2", {"data": "value2"}) @@ -116,7 +116,7 @@ def test_batch_order(store: PyMySQLStore) -> None: assert results_reordered[4].key == "key1" -def test_batch_get_ops(store: PyMySQLStore) -> None: +def test_batch_get_ops(store: PyOceanBaseStore) -> None: # Setup test data store.put(("test",), "key1", {"data": "value1"}) store.put(("test",), "key2", {"data": "value2"}) @@ -137,7 +137,7 @@ def test_batch_get_ops(store: PyMySQLStore) -> None: assert results[1].key == "key2" -def test_batch_put_ops(store: PyMySQLStore) -> None: +def test_batch_put_ops(store: PyOceanBaseStore) -> None: ops = [ PutOp(namespace=("test",), key="key1", value={"data": "value1"}), PutOp(namespace=("test",), key="key2", value={"data": "value2"}), @@ -158,7 +158,7 @@ def test_batch_put_ops(store: PyMySQLStore) -> None: assert item3 is None -def test_batch_search_ops(store: PyMySQLStore) -> None: +def test_batch_search_ops(store: PyOceanBaseStore) -> None: # Setup test data test_data = [ (("test", "foo"), "key1", {"data": "value1", "tag": "a"}), @@ -189,7 +189,7 @@ def test_batch_search_ops(store: PyMySQLStore) -> None: assert results[2][0].namespace == ("test", "foo") -def test_batch_list_namespaces_ops(store: PyMySQLStore) -> None: +def test_batch_list_namespaces_ops(store: PyOceanBaseStore) -> None: # Setup test data with various namespaces test_data = [ (("test", "documents", "public"), "doc1", {"content": "public doc"}), @@ -224,14 +224,14 @@ def test_batch_list_namespaces_ops(store: PyMySQLStore) -> None: assert all(ns[-1] == "public" for ns in results[2]) -class TestPyMySQLStore: +class TestPyOceanBaseStore: @pytest.fixture(autouse=True) def setup(self) -> None: - with PyMySQLStore.from_conn_string(DEFAULT_URI) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_URI) as store: store.setup() def test_basic_store_ops(self) -> None: - with PyMySQLStore.from_conn_string(DEFAULT_URI) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_URI) as store: namespace = ("test", "documents") item_id = "doc1" item_value = {"title": "Test Document", "content": "Hello, World!"} @@ -264,7 +264,7 @@ def test_basic_store_ops(self) -> None: assert deleted_item is None def test_list_namespaces(self) -> None: - with PyMySQLStore.from_conn_string(DEFAULT_URI) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_URI) as store: # Create test data with various namespaces test_namespaces = [ ("test", "documents", "public"), @@ -306,7 +306,7 @@ def test_list_namespaces(self) -> None: store.delete(namespace, "dummy") def test_search(self) -> None: - with PyMySQLStore.from_conn_string(DEFAULT_URI) as store: + with PyOceanBaseStore.from_conn_string(DEFAULT_URI) as store: # Create test data test_data = [ ( @@ -357,6 +357,6 @@ def test_search(self) -> None: def test_nonnull_migrations() -> None: _leading_comment_remover = re.compile(r"^/\*.*?\*/") - for migration in PyMySQLStore.MIGRATIONS: + for migration in PyOceanBaseStore.MIGRATIONS: statement = _leading_comment_remover.sub("", migration).split()[0] assert statement.strip() diff --git a/langgraph-checkpoint-oceanbase-plugin/tests/test_sync.py b/langgraph-checkpoint-oceanbase-plugin/tests/test_sync.py index 7eefa9c9..e46746db 100644 --- a/langgraph-checkpoint-oceanbase-plugin/tests/test_sync.py +++ b/langgraph-checkpoint-oceanbase-plugin/tests/test_sync.py @@ -19,7 +19,7 @@ create_checkpoint, empty_checkpoint, ) -from langgraph.checkpoint.mysql.pymysql import PyMySQLSaver, ShallowPyMySQLSaver +from langgraph.checkpoint.oceanbase.pyoceanbase import PyOceanBaseSaver, ShallowPyMySQLSaver from langgraph.checkpoint.serde.types import TASKS from tests.conftest import ( DEFAULT_BASE_URI, @@ -49,7 +49,7 @@ def _database() -> Iterator[str]: # create unique db with pymysql.connect( - **PyMySQLSaver.parse_conn_string(DEFAULT_BASE_URI), autocommit=True + **PyOceanBaseSaver.parse_conn_string(DEFAULT_BASE_URI), autocommit=True ) as conn: with conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE {database}") @@ -58,35 +58,35 @@ def _database() -> Iterator[str]: finally: # drop unique db with pymysql.connect( - **PyMySQLSaver.parse_conn_string(DEFAULT_BASE_URI), autocommit=True + **PyOceanBaseSaver.parse_conn_string(DEFAULT_BASE_URI), autocommit=True ) as conn: with conn.cursor() as cursor: cursor.execute(f"DROP DATABASE {database}") @contextmanager -def _base_saver() -> Iterator[PyMySQLSaver]: +def _base_saver() -> Iterator[PyOceanBaseSaver]: with _database() as database: - with PyMySQLSaver.from_conn_string(DEFAULT_BASE_URI + database) as checkpointer: + with PyOceanBaseSaver.from_conn_string(DEFAULT_BASE_URI + database) as checkpointer: yield checkpointer @contextmanager -def _sqlalchemy_engine_saver() -> Iterator[PyMySQLSaver]: +def _sqlalchemy_engine_saver() -> Iterator[PyOceanBaseSaver]: with _database() as database: engine = get_pymysql_sqlalchemy_engine(DEFAULT_BASE_URI + database) try: - yield PyMySQLSaver(engine.raw_connection) + yield PyOceanBaseSaver(engine.raw_connection) finally: engine.dispose() @contextmanager -def _sqlalchemy_pool_saver() -> Iterator[PyMySQLSaver]: +def _sqlalchemy_pool_saver() -> Iterator[PyOceanBaseSaver]: with _database() as database: pool = get_pymysql_sqlalchemy_pool(DEFAULT_BASE_URI + database) try: - yield PyMySQLSaver(pool.connect) + yield PyOceanBaseSaver(pool.connect) finally: pool.dispose() @@ -103,7 +103,7 @@ def _shallow_saver() -> Iterator[ShallowPyMySQLSaver]: @contextmanager -def _saver(name: str) -> Iterator[PyMySQLSaver | ShallowPyMySQLSaver]: +def _saver(name: str) -> Iterator[PyOceanBaseSaver | ShallowPyMySQLSaver]: if name == "base": factory = _base_saver elif name == "shallow": @@ -408,7 +408,7 @@ def test_write_with_same_checkpoint_ns_updates(saver_name: str) -> None: def test_nonnull_migrations() -> None: _leading_comment_remover = re.compile(r"^/\*.*?\*/") - for migration in PyMySQLSaver.MIGRATIONS: + for migration in PyOceanBaseSaver.MIGRATIONS: statement = _leading_comment_remover.sub("", migration).split()[0] assert statement.strip()