nvidia-nat-mysql 1.3.dev0__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/plugins/mysql/mysql_object_store.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import logging
-import pickle
+import re
 
 import aiomysql
 from aiomysql.pool import Pool
@@ -23,7 +23,6 @@ from nat.data_models.object_store import KeyAlreadyExistsError
 from nat.data_models.object_store import NoSuchKeyError
 from nat.object_store.interfaces import ObjectStore
 from nat.object_store.models import ObjectStoreItem
-from nat.plugins.mysql.object_store import MySQLObjectStoreClientConfig
 from nat.utils.type_utils import override
 
 logger = logging.getLogger(__name__)
@@ -34,66 +33,78 @@ class MySQLObjectStore(ObjectStore):
     Implementation of ObjectStore that stores objects in a MySQL database.
     """
 
-    def __init__(self, config: MySQLObjectStoreClientConfig):
-
+    def __init__(self, *, bucket_name: str, host: str, port: int, username: str | None, password: str | None):
         super().__init__()
 
-        self._config = config
+        if not re.fullmatch(r"[A-Za-z0-9_-]+", bucket_name):
+            raise ValueError("bucket_name must match [A-Za-z0-9_-]+")
+
+        self._bucket_name = bucket_name
+        self._host = host
+        self._port = port
+        self._username = username
+        self._password = password
+
         self._conn_pool: Pool | None = None
 
-        self._schema = f"`bucket_{self._config.bucket_name}`"
+    @property
+    def _schema(self) -> str:
+        return f"`bucket_{self._bucket_name}`"
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "MySQLObjectStore":
 
         if self._conn_pool is not None:
             raise RuntimeError("Connection already established")
 
         self._conn_pool = await aiomysql.create_pool(
-            host=self._config.host,
-            port=self._config.port,
-            user=self._config.username,
-            password=self._config.password,
+            host=self._host,
+            port=self._port,
+            user=self._username,
+            password=self._password,
             autocommit=False,  # disable autocommit for transactions
         )
         assert self._conn_pool is not None
 
-        logger.info(
-            f"Created connection pool for {self._config.bucket_name} at {self._config.host}:{self._config.port}")
+        logger.info("Created connection pool for %s at %s:%s", self._bucket_name, self._host, self._port)
 
         async with self._conn_pool.acquire() as conn:
             async with conn.cursor() as cur:
 
-                # Create schema (database) if doesn't exist
-                await cur.execute(f"CREATE SCHEMA IF NOT EXISTS {self._schema} DEFAULT CHARACTER SET utf8mb4;")
-                await cur.execute(f"USE {self._schema};")
+                # Suppress MySQL "IF NOT EXISTS" notes that surface as warnings in the driver
+                await cur.execute("SET sql_notes = 0;")
+                try:
+                    # Create schema (database) if doesn't exist
+                    await cur.execute(f"CREATE SCHEMA IF NOT EXISTS {self._schema} DEFAULT CHARACTER SET utf8mb4;")
+                    await cur.execute(f"USE {self._schema};")
 
-                # Create metadata table_schema
-                await cur.execute("""
-                CREATE TABLE IF NOT EXISTS object_meta (
-                    id INT AUTO_INCREMENT PRIMARY KEY,
-                    path VARCHAR(768) NOT NULL UNIQUE,
-                    size BIGINT NOT NULL,
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-                ) ENGINE=InnoDB;
-                """)
-
-                # Create blob data table
-                await cur.execute("""
-                CREATE TABLE IF NOT EXISTS object_data (
-                    id INT PRIMARY KEY,
-                    data LONGBLOB NOT NULL,
-                    FOREIGN KEY (id) REFERENCES object_meta(id) ON DELETE CASCADE
-                ) ENGINE=InnoDB ROW_FORMAT=DYNAMIC;
-                """)
+                    # Create metadata table_schema
+                    await cur.execute("""
+                    CREATE TABLE IF NOT EXISTS object_meta (
+                        id INT AUTO_INCREMENT PRIMARY KEY,
+                        path VARCHAR(768) NOT NULL UNIQUE,
+                        size BIGINT NOT NULL,
+                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                    ) ENGINE=InnoDB;
+                    """)
+
+                    # Create blob data table
+                    await cur.execute("""
+                    CREATE TABLE IF NOT EXISTS object_data (
+                        id INT PRIMARY KEY,
+                        data LONGBLOB NOT NULL,
+                        FOREIGN KEY (id) REFERENCES object_meta(id) ON DELETE CASCADE
+                    ) ENGINE=InnoDB ROW_FORMAT=DYNAMIC;
+                    """)
+                finally:
+                    await cur.execute("SET sql_notes = 1;")
 
             await conn.commit()
 
-        logger.info(
-            f"Created schema and tables for {self._config.bucket_name} at {self._config.host}:{self._config.port}")
+        logger.info("Created schema and tables for %s at %s:%s", self._bucket_name, self._host, self._port)
 
         return self
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
 
         if not self._conn_pool:
             raise RuntimeError("Connection not established")
@@ -120,11 +131,10 @@ class MySQLObjectStore(ObjectStore):
                                       (key, len(item.data)))
                     if cur.rowcount == 0:
                         raise KeyAlreadyExistsError(
-                            key=key, additional_message=f"MySQL table {self._config.bucket_name} already has key {key}")
+                            key=key, additional_message=f"MySQL table {self._bucket_name} already has key {key}")
                     await cur.execute("SELECT id FROM object_meta WHERE path=%s FOR UPDATE;", (key, ))
                     (obj_id, ) = await cur.fetchone()
-
-                    blob = pickle.dumps(item)
+                    blob = item.model_dump_json()
                     await cur.execute("INSERT INTO object_data (id, data) VALUES (%s, %s)", (obj_id, blob))
                     await conn.commit()
                 except Exception:
@@ -145,13 +155,13 @@ class MySQLObjectStore(ObjectStore):
                     await cur.execute(
                         """
                         INSERT INTO object_meta (path, size)
-                        VALUES (%s, %s)
-                        ON DUPLICATE KEY UPDATE size=VALUES(size), created_at=CURRENT_TIMESTAMP
+                        VALUES (%s, %s) AS new
+                        ON DUPLICATE KEY UPDATE size=new.size, created_at=CURRENT_TIMESTAMP
                         """, (key, len(item.data)))
                     await cur.execute("SELECT id FROM object_meta WHERE path=%s FOR UPDATE;", (key, ))
                     (obj_id, ) = await cur.fetchone()
 
-                    blob = pickle.dumps(item)
+                    blob = item.model_dump_json()
                     await cur.execute("REPLACE INTO object_data (id, data) VALUES (%s, %s)", (obj_id, blob))
                     await conn.commit()
                 except Exception:
@@ -176,9 +186,9 @@ class MySQLObjectStore(ObjectStore):
                 """, (key, ))
                 row = await cur.fetchone()
                 if not row:
-                    raise NoSuchKeyError(
-                        key=key, additional_message=f"MySQL table {self._config.bucket_name} does not have key {key}")
-                return pickle.loads(row[0])
+                    raise NoSuchKeyError(key=key,
+                                         additional_message=f"MySQL table {self._bucket_name} does not have key {key}")
+                return ObjectStoreItem.model_validate_json(row[0].decode("utf-8"))
 
     @override
     async def delete_object(self, key: str):
@@ -200,8 +210,7 @@ class MySQLObjectStore(ObjectStore):
 
                     if cur.rowcount == 0:
                         raise NoSuchKeyError(
-                            key=key,
-                            additional_message=f"MySQL table {self._config.bucket_name} does not have key {key}")
+                            key=key, additional_message=f"MySQL table {self._bucket_name} does not have key {key}")
 
                     await conn.commit()
                 except Exception:

nat/plugins/mysql/object_store.py

@@ -58,9 +58,9 @@ class MySQLObjectStoreClientConfig(ObjectStoreBaseConfig, name="mysql"):
 
 
 @register_object_store(config_type=MySQLObjectStoreClientConfig)
-async def mysql_object_store_client(config: MySQLObjectStoreClientConfig, builder: Builder):
+async def mysql_object_store_client(config: MySQLObjectStoreClientConfig, _builder: Builder):
 
     from .mysql_object_store import MySQLObjectStore
 
-    async with MySQLObjectStore(config) as store:
+    async with MySQLObjectStore(**config.model_dump(exclude={"type"}, exclude_none=True)) as store:
         yield store

nat/plugins/mysql/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

{nvidia_nat_mysql-1.3.dev0.dist-info → nvidia_nat_mysql-1.3.0rc1.dist-info}/METADATA

@@ -1,10 +1,13 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-mysql
-Version: 1.3.dev0
+Version: 1.3.0rc1
 Summary: Subpackage for MySQL integration in NeMo Agent toolkit
 Keywords: ai,agents,memory,data store
 Classifier: Programming Language :: Python
-Requires-Python: >=3.12
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3-dev
+Requires-Dist: nvidia-nat==v1.3.0-rc1
 Requires-Dist: aiomysql>=0.2.0

nvidia_nat_mysql-1.3.0rc1.dist-info/RECORD

@@ -0,0 +1,9 @@
+nat/plugins/mysql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/mysql/mysql_object_store.py,sha256=OhOOCXMBcEBITfTt1mqpCAcDUgwIc0_HzkcTKO1Kezc,8602
+nat/plugins/mysql/object_store.py,sha256=iEkWrlooEXqAfJYdg39cKe9EMCvQJ6CqCeeoL_cmOrw,2651
+nat/plugins/mysql/register.py,sha256=iBGq69m0rbZ3BTisA4Lt5UCHtW0Pc7m-l8G8dz-_pBc,813
+nvidia_nat_mysql-1.3.0rc1.dist-info/METADATA,sha256=g9STcuk1JaCUmdhgrWO1epoJ89NQ5DjAXEzgsV3JwaQ,500
+nvidia_nat_mysql-1.3.0rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_mysql-1.3.0rc1.dist-info/entry_points.txt,sha256=ZI7ielmDX-k_OPXGvpRB5tKFVR5kMCHmHGwLWqxRKh0,56
+nvidia_nat_mysql-1.3.0rc1.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_mysql-1.3.0rc1.dist-info/RECORD,,

nvidia_nat_mysql-1.3.dev0.dist-info/RECORD

@@ -1,9 +0,0 @@
-nat/plugins/mysql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nat/plugins/mysql/mysql_object_store.py,sha256=EXHuA6z28__lH6qhSyGuIPWLPZ-5BzH9EHVCkpDYyvE,8042
-nat/plugins/mysql/object_store.py,sha256=IWMJFKBhn0lH7UGHuYkYht82TwIe5KXS7VUhQrlNRyI,2600
-nat/plugins/mysql/register.py,sha256=7gqnwyDrYttIlEaa7lo9AASYt-2GrZJE0YT2jpKjepo,845
-nvidia_nat_mysql-1.3.dev0.dist-info/METADATA,sha256=g11nXToNcwBEvbspx29N-29LrI_sGnp3Z0h2z2Wh7XI,339
-nvidia_nat_mysql-1.3.dev0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_mysql-1.3.dev0.dist-info/entry_points.txt,sha256=ZI7ielmDX-k_OPXGvpRB5tKFVR5kMCHmHGwLWqxRKh0,56
-nvidia_nat_mysql-1.3.dev0.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_mysql-1.3.dev0.dist-info/RECORD,,