feat: assuming roles

Author: Vixtir

config_examples/s3.yaml

@@ -11,6 +11,9 @@ plugins:
     aws_access_key_id: <aws_access_key_id> # Optional.
     aws_session_token: <aws_session_token> # Optional.
     aws_region: <aws_region> # Optional.
+    filename_filter: # Optional. Default filter allows each file to be ingested to platform.
+      include: [ '.*.parquet' ]
+      exclude: [ 'dev_.*' ]
     datasets:
       # Recursive fetch for all objects in the bucket.
       - bucket: my_bucket
@@ -41,10 +44,9 @@ token: ""
 plugins:
   - type: s3
     name: s3_minio_adapter
-    endpoint_url: <some_endpoint>
-    aws_secret_access_key: <aws_secret_access_key>
-    aws_access_key_id: <aws_access_key_id>
-    aws_region: <aws_region>
+    endpoint_url: http://localhost:9000
+    aws_secret_access_key: minioadmin
+    aws_access_key_id: minioadmin
     datasets:
       - bucket: my_bucket
         prefix: partitioned_data

odd_collector_aws/adapters/s3/clients/__init__.py

@@ -1,55 +1,23 @@
 from typing import Union
 
-import pyarrow.dataset as ds
-from funcy import iffy, lmap
-from pyarrow._fs import FileInfo, FileSelector
-from pyarrow.fs import S3FileSystem
-
-from odd_collector_aws.domain.plugin import AwsPlugin
+from odd_collector_aws.domain.plugin import S3Plugin
 
 from ...domain.dataset_config import DatasetConfig
+from ...filesystem.pyarrow_fs import FileSystem as PyarrowFs
+from ...utils.remove_s3_protocol import remove_protocol
 from .domain.models import Bucket, File, Folder
 from .logger import logger
 from .utils import file_format
 
 
-class FileSystem:
+class FileSystem(PyarrowFs):
     """
     FileSystem hides pyarrow.fs implementation details.
     """
 
-    def __init__(self, config: AwsPlugin):
-        params = {}
-
-        if config.aws_access_key_id:
-            params["access_key"] = config.aws_access_key_id
-        if config.aws_secret_access_key:
-            params["secret_key"] = config.aws_secret_access_key
-        if config.aws_session_token:
-            params["session_token"] = config.aws_session_token
-        if config.aws_region:
-            params["region"] = config.aws_region
-        if config.endpoint_url:
-            params["endpoint_override"] = config.endpoint_url
-
-        self.fs = S3FileSystem(**params)
-
-    def get_file_info(self, path: str) -> list[FileInfo]:
-        """
-        Get file info from path.
-        @param path: s3 path to file or folder
-        @return: FileInfo
-        """
-        return self.fs.get_file_info(FileSelector(base_dir=path))
-
-    def get_dataset(self, file_path: str, format: str) -> ds.Dataset:
-        """
-        Get dataset from file path.
-        @param file_path:
-        @param format: Should be one of available formats: https://arrow.apache.org/docs/python/api/dataset.html#file-format
-        @return: Dataset
-        """
-        return ds.dataset(source=file_path, filesystem=self.fs, format=format)
+    def __init__(self, config: S3Plugin):
+        self.fs = PyarrowFs(config)
+        self.filename_filter = config.filename_filter
 
     def get_folder_as_file(self, dataset_config: DatasetConfig) -> File:
         """
@@ -59,15 +27,13 @@ def get_folder_as_file(self, dataset_config: DatasetConfig) -> File:
         """
         logger.debug(f"Getting folder dataset for {dataset_config=}")
 
-        dataset = ds.dataset(
-            source=dataset_config.full_path,
+        dataset = self.get_dataset(
+            path=dataset_config.full_path,
             format=dataset_config.folder_as_dataset.file_format,
-            partitioning=ds.partitioning(
-                flavor=dataset_config.folder_as_dataset.flavor,
-                field_names=dataset_config.folder_as_dataset.field_names,
-            ),
-            filesystem=self.fs,
+            partitioning_flavor=dataset_config.folder_as_dataset.flavor,
+            field_names=dataset_config.folder_as_dataset.field_names,
         )
+
         return File(
             path=dataset_config.full_path,
             base_name=dataset_config.full_path,
@@ -103,14 +69,18 @@ def list_objects(self, path: str) -> list[Union[File, Folder]]:
         @return: list of either File or Folder
         """
         logger.debug(f"Getting objects for {path=}")
-        return lmap(
-            iffy(
-                lambda x: x.is_file,
-                lambda x: self.get_file(x.path, x.base_name),
-                lambda x: self.get_folder(x.path),
-            ),
-            self.get_file_info(path),
-        )
+        objects = []
+
+        for obj in self.fs.get_file_info(path):
+            if obj.is_file:
+                if not self.filename_filter.is_allowed(obj.base_name):
+                    continue
+
+                objects.append(self.get_file(obj.path, obj.base_name))
+            else:
+                objects.append(self.get_folder(obj.path))
+
+        return objects
 
     def get_file(self, path: str, file_name: str = None) -> File:
         """
@@ -125,7 +95,7 @@ def get_file(self, path: str, file_name: str = None) -> File:
 
         try:
             file_fmt = file_format(file_name)
-            dataset = self.get_dataset(path, file_fmt)
+            dataset = self.fs.get_dataset(path, file_fmt)
 
             return File.dataset(
                 path=path,
@@ -146,17 +116,9 @@ def get_folder(self, path: str, recursive: bool = True) -> Folder:
         """
         Get Folder with objects recursively.
         @param path: s3 path to
+        @param recursive: Flag to recursively search nested objects
         @return: Folder class with objects and path
         """
         path = remove_protocol(path)
         objects = self.list_objects(path) if recursive else []
         return Folder(path, objects)
-
-
-def remove_protocol(path: str) -> str:
-    if path.startswith("s3://"):
-        return path.removeprefix("s3://")
-    elif path.startswith(("s3a://", "s3n://")):
-        return path[6:]
-    else:
-        return path

odd_collector_aws/adapters/s3/clients/s3_client.py

@@ -1,4 +1,3 @@
-import datetime
 import traceback as tb
 from dataclasses import asdict, dataclass
 from typing import Any, Iterable, Optional
@@ -9,19 +8,15 @@
 from odd_collector_aws.domain.plugin import DeltaTableConfig, S3DeltaPlugin
 from odd_collector_aws.filesystem.pyarrow_fs import FileSystem
 
+from ...utils.dates import add_utc_timezone, from_ms
+from ...utils.remove_s3_protocol import remove_protocol
 from .logger import logger
 from .models.table import DTable
 
 
-def from_ms(ms) -> datetime.datetime:
-    return datetime.datetime.fromtimestamp(ms / 1000, tz=datetime.timezone.utc)
-
-
-def add_utc_timezone(dt: datetime.datetime) -> datetime.datetime:
-    return dt.replace(tzinfo=datetime.timezone.utc)
-
-
-def handle_values(obj: dict, handler: tuple[str, callable]) -> dict[str, Optional[any]]:
+def handle_values(
+    obj: dict, handler: tuple[str, callable]
+) -> tuple[str, Optional[any]]:
     key, callback = handler
     return key, silent(callback)(obj.get(key))
 
@@ -34,8 +29,10 @@ class StorageOptions:
     aws_secret_access_key: str = None
     aws_region: str = None
     aws_session_token: str = None
-    aws_storage_allow_http: bool = None
+    aws_storage_allow_http: str = None
     endpoint_url: str = None
+    aws_profile: str = None
+    aws_role_session_name: str = None
 
     @classmethod
     def from_config(cls, config: S3DeltaPlugin) -> "StorageOptions":
@@ -44,8 +41,10 @@ def from_config(cls, config: S3DeltaPlugin) -> "StorageOptions":
             aws_secret_access_key=config.aws_secret_access_key,
             aws_region=config.aws_region or cls.DEFAULT_REGION,
             aws_session_token=config.aws_session_token,
-            aws_storage_allow_http="true" if config.aws_storage_allow_http else None,
             endpoint_url=config.endpoint_url,
+            aws_storage_allow_http="true" if config.aws_storage_allow_http else None,
+            aws_profile=config.profile_name,
+            aws_role_session_name=config.aws_role_session_name,
         )
 
     def to_dict(self) -> dict[str, str]:
@@ -75,10 +74,8 @@ def handle_folder(self, config: DeltaTableConfig) -> Iterable[DTable]:
 
         objects = self.fs.get_file_info(remove_protocol(config.path))
 
-        allowed = filter(
-            lambda obj: not obj.is_file and config.allow(obj.base_name),
-            objects,
-        )
+        folders = filter(lambda obj: not obj.is_file, objects)
+        allowed = filter(lambda folder: folder.base_name, folders)
 
         for obj in allowed:
             new_config = config.append_prefix(obj.base_name)
@@ -142,12 +139,3 @@ def get_metadata(table: DeltaTable) -> dict[str, Any]:
         logger.error(f"Failed to get metadata for {table.table_uri}. {e}")
 
     return metadata
-
-
-def remove_protocol(path: str) -> str:
-    if path.startswith("s3://"):
-        return path.removeprefix("s3://")
-    elif path.startswith(("s3a://", "s3n://")):
-        return path[6:]
-    else:
-        return path