Enhancements in Filesystem Configuration

dlt/common/configuration/specs/gcp_credentials.py

@@ -1,5 +1,5 @@
 import sys
-from typing import Any, Final, List, Tuple, Union
+from typing import Any, Final, List, Tuple, Union, Dict
 
 from dlt.common import json, pendulum
 from dlt.common.configuration.specs.api_credentials import OAuth2Credentials
@@ -48,6 +48,20 @@ def _from_info_dict(self, info: StrAny) -> None:
     def __str__(self) -> str:
         return f"{self.project_id}"
 
+    def to_gcs_credentials(self) -> Dict[str, Any]:
+        """
+        Dict of keyword arguments can be passed to gcsfs.
+        Delegates default GCS credential handling to gcsfs.
+        """
+        return {
+            "project": self.project_id,
+            "token": (
+                None
+                if isinstance(self, CredentialsWithDefault) and self.has_default_credentials()
+                else dict(self)
+            ),
+        }
+
 
 @configspec
 class GcpServiceAccountCredentialsWithoutDefaults(GcpCredentials):

dlt/common/storages/configuration.py

@@ -1,9 +1,10 @@
 import os
-from urllib.parse import urlparse
 from typing import TYPE_CHECKING, Any, Literal, Optional, Type, get_args, ClassVar, Dict, Union
+from urllib.parse import urlparse
 
-from dlt.common.configuration.specs import BaseConfiguration, configspec, CredentialsConfiguration
 from dlt.common.configuration import configspec, resolve_type
+from dlt.common.configuration.exceptions import ConfigurationValueError
+from dlt.common.configuration.specs import CredentialsConfiguration
 from dlt.common.configuration.specs import (
     GcpServiceAccountCredentials,
     AwsCredentials,
@@ -12,8 +13,9 @@
     AzureCredentialsWithoutDefaults,
     BaseConfiguration,
 )
+from dlt.common.typing import DictStrAny
 from dlt.common.utils import digest128
-from dlt.common.configuration.exceptions import ConfigurationValueError
+
 
 TSchemaFileFormat = Literal["json", "yaml"]
 SchemaFileExtensions = get_args(TSchemaFileFormat)
@@ -92,9 +94,13 @@ class FilesystemConfiguration(BaseConfiguration):
     }
 
     bucket_url: str = None
-    # should be an union of all possible credentials as found in PROTOCOL_CREDENTIALS
+
+    # should be a union of all possible credentials as found in PROTOCOL_CREDENTIALS
     credentials: FileSystemCredentials
 
+    kwargs: Optional[DictStrAny] = None
+    client_kwargs: Optional[DictStrAny] = None
+
     @property
     def protocol(self) -> str:
         """`bucket_url` protocol"""
@@ -112,7 +118,7 @@ def on_resolved(self) -> None:
                 "File path or netloc missing. Field bucket_url of FilesystemClientConfiguration"
                 " must contain valid url with a path or host:password component."
             )
-        # this is just a path in local file system
+        # this is just a path in a local file system
         if url.path == self.bucket_url:
             url = url._replace(scheme="file")
             self.bucket_url = url.geturl()
@@ -124,9 +130,7 @@ def resolve_credentials_type(self) -> Type[CredentialsConfiguration]:
 
     def fingerprint(self) -> str:
         """Returns a fingerprint of bucket_url"""
-        if self.bucket_url:
-            return digest128(self.bucket_url)
-        return ""
+        return digest128(self.bucket_url) if self.bucket_url else ""
 
     def __str__(self) -> str:
         """Return displayable destination location"""
@@ -141,4 +145,15 @@ def __str__(self) -> str:
 
     if TYPE_CHECKING:
 
-        def __init__(self, bucket_url: str, credentials: FileSystemCredentials = None) -> None: ...
+        def __init__(
+            self,
+            bucket_url: str,
+            credentials: FileSystemCredentials = None,
+            kwargs: Optional[DictStrAny] = None,
+            client_kwargs: Optional[DictStrAny] = None,
+        ) -> None:
+            self.bucket_url = bucket_url
+            self.credentials = credentials
+            self.kwargs = kwargs
+            self.client_kwargs = client_kwargs
+            ...

dlt/common/storages/fsspec_filesystem.py

@@ -1,27 +1,25 @@
 import io
 import mimetypes
-import posixpath
 import pathlib
-from urllib.parse import urlparse
+import posixpath
 from io import BytesIO
-from typing import cast, Tuple, TypedDict, Optional, Union, Iterator, Any, IO
+from typing import cast, Tuple, TypedDict, Optional, Union, Iterator, Any, IO, Dict, Callable
+from urllib.parse import urlparse
 
-from fsspec.core import url_to_fs
 from fsspec import AbstractFileSystem
+from fsspec.core import url_to_fs
 
+from dlt import version
 from dlt.common import pendulum
-from dlt.common.exceptions import MissingDependencyException
-from dlt.common.time import ensure_pendulum_datetime
-from dlt.common.typing import DictStrAny
 from dlt.common.configuration.specs import (
-    CredentialsWithDefault,
     GcpCredentials,
     AwsCredentials,
     AzureCredentials,
 )
+from dlt.common.exceptions import MissingDependencyException
 from dlt.common.storages.configuration import FileSystemCredentials, FilesystemConfiguration
-
-from dlt import version
+from dlt.common.time import ensure_pendulum_datetime
+from dlt.common.typing import DictStrAny
 
 
 class FileItem(TypedDict, total=False):
@@ -50,56 +48,75 @@ class FileItem(TypedDict, total=False):
 MTIME_DISPATCH["s3a"] = MTIME_DISPATCH["s3"]
 MTIME_DISPATCH["abfs"] = MTIME_DISPATCH["az"]
 
+# Map of protocol to a filesystem type
+CREDENTIALS_DISPATCH: Dict[str, Callable[[FilesystemConfiguration], DictStrAny]] = {
+    "s3": lambda config: cast(AwsCredentials, config.credentials).to_s3fs_credentials(),
+    "adl": lambda config: cast(AzureCredentials, config.credentials).to_adlfs_credentials(),
+    "az": lambda config: cast(AzureCredentials, config.credentials).to_adlfs_credentials(),
+    "gcs": lambda config: cast(GcpCredentials, config.credentials).to_gcs_credentials(),
+    "gs": lambda config: cast(GcpCredentials, config.credentials).to_gcs_credentials(),
+    "abfs": lambda config: cast(AzureCredentials, config.credentials).to_adlfs_credentials(),
+    "azure": lambda config: cast(AzureCredentials, config.credentials).to_adlfs_credentials(),
+}
+
 
 def fsspec_filesystem(
-    protocol: str, credentials: FileSystemCredentials = None
+    protocol: str,
+    credentials: FileSystemCredentials = None,
+    kwargs: Optional[DictStrAny] = None,
+    client_kwargs: Optional[DictStrAny] = None,
 ) -> Tuple[AbstractFileSystem, str]:
     """Instantiates an authenticated fsspec `FileSystem` for a given `protocol` and credentials.
 
-    Please supply credentials instance corresponding to the protocol. The `protocol` is just the code name of the filesystem ie:
+    Please supply credentials instance corresponding to the protocol.
+    The `protocol` is just the code name of the filesystem i.e.:
     * s3
     * az, abfs
     * gcs, gs
 
     also see filesystem_from_config
     """
-    return fsspec_from_config(FilesystemConfiguration(protocol, credentials))
+    return fsspec_from_config(
+        FilesystemConfiguration(protocol, credentials, kwargs=kwargs, client_kwargs=client_kwargs)
+    )
 
 
 def fsspec_from_config(config: FilesystemConfiguration) -> Tuple[AbstractFileSystem, str]:
     """Instantiates an authenticated fsspec `FileSystem` from `config` argument.
 
-    Authenticates following filesystems:
+    Authenticates the following filesystems:
     * s3
     * az, abfs
     * gcs, gs
 
+    Additional fsspec filesystem arguments and client arguments are gathered from the
+    FilesystemConfiguration object and passed to the `url_to_fs` factory.
+
     All other filesystems are not authenticated
 
     Returns: (fsspec filesystem, normalized url)
 
     """
     proto = config.protocol
-    fs_kwargs: DictStrAny = {}
-    if proto == "s3":
-        fs_kwargs.update(cast(AwsCredentials, config.credentials).to_s3fs_credentials())
-    elif proto in ["az", "abfs", "adl", "azure"]:
-        fs_kwargs.update(cast(AzureCredentials, config.credentials).to_adlfs_credentials())
-    elif proto in ["gcs", "gs"]:
-        assert isinstance(config.credentials, GcpCredentials)
-        # Default credentials are handled by gcsfs
-        if (
-            isinstance(config.credentials, CredentialsWithDefault)
-            and config.credentials.has_default_credentials()
-        ):
-            fs_kwargs["token"] = None
-        else:
-            fs_kwargs["token"] = dict(config.credentials)
-        fs_kwargs["project"] = config.credentials.project_id
+    fs_kwargs: DictStrAny = {"use_listings_cache": False}
+    credentials = CREDENTIALS_DISPATCH.get(proto, lambda _: {})(config)
+
+    if config.kwargs is not None:
+        fs_kwargs.update(config.kwargs)
+    if config.client_kwargs is not None:
+        fs_kwargs["client_kwargs"] = config.client_kwargs
+
+    if "client_kwargs" in fs_kwargs and "client_kwargs" in credentials:
+        fs_kwargs["client_kwargs"].update(credentials.pop("client_kwargs"))
+
+    fs_kwargs.update(credentials)
+
     try:
-        return url_to_fs(config.bucket_url, use_listings_cache=False, **fs_kwargs)  # type: ignore[no-any-return]
+        return url_to_fs(config.bucket_url, **fs_kwargs)  # type: ignore
     except ModuleNotFoundError as e:
-        raise MissingDependencyException("filesystem", [f"{version.DLT_PKG_NAME}[{proto}]"]) from e
+        raise MissingDependencyException(
+            "filesystem", [f"{version.DLT_PKG_NAME}[{config.protocol}]"]
+        ) from e
 
 
 class FileItemDict(DictStrAny):
@@ -122,7 +139,7 @@ def __init__(
 
     @property
     def fsspec(self) -> AbstractFileSystem:
-        """The filesystem client based on the given credentials.
+        """The filesystem client is based on the given credentials.
 
         Returns:
             AbstractFileSystem: The fsspec client.
@@ -150,16 +167,15 @@ def open(self, mode: str = "rb", **kwargs: Any) -> IO[Any]:  # noqa: A003
         if "file_content" in self:
             bytes_io = BytesIO(self["file_content"])
 
-            if "t" in mode:
-                text_kwargs = {
-                    k: kwargs.pop(k) for k in ["encoding", "errors", "newline"] if k in kwargs
-                }
-                return io.TextIOWrapper(
-                    bytes_io,
-                    **text_kwargs,
-                )
-            else:
+            if "t" not in mode:
                 return bytes_io
+            text_kwargs = {
+                k: kwargs.pop(k) for k in ["encoding", "errors", "newline"] if k in kwargs
+            }
+            return io.TextIOWrapper(
+                bytes_io,
+                **text_kwargs,
+            )
         else:
             opened_file = self.fsspec.open(self["file_url"], mode=mode, **kwargs)
         return opened_file
@@ -171,19 +187,17 @@ def read_bytes(self) -> bytes:
             bytes: The file content.
         """
         content: bytes
-        # same as open, if the user has already extracted the content, we use it.
-        if "file_content" in self and self["file_content"] is not None:
-            content = self["file_content"]
-        else:
-            content = self.fsspec.read_bytes(self["file_url"])
-        return content
+        return (  # type: ignore
+            self["file_content"]
+            if "file_content" in self and self["file_content"] is not None
+            else self.fsspec.read_bytes(self["file_url"])
+        )
 
 
 def guess_mime_type(file_name: str) -> str:
-    mime_type = mimetypes.guess_type(posixpath.basename(file_name), strict=False)[0]
-    if not mime_type:
-        mime_type = "application/" + (posixpath.splitext(file_name)[1][1:] or "octet-stream")
-    return mime_type
+    return mimetypes.guess_type(posixpath.basename(file_name), strict=False)[
+        0
+    ] or "application/" + (posixpath.splitext(file_name)[1][1:] or "octet-stream")
 
 
 def glob_files(
@@ -202,7 +216,7 @@ def glob_files(
     import os
 
     bucket_url_parsed = urlparse(bucket_url)
-    # if this is file path without scheme
+    # if this is a file path without a scheme
     if not bucket_url_parsed.scheme or (os.path.isabs(bucket_url) and "\\" in bucket_url):
         # this is a file so create a proper file url
         bucket_url = pathlib.Path(bucket_url).absolute().as_uri()
@@ -224,9 +238,9 @@ def glob_files(
             continue
         # make that absolute path on a file://
         if bucket_url_parsed.scheme == "file" and not file.startswith("/"):
-            file = "/" + file
+            file = f"/{file}"
         file_name = posixpath.relpath(file, bucket_path)
-        file_url = bucket_url_parsed.scheme + "://" + file
+        file_url = f"{bucket_url_parsed.scheme}://{file}"
         yield FileItem(
             file_name=file_name,
             file_url=file_url,

docs/website/docs/conftest.py

@@ -52,6 +52,6 @@ def _initial_providers():
 
 def pytest_configure(config):
     # push sentry to ci
-    os.environ[
-        "RUNTIME__SENTRY_DSN"
-    ] = "https://[email protected]/4504819859914752"
+    os.environ["RUNTIME__SENTRY_DSN"] = (
+        "https://[email protected]/4504819859914752"
+    )