Upload folder implementation.

src/tiledb/cloud/files/utils.py

@@ -1,22 +1,28 @@
 import base64
+import logging
 import os
 import re
 import urllib.parse
 import warnings
 from fnmatch import fnmatch
-from typing import Mapping, Optional, Tuple, Union
+from typing import Dict, List, Mapping, Optional, Tuple, Union
+
+import attrs
 
 import tiledb
+import tiledb.cloud
 import tiledb.cloud.tiledb_cloud_error as tce
 from tiledb.cloud import array
 from tiledb.cloud import client
 from tiledb.cloud import config
+from tiledb.cloud import groups
 from tiledb.cloud import rest_api
 from tiledb.cloud import tiledb_cloud_error
 from tiledb.cloud._common import utils
 from tiledb.cloud.rest_api import ApiException as GenApiException
 from tiledb.cloud.rest_api import configuration
 from tiledb.cloud.rest_api import models
+from tiledb.cloud.utilities import get_logger_wrapper
 
 
 def sanitize_filename(fname: str) -> str:
@@ -262,3 +268,170 @@ def _auth_headers(cfg: configuration.Configuration) -> Mapping[str, object]:
  return {"authorization": basic}
  return {}
  # No authentication has been provided. Do nothing.
+
+
+@attrs.define
+class UploadFoldersResults:
+ directory: str
+ report: str
+ errors: dict
+ sub_folders: list["UploadFoldersResults"]
+
+
+def upload_folder(
+ input_uri: str,
+ output_uri: str,
+ *,
+ parent_group_uri: Optional[str] = None,
+ access_credentials_name: Optional[str] = None,
+ config: Optional[dict] = None,
+ flatten: bool = False,
+ serializable: bool = True,
+ logger: Optional[logging.Logger] = None,
+ verbose: bool = False,
+) -> Union[dict, UploadFoldersResults]:
+ """
+ Uploads a folder to TileDB Cloud.
+ By default respects the initial folder structure in the destination.
+
+ :param input_uri: The URI or path of the input folder. May be an ordinary path
+ or any URI accessible via TileDB VFS.
+ :param output_uri: The TileDB URI to write the folder into.
+ :param parent_group_uri: A TileDB Group URI to add folder under,
+ defaults to None
+ :param access_credentials_name: If present, the name of the credentials
+ to use when writing the uploaded file to backend storage instead of
+ the defaults.
+ :param config: Config dictionary, defaults to None
+ :param flatten: Flag. If set to True, the upload will flatten the folder
+ structure instead of recreating it. (Not Implemented yet)
+ :param serializable: Flag. If set to True the function returns a dictionary
+ report. Differently it returns an UploadFoldersResults attrs class.
+ Defaults to True.
+ :param logger: A logging.Logger instance, defaults to None.
+ :param verbose: Verbose logging, defaults to None
+ :return: A dictionary containing a report message
+ and an upload errors dictionary (if any)
+ """
+ logger = logger or get_logger_wrapper(verbose)
+ logger.info("=====")
+
+ if flatten:
+ raise NotImplementedError(
+ "The option to flatten a folder structure is not yet implemented"
+ )
+
+ # Prepare and sanitize arguments
+ output_uri = output_uri.strip("/")
+
+ input_uri = input_uri if input_uri.endswith(os.sep) else input_uri + os.sep
+ base_dir = os.path.dirname(input_uri)
+ base_dir = os.path.basename(base_dir)
+
+ namespace, name = utils.split_uri(output_uri)
+ _, sp, acn = groups._default_ns_path_cred(namespace=namespace)
+
+ # If `name` is a URL, assume it points to a cloud storage
+ storage_path = name if "://" in name else sp
+ storage_path = f"{storage_path.strip('/')}/{base_dir}"
+ tb_storage_uri = f"tiledb://{namespace}/{storage_path}"
+ logger.debug("Output storage path: %s", storage_path)
+ logger.debug("TileDB Storage URI: %r", tb_storage_uri)
+
+ access_credentials_name = access_credentials_name or acn
+ logger.debug("Storage path: %r", storage_path)
+ logger.debug("ACN: %s", access_credentials_name)
+
+ # Group check and/or creation
+ if "://" not in name:
+ group_uri = output_uri
+ else:
+ group_uri = f"tiledb://{namespace}/{base_dir}"
+ logger.debug("Group URI: %r", group_uri)
+
+ group_created = False
+ if not tiledb.object_type(group_uri, ctx=tiledb.cloud.Ctx()) == "group":
+ group_namespace, group_name = utils.split_uri(group_uri)
+ groups.create(
+ name=group_name,
+ namespace=group_namespace,
+ storage_uri=storage_path,
+ parent_uri=parent_group_uri,
+ credentials_name=access_credentials_name,
+ )
+ group_created = True
+ logger.debug("Group URI: %r created", group_uri)
+
+ # Upload Stats
+ logger.info("-----")
+ if parent_group_uri:
+ logger.info("Sub-Folder %r Upload Stats", base_dir)
+ else:
+ logger.info("Folder %r Upload Stats", base_dir)
+ logger.info("-----")
+ logger.info("- Input URI: %r", input_uri)
+ logger.info("- Output URI: %r", output_uri)
+ logger.info("-- Storage Path: %r", storage_path)
+ logger.info("- Group URI: %r", group_uri)
+ logger.info("-- Group Created: %s", group_created)
+
+ # Report results
+ uploaded = 0
+ dir_count = 0
+ upload_errors: Dict[str, str] = {}
+ results = UploadFoldersResults(
+ directory=base_dir, report="", errors={}, sub_folders=[]
+ )
+
+ vfs = tiledb.VFS(config=config)
+ # List local folder
+ input_ls: List[str] = vfs.ls(input_uri)
+ with tiledb.Group(
+ group_uri, mode="w", config=config, ctx=tiledb.cloud.Ctx()
+ ) as grp:
+ for fname in input_ls:
+ if vfs.is_dir(fname):
+ dir_count += 1
+
+ results.sub_folders.append(
+ upload_folder(
+ input_uri=fname,
+ output_uri=tb_storage_uri,
+ parent_group_uri=group_uri,
+ access_credentials_name=access_credentials_name,
+ config=config,
+ # Serialization concerns the final results.
+ serializable=False,
+ logger=logger,
+ verbose=verbose,
+ )
+ )
+ else:
+ logger.info("Uploading %r to Group %r" % (fname, group_uri))
+ filename = fname.split(base_dir + os.sep)[1]
+ tb_output_uri = f"{tb_storage_uri}/{filename}"
+ try:
+ uploaded_uri = upload_file(
+ input_uri=fname,
+ output_uri=tb_output_uri,
+ filename=filename,
+ access_credentials_name=access_credentials_name,
+ )
+ grp.add(uploaded_uri)
+ uploaded += 1
+ except Exception as exc:
+ logger.exception(
+ "File %r while uploading to %r raised an exception"
+ % (filename, tb_output_uri)
+ )
+ upload_errors[fname] = str(exc)
+ continue
+
+ results.report = f"Uploaded {uploaded}/{len(input_ls) - dir_count} files"
+ results.errors = upload_errors
+ logger.info(results.report)
+ logger.info("=====")
+
+ if serializable:
+ return attrs.asdict(results)
+ return results

tests/test_file.py

@@ -2,6 +2,7 @@
 import os
 import pathlib
 import tempfile
+import time
 import unittest
 from typing import List
 
@@ -191,8 +192,91 @@ def test_round_trip(self):
  array.delete_array(uri)
 
 
-# FIXME: Will be fixed with #595 implementation
-# Disable until then.
+class UploadFolderTest(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ cls.ASSERT_EXIST_TRIES = 5
+ cls.folder_basename = "upload_folder_test"
+ cls.test_folder = os.path.join(CURRENT_DIR, "data", cls.folder_basename)
+ cls.namespace, cls.storage_path, cls.acn = groups._default_ns_path_cred()
+ return super().setUpClass()
+
+ def setUp(self):
+ self.group_uri = (
+ f"tiledb://{self.namespace}/{testonly.random_name('upload_folder')}"
+ )
+ return super().setUp()
+
+ def tearDown(self):
+ tiledb.cloud.groups.delete(self.group_uri, recursive=True)
+ return super().tearDown()
+
+ def group_info(self, uri: str):
+ for _ in range(self.ASSERT_EXIST_TRIES):
+ try:
+ return groups.info(uri)
+ except Exception:
+ pass
+ time.sleep(1)
+ self.fail(f"Group '{uri}' does not exist")
+
+ def test_upload_folder(self):
+ report = file_utils.upload_folder(
+ input_uri=self.test_folder,
+ output_uri=self.group_uri,
+ access_credentials_name=self.acn,
+ serializable=True,
+ )
+
+ info = self.group_info(self.group_uri)
+ self.assertEqual(info.size, 4) # 2 sub-folders, 2 files
+
+ base_storage_uri = f"{self.storage_path}/{self.folder_basename}"
+ info = self.group_info(f"tiledb://{self.namespace}/nested_folder_1")
+ self.assertEqual(info.size, 2) # 2 files
+ self.assertEqual(info.uri, f"{base_storage_uri}/nested_folder_1")
+
+ info = self.group_info(f"tiledb://{self.namespace}/nested_folder_2")
+ self.assertEqual(info.size, 3) # 1 sub-folder, 2 files
+ self.assertEqual(info.uri, f"{base_storage_uri}/nested_folder_2")
+
+ info = self.group_info(f"tiledb://{self.namespace}/double_nested_folder")
+ self.assertEqual(info.size, 2) # 2 files
+ self.assertEqual(
+ info.uri, f"{base_storage_uri}/nested_folder_2/double_nested_folder"
+ )
+
+ self.assertDictEqual(
+ report,
+ {
+ "directory": self.folder_basename,
+ "report": "Uploaded 2/2 files",
+ "errors": {},
+ "sub_folders": [
+ {
+ "directory": "nested_folder_1",
+ "report": "Uploaded 2/2 files",
+ "errors": {},
+ "sub_folders": [],
+ },
+ {
+ "directory": "nested_folder_2",
+ "report": "Uploaded 2/2 files",
+ "errors": {},
+ "sub_folders": [
+ {
+ "directory": "double_nested_folder",
+ "report": "Uploaded 2/2 files",
+ "errors": {},
+ "sub_folders": [],
+ }
+ ],
+ },
+ ],
+ },
+ )
+
+
 @unittest.skip("Skip until fixed VFS access")
 class TestFileIngestion(unittest.TestCase):
  @classmethod