Skip to content

Commit a43238f

Browse files
willkill07willkill07
committed
Address coderabbit review
Signed-off-by: Will Killian <[email protected]>
1 parent b99a122 commit a43238f

File tree

7 files changed

+48
-32
lines changed

7 files changed

+48
-32
lines changed

‎docs/source/store-and-retrieve/object-store.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,17 +177,17 @@ object_stores:
177177
```
178178
179179
This enables HTTP endpoints for object store operations:
180-
- **PUT** `/static/{file_path}` - Update an existing object
180+
- **PUT** `/static/{file_path}` - Create or replace an object at the given path (upsert)
181181
```console
182-
$ curl -X PUT -d @data.txt http://localhost:9000/static/folder/data.txt
182+
$ curl -X PUT --data-binary @data.txt http://localhost:9000/static/folder/data.txt
183183
```
184184
- **GET** `/static/{file_path}` - Download an object
185185
```console
186186
$ curl -X GET http://localhost:9000/static/folder/data.txt
187187
```
188188
- **POST** `/static/{file_path}` - Upload a new object
189189
```console
190-
$ curl -X POST -d @data_new.txt http://localhost:9000/static/folder/data.txt
190+
$ curl -X POST --data-binary @data_new.txt http://localhost:9000/static/folder/data.txt
191191
```
192192
- **DELETE** `/static/{file_path}` - Delete an object
193193
```console

‎examples/object_store/user_report/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -163,8 +163,8 @@ The above command will use the S3-compatible object store. Other configuration f
163163
### Using the Object Store Backed File Server (Optional)
164164

165165
- Download an object: `curl -X GET http://<hostname>:<port>/static/{file_path} -o {filename}`
166-
- Upload an object: `curl -X POST http://<hostname>:<port>/static/{file_path} -d @{filename}`
167-
- Upsert an object: `curl -X PUT http://<hostname>:<port>/static/{file_path} -d @{filename}`
166+
- Upload an object: `curl -X POST http://<hostname>:<port>/static/{file_path} --data-binary @{filename}`
167+
- Upsert an object: `curl -X PUT http://<hostname>:<port>/static/{file_path} --data-binary @{filename}`
168168
- Delete an object: `curl -X DELETE http://<hostname>:<port>/static/{file_path}`
169169

170170
If any of the loading scripts were run and the files are in the object store, example commands are:

‎packages/nvidia_nat_mysql/src/nat/plugins/mysql/mysql_object_store.py‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515

1616
import logging
17+
import re
1718

1819
import aiomysql
1920
from aiomysql.pool import Pool
@@ -33,19 +34,24 @@ class MySQLObjectStore(ObjectStore):
3334
"""
3435

3536
def __init__(self, *, bucket_name: str, host: str, port: int, username: str | None, password: str | None):
37+
super().__init__()
38+
39+
if not re.fullmatch(r"[A-Za-z0-9_-]+", bucket_name):
40+
raise ValueError("bucket_name must match [A-Za-z0-9_-]+")
41+
3642
self._bucket_name = bucket_name
3743
self._host = host
3844
self._port = port
3945
self._username = username
4046
self._password = password
4147

42-
super().__init__()
43-
4448
self._conn_pool: Pool | None = None
4549

46-
self._schema = f"`bucket_{self._bucket_name}`"
50+
@property
51+
def _schema(self) -> str:
52+
return f"`bucket_{self._bucket_name}`"
4753

48-
async def __aenter__(self):
54+
async def __aenter__(self) -> "MySQLObjectStore":
4955

5056
if self._conn_pool is not None:
5157
raise RuntimeError("Connection already established")
@@ -98,7 +104,7 @@ async def __aenter__(self):
98104

99105
return self
100106

101-
async def __aexit__(self, exc_type, exc_value, traceback):
107+
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
102108

103109
if not self._conn_pool:
104110
raise RuntimeError("Connection not established")

‎packages/nvidia_nat_redis/src/nat/plugins/redis/redis_object_store.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,8 @@ def __init__(self, *, bucket_name: str, host: str, port: int, db: int):
4242
self._port = port
4343
self._db = db
4444
self._client: redis.Redis | None = None
45-
self._key_prefix = f"nat:object_store:bucket:{self._bucket_name}"
4645

47-
async def __aenter__(self):
46+
async def __aenter__(self) -> "RedisObjectStore":
4847

4948
if self._client is not None:
5049
raise RuntimeError("Connection already established")
@@ -66,7 +65,7 @@ async def __aenter__(self):
6665

6766
return self
6867

69-
async def __aexit__(self, exc_type, exc_value, traceback):
68+
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
7069

7170
if not self._client:
7271
raise RuntimeError("Connection not established")
@@ -75,7 +74,7 @@ async def __aexit__(self, exc_type, exc_value, traceback):
7574
self._client = None
7675

7776
def _make_key(self, key: str) -> str:
78-
return f"{self._key_prefix}:{key}"
77+
return f"nat/object_store/{self._bucket_name}/{key}"
7978

8079
@override
8180
async def put_object(self, key: str, item: ObjectStoreItem):

‎packages/nvidia_nat_redis/tests/test_redis_object_store.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16+
import socket
1617
from contextlib import asynccontextmanager
1718

1819
import pytest
@@ -26,7 +27,18 @@
2627
# docker run --rm -ti --name test-redis -p 6379:6379 redis:7-alpine
2728

2829

30+
def _redis_available(host: str = "localhost", port: int = 6379) -> bool:
31+
with socket.socket() as s:
32+
s.settimeout(0.25)
33+
try:
34+
s.connect((host, port))
35+
return True
36+
except OSError:
37+
return False
38+
39+
2940
@pytest.mark.integration
41+
@pytest.mark.skipif(not _redis_available(), reason="Redis server not available")
3042
class TestRedisObjectStore(ObjectStoreTests):
3143

3244
@asynccontextmanager

‎packages/nvidia_nat_s3/src/nat/plugins/s3/s3_object_store.py‎

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -47,20 +47,16 @@ def __init__(self,
4747
self._client: BaseClient | None = None
4848
self._client_context = None
4949

50-
if not access_key:
51-
raise ValueError("Access key is not set")
50+
self._client_args: dict = {}
51+
if access_key and secret_key:
52+
self._client_args["aws_access_key_id"] = access_key
53+
self._client_args["aws_secret_access_key"] = secret_key
54+
if region:
55+
self._client_args["region_name"] = region
56+
if endpoint_url:
57+
self._client_args["endpoint_url"] = endpoint_url
5258

53-
if not secret_key:
54-
raise ValueError("Secret key is not set")
55-
56-
self._client_args = {
57-
"aws_access_key_id": access_key,
58-
"aws_secret_access_key": secret_key,
59-
"region_name": region,
60-
"endpoint_url": endpoint_url,
61-
}
62-
63-
async def __aenter__(self):
59+
async def __aenter__(self) -> "S3ObjectStore":
6460

6561
if self._client_context is not None:
6662
raise RuntimeError("Connection already established")
@@ -82,7 +78,7 @@ async def __aenter__(self):
8278

8379
return self
8480

85-
async def __aexit__(self, exc_type, exc_value, traceback):
81+
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
8682

8783
if self._client_context is None:
8884
raise RuntimeError("Connection not established")

‎src/nat/cli/commands/object_store/object_store.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import importlib
1818
import logging
1919
import mimetypes
20+
import time
2021
from pathlib import Path
2122

2223
import click
@@ -60,16 +61,15 @@ async def upload_file(object_store: ObjectStore, file_path: Path, key: str):
6061
key: The key to upload the file to.
6162
"""
6263
try:
63-
with open(file_path, "rb") as f:
64-
data = f.read()
64+
data = await asyncio.to_thread(file_path.read_bytes)
6565

6666
item = ObjectStoreItem(data=data,
6767
content_type=mimetypes.guess_type(str(file_path))[0],
6868
metadata={
6969
"original_filename": file_path.name,
7070
"file_size": str(len(data)),
7171
"file_extension": file_path.suffix,
72-
"upload_timestamp": str(int(asyncio.get_event_loop().time()))
72+
"upload_timestamp": str(int(time.time()))
7373
})
7474

7575
# Upload using upsert to allow overwriting
@@ -186,17 +186,19 @@ def register_object_store_commands():
186186
@click.option("--region", type=str, help="S3 region")
187187
@click.pass_context
188188
def s3(ctx: click.Context, **kwargs):
189+
ctx.ensure_object(dict)
189190
ctx.obj["store_config"] = get_object_store_config(store_type="s3", **kwargs)
190191

191192
@click.group(name="mysql", invoke_without_command=False, help="MySQL object store operations.")
192193
@click.argument("bucket_name", type=str, required=True)
193194
@click.option("--host", type=str, help="MySQL host")
194195
@click.option("--port", type=int, help="MySQL port")
195-
@click.option("--db", type=int, help="MySQL db")
196+
@click.option("--db", type=str, help="MySQL database name")
196197
@click.option("--username", type=str, help="MySQL username")
197198
@click.option("--password", type=str, help="MySQL password")
198199
@click.pass_context
199200
def mysql(ctx: click.Context, **kwargs):
201+
ctx.ensure_object(dict)
200202
ctx.obj["store_config"] = get_object_store_config(store_type="mysql", **kwargs)
201203

202204
@click.group(name="redis", invoke_without_command=False, help="Redis object store operations.")
@@ -206,6 +208,7 @@ def mysql(ctx: click.Context, **kwargs):
206208
@click.option("--db", type=int, help="Redis db")
207209
@click.pass_context
208210
def redis(ctx: click.Context, **kwargs):
211+
ctx.ensure_object(dict)
209212
ctx.obj["store_config"] = get_object_store_config(store_type="redis", **kwargs)
210213

211214
commands = {"s3": s3, "mysql": mysql, "redis": redis}

0 commit comments

Comments
 (0)