From 65161d96c650980acff14abdd633a036f1e9090c Mon Sep 17 00:00:00 2001
From: Theodore Ni <3806110+tjni@users.noreply.github.com>
Date: Fri, 20 Dec 2024 20:26:52 -0800
Subject: [PATCH] Remove checkpoint_ns from primary keys.

At the same time, increase its max size to 2000 characters. We do this
because with nested subgraphs, we hit the original 100 character limit
way too quickly.

We have to remove it from the primary key because primary keys are
limited to 767 bytes in MySQL.
---
 langgraph/checkpoint/mysql/base.py | 18 ++++++++++++++++++
 pyproject.toml                     |  2 +-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/langgraph/checkpoint/mysql/base.py b/langgraph/checkpoint/mysql/base.py
index 8eb7c5e..727ee52 100644
--- a/langgraph/checkpoint/mysql/base.py
+++ b/langgraph/checkpoint/mysql/base.py
@@ -72,6 +72,24 @@
     "ALTER TABLE checkpoints MODIFY COLUMN `checkpoint_ns` VARCHAR(255) NOT NULL DEFAULT '';",
     "ALTER TABLE checkpoint_blobs MODIFY COLUMN `checkpoint_ns` VARCHAR(255) NOT NULL DEFAULT '';",
     "ALTER TABLE checkpoint_writes MODIFY COLUMN `checkpoint_ns` VARCHAR(255) NOT NULL DEFAULT '';",
+    """
+    ALTER TABLE checkpoints
+    DROP PRIMARY KEY,
+    ADD PRIMARY KEY (thread_id, checkpoint_id),
+    MODIFY COLUMN `checkpoint_ns` VARCHAR(2000) NOT NULL DEFAULT '';
+    """,
+    """
+    ALTER TABLE checkpoint_blobs
+    DROP PRIMARY KEY,
+    ADD PRIMARY KEY (thread_id, channel, version),
+    MODIFY COLUMN `checkpoint_ns` VARCHAR(2000) NOT NULL DEFAULT '';
+    """,
+    """
+    ALTER TABLE checkpoint_writes
+    DROP PRIMARY KEY,
+    ADD PRIMARY KEY (thread_id, checkpoint_id, task_id, idx),
+    MODIFY COLUMN `checkpoint_ns` VARCHAR(2000) NOT NULL DEFAULT '';
+    """,
 ]
 
 SELECT_SQL = f"""
diff --git a/pyproject.toml b/pyproject.toml
index a357dfd..a118b65 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-mysql"
-version = "2.0.8"
+version = "2.0.9"
 description = "Library with a MySQL implementation of LangGraph checkpoint saver."
 authors = ["Theodore Ni <dev@ted.bio>"]
 license = "MIT"