From 4c9550964febd8fe2b89914ac241821ffdc861bd Mon Sep 17 00:00:00 2001
From: Adam Kariv <adam.kariv@gmail.com>
Date: Sat, 8 Jun 2024 17:45:13 +0300
Subject: [PATCH] Simplify S3 code

---
 odds/common/store/s3/s3_store.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/odds/common/store/s3/s3_store.py b/odds/common/store/s3/s3_store.py
index 43dc761..4a6b5d2 100644
--- a/odds/common/store/s3/s3_store.py
+++ b/odds/common/store/s3/s3_store.py
@@ -28,19 +28,18 @@ async def bucket(self):
             region_name=config.credentials.s3_store.region,
             endpoint_url=config.credentials.s3_store.endpoint_url,
         ) as s3:
-            yield s3, await s3.Bucket(config.credentials.s3_store.bucket)
+            yield await s3.Bucket(config.credentials.s3_store.bucket)
 
     async def storeDataset(self, dataset: Dataset, ctx: str) -> None:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             id = dataset.storeId()
             key = self.get_key('dataset', id, 'json')
             rts.set(ctx, f'STORING DATASET {dataset.title} -> {key}')
             obj = await bucket.Object(key)
             await obj.put(Body=json.dumps(dataclasses.asdict(dataset), indent=2, ensure_ascii=False).encode('utf-8'))
-            # await s3.put_object(Bucket=bucket, Key=key, Body=json.dumps(dataclasses.asdict(dataset), indent=2, ensure_ascii=False))
 
     async def storeDB(self, resource: Resource, dataset: Dataset, dbFile, ctx: str) -> None:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             id = '{}/{}'.format(dataset.storeId(), resource.url)
             key = self.get_key('db', id, 'sqlite')
             rts.set(ctx, f'STORING RES-DB {resource.title} -> {key}')
@@ -48,7 +47,7 @@ async def storeDB(self, resource: Resource, dataset: Dataset, dbFile, ctx: str)
             await obj.upload_file(dbFile)
 
     async def storeEmbedding(self, dataset: Dataset, embedding: Embedding, ctx: str) -> None:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             id = dataset.storeId()
             key = self.get_key('embedding', id, 'npy')
             rts.set(ctx, f'STORING EMBEDDING -> {key}')
@@ -59,7 +58,7 @@ async def storeEmbedding(self, dataset: Dataset, embedding: Embedding, ctx: str)
             await obj.upload_fileobj(filename)
         
     async def getDataset(self, datasetId: str) -> Dataset:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             key = self.get_key('dataset', datasetId, 'json')
             try:
                 obj = await bucket.Object(key)
@@ -80,7 +79,7 @@ async def getDataset(self, datasetId: str) -> Dataset:
             return dataset
     
     async def hasDataset(self, datasetId: str) -> bool:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             key = self.get_key('dataset', datasetId, 'json')
             obj = await bucket.Object(key)
             try:
@@ -90,7 +89,7 @@ async def hasDataset(self, datasetId: str) -> bool:
                 return False
     
     async def getDB(self, resource: Resource, dataset: Dataset) -> str:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             id = '{}/{}'.format(dataset.storeId(), resource.url)
             key = self.get_key('db', id, 'sqlite')
             print('GETTING DB', dataset.catalogId, dataset.id, resource.title, key)
@@ -107,7 +106,7 @@ async def getDB(self, resource: Resource, dataset: Dataset) -> str:
         return None
     
     async def getEmbedding(self, dataset: Dataset) -> Embedding:
-        async with self.bucket() as (s3, bucket):
+        async with self.bucket() as bucket:
             id = dataset.storeId()
             key = self.get_key('embedding', id, 'npy')
             try: