Skip to content

dqm_ml_job.utils

DQM-ML job utilities.

__all__ = ['get_s3_filesystem'] module-attribute

get_s3_filesystem(storage_config: StorageConfig) -> pa.fs.S3FileSystem | None

Create and return an S3 filesystem instance from StorageConfig.

Credentials fall back to environment variables

S3_ACCESS_KEY, S3_SECRET_KEY, S3_ENDPOINT, S3_REGION

Parameters:

Name Type Description Default
storage_config StorageConfig

Validated StorageConfig instance.

 required

Returns:

Type Description
S3FileSystem | None

Configured S3 filesystem or None if credentials not available.
Source code in packages/dqm-ml-job/src/dqm_ml_job/utils/s3.py 
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def get_s3_filesystem(
    storage_config: StorageConfig,
) -> pa.fs.S3FileSystem | None:
    """Create and return an S3 filesystem instance from StorageConfig.

    Credentials fall back to environment variables:
      S3_ACCESS_KEY, S3_SECRET_KEY, S3_ENDPOINT, S3_REGION

    Args:
        storage_config: Validated StorageConfig instance.

    Returns:
        Configured S3 filesystem or None if credentials not available.
    """
    os.environ["AWS_RESPONSE_CHECKSUM_VALIDATION"] = storage_config.checksum_validation

    access_key = storage_config.access_key or os.getenv("S3_ACCESS_KEY")
    secret_key = storage_config.secret_key or os.getenv("S3_SECRET_KEY")

    if not access_key or not secret_key:
        return None

    kwargs: dict[str, Any] = {
        "access_key": access_key,
        "secret_key": secret_key,
        "endpoint_override": storage_config.endpoint or os.getenv("S3_ENDPOINT", ""),
    }

    _apply_simple_kwargs(storage_config, kwargs)

    if storage_config.role_arn:
        _apply_role_kwargs(storage_config, kwargs)

    if storage_config.retry:
        kwargs["retry_strategy"] = _build_retry_strategy(storage_config.retry)

    if os.getenv("ENVIRONMENT") == "mock":
        kwargs.setdefault("background_writes", False)
        kwargs.setdefault("retry_strategy", pa.fs.AwsStandardS3RetryStrategy(max_attempts=10))

    return pa.fs.S3FileSystem(**kwargs)