nvidia-nat-mysql 1.2.0a20250813__py3-none-any.whl

nat/plugins/mysql/mysql_object_store.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import pickle
+
+import aiomysql
+from aiomysql.pool import Pool
+
+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__)
+
+
+class MySQLObjectStore(ObjectStore):
+    """
+    Implementation of ObjectStore that stores objects in a MySQL database.
+    """
+
+    def __init__(self, config: MySQLObjectStoreClientConfig):
+
+        super().__init__()
+
+        self._config = config
+        self._conn_pool: Pool | None = None
+
+        self._schema = f"`bucket_{self._config.bucket_name}`"
+
+    async def __aenter__(self):
+
+        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,
+            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}")
+
+        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};")
+
+                # 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;
+                """)
+
+            await conn.commit()
+
+        logger.info(
+            f"Created schema and tables for {self._config.bucket_name} at {self._config.host}:{self._config.port}")
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+
+        if not self._conn_pool:
+            raise RuntimeError("Connection not established")
+
+        # Trigger the non-async close method then wait for the pool to close
+        self._conn_pool.close()
+
+        await self._conn_pool.wait_closed()
+
+        self._conn_pool = None
+
+    @override
+    async def put_object(self, key: str, item: ObjectStoreItem):
+
+        if not self._conn_pool:
+            raise RuntimeError("Connection not established")
+
+        async with self._conn_pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(f"USE {self._schema};")
+                try:
+                    await cur.execute("START TRANSACTION;")
+                    await cur.execute("INSERT IGNORE INTO object_meta (path, size) VALUES (%s, %s)",
+                                      (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}")
+                    await cur.execute("SELECT id FROM object_meta WHERE path=%s FOR UPDATE;", (key, ))
+                    (obj_id, ) = await cur.fetchone()
+
+                    blob = pickle.dumps(item)
+                    await cur.execute("INSERT INTO object_data (id, data) VALUES (%s, %s)", (obj_id, blob))
+                    await conn.commit()
+                except Exception:
+                    await conn.rollback()
+                    raise
+
+    @override
+    async def upsert_object(self, key: str, item: ObjectStoreItem):
+
+        if not self._conn_pool:
+            raise RuntimeError("Connection not established")
+
+        async with self._conn_pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(f"USE {self._schema};")
+                try:
+                    await cur.execute("START TRANSACTION;")
+                    await cur.execute(
+                        """
+                        INSERT INTO object_meta (path, size)
+                        VALUES (%s, %s)
+                        ON DUPLICATE KEY UPDATE size=VALUES(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)
+                    await cur.execute("REPLACE INTO object_data (id, data) VALUES (%s, %s)", (obj_id, blob))
+                    await conn.commit()
+                except Exception:
+                    await conn.rollback()
+                    raise
+
+    @override
+    async def get_object(self, key: str) -> ObjectStoreItem:
+
+        if not self._conn_pool:
+            raise RuntimeError("Connection not established")
+
+        async with self._conn_pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(f"USE {self._schema};")
+                await cur.execute(
+                    """
+                    SELECT d.data
+                    FROM object_data d
+                    JOIN object_meta m USING(id)
+                    WHERE m.path=%s
+                """, (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])
+
+    @override
+    async def delete_object(self, key: str):
+
+        if not self._conn_pool:
+            raise RuntimeError("Connection not established")
+
+        async with self._conn_pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                try:
+                    await cur.execute(f"USE {self._schema};")
+                    await cur.execute(
+                        """
+                        DELETE m, d
+                        FROM object_meta m
+                        JOIN object_data d USING(id)
+                        WHERE m.path=%s
+                    """, (key, ))
+
+                    if cur.rowcount == 0:
+                        raise NoSuchKeyError(
+                            key=key,
+                            additional_message=f"MySQL table {self._config.bucket_name} does not have key {key}")
+
+                    await conn.commit()
+                except Exception:
+                    await conn.rollback()
+                    raise

nat/plugins/mysql/object_store.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import ClassVar
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_object_store
+from nat.data_models.object_store import ObjectStoreBaseConfig
+
+
+class MySQLObjectStoreClientConfig(ObjectStoreBaseConfig, name="mysql"):
+    """
+    Object store that stores objects in a MySQL database.
+    """
+
+    DEFAULT_HOST: ClassVar[str] = "localhost"
+    DEFAULT_PORT: ClassVar[int] = 3306
+
+    HOST_ENV: ClassVar[str] = "AIQ_MYSQL_OBJECT_STORE_HOST"
+    PORT_ENV: ClassVar[str] = "AIQ_MYSQL_OBJECT_STORE_PORT"
+    USERNAME_ENV: ClassVar[str] = "AIQ_MYSQL_OBJECT_STORE_USERNAME"
+    PASSWORD_ENV: ClassVar[str] = "AIQ_MYSQL_OBJECT_STORE_PASSWORD"
+
+    bucket_name: str = Field(description="The name of the bucket to use for the object store")
+    host: str = Field(
+        default=os.environ.get(HOST_ENV, DEFAULT_HOST),
+        description="The host of the MySQL server"
+        " (uses {HOST_ENV} if unspecified; falls back to {DEFAULT_HOST})",
+    )
+    port: int = Field(
+        default=int(os.environ.get(PORT_ENV, DEFAULT_PORT)),
+        description="The port of the MySQL server"
+        " (uses {PORT_ENV} if unspecified; falls back to {DEFAULT_PORT})",
+    )
+    username: str | None = Field(
+        default=os.environ.get(USERNAME_ENV),
+        description=f"The username used to connect to the MySQL server (uses {USERNAME_ENV} if unspecifed)",
+    )
+    password: str | None = Field(
+        default=os.environ.get(PASSWORD_ENV),
+        description="The password used to connect to the MySQL server (uses {PASSWORD_ENV} if unspecifed)",
+    )
+
+
+@register_object_store(config_type=MySQLObjectStoreClientConfig)
+async def mysql_object_store_client(config: MySQLObjectStoreClientConfig, builder: Builder):
+
+    from .mysql_object_store import MySQLObjectStore
+
+    async with MySQLObjectStore(config) as store:
+        yield store

nat/plugins/mysql/register.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=unused-import
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+
+from . import object_store

nvidia_nat_mysql-1.2.0a20250813.dist-info/METADATA

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: nvidia-nat-mysql
+Version: 1.2.0a20250813
+Summary: Subpackage for MySQL integration in NeMo Agent toolkit
+Keywords: ai,agents,memory,data store
+Classifier: Programming Language :: Python
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: nvidia-nat==v1.2.0a20250813
+Requires-Dist: aiomysql>=0.2.0

nvidia_nat_mysql-1.2.0a20250813.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=EXHuA6z28__lH6qhSyGuIPWLPZ-5BzH9EHVCkpDYyvE,8042
+nat/plugins/mysql/object_store.py,sha256=_9uPdkXS5C6LMgODizyiyhwObOF5AfNF2vZZr-h6Yc0,2600
+nat/plugins/mysql/register.py,sha256=7gqnwyDrYttIlEaa7lo9AASYt-2GrZJE0YT2jpKjepo,845
+nvidia_nat_mysql-1.2.0a20250813.dist-info/METADATA,sha256=4IC6fs-Bs8S_xAGWl3-wf_uF-fLNi6i4tYNa-i7KN20,352
+nvidia_nat_mysql-1.2.0a20250813.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_mysql-1.2.0a20250813.dist-info/entry_points.txt,sha256=ZI7ielmDX-k_OPXGvpRB5tKFVR5kMCHmHGwLWqxRKh0,56
+nvidia_nat_mysql-1.2.0a20250813.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_mysql-1.2.0a20250813.dist-info/RECORD,,

nvidia_nat_mysql-1.2.0a20250813.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nvidia_nat_mysql-1.2.0a20250813.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[nat.components]
+nat_mysql = nat.plugins.mysql.register

nvidia_nat_mysql-1.2.0a20250813.dist-info/top_level.txt

@@ -0,0 +1 @@
+nat