PrefectHQ / prefect-aws

Prefect integrations with AWS.
https://PrefectHQ.github.io/prefect-aws/
Apache License 2.0
84 stars 40 forks source link

Fix `S3Bucket.load()` for nested MinIO Credentials block #359

Closed kevingrismore closed 9 months ago

kevingrismore commented 9 months ago

Closes #210

Because every field in AwsCredentials has a default or default factory, Pydantic considers non-matching credentials dictionaries valid and selects AwsCredentials when calling S3Bucket.parse_obj().

Example

from pydantic import BaseModel, Field
from typing import Union
from prefect_aws import AwsCredentials, MinIOCredentials

class S3Bucket(BaseModel):
    bucket_name: str = Field(default=...)
    credentials: Union[AwsCredentials, MinIOCredentials] = Field(default_factory=AwsCredentials)

if __name__=="__main__":
    minio_bucket_parsed = S3Bucket.parse_obj({
        "bucket_name": "my_bucket",
        "credentials": {
            "minio_root_user": "Kevin",
            "minio_root_password": "password",
        }
    })

    s3_bucket_parsed = S3Bucket.parse_obj({
        "bucket_name": "my_bucket",
        "credentials": {
            "aws_access_key_id": "1234",
            "aws_access_key_secret": "5678",
        }
    })

    s3_bucket_no_creds_parsed = S3Bucket.parse_obj({
        "bucket_name": "my_bucket"
    })

    print(type(minio_bucket_parsed.credentials))
    print(type(s3_bucket_parsed.credentials))
    print(type(s3_bucket_no_creds_parsed.credentials))
<class 'prefect_aws.credentials.AwsCredentials'>
<class 'prefect_aws.credentials.AwsCredentials'>
<class 'prefect_aws.credentials.AwsCredentials'>

Reordered:

from pydantic import BaseModel, Field
from typing import Union
from prefect_aws import AwsCredentials, MinIOCredentials

class S3Bucket(BaseModel):
    bucket_name: str = Field(default=...)
    credentials: Union[MinIOCredentials, AwsCredentials] = Field(default_factory=AwsCredentials)

if __name__=="__main__":
    minio_bucket_parsed = S3Bucket.parse_obj({
        "bucket_name": "my_bucket",
        "credentials": {
            "minio_root_user": "Kevin",
            "minio_root_password": "password",
        }
    })

    s3_bucket_parsed = S3Bucket.parse_obj({
        "bucket_name": "my_bucket",
        "credentials": {
            "aws_access_key_id": "1234",
            "aws_access_key_secret": "5678",
        }
    })

    s3_bucket_no_creds_parsed = S3Bucket.parse_obj({
        "bucket_name": "my_bucket"
    })

    print(type(minio_bucket_parsed.credentials))
    print(type(s3_bucket_parsed.credentials))
    print(type(s3_bucket_no_creds_parsed.credentials))
<class 'prefect_aws.credentials.MinIOCredentials'>
<class 'prefect_aws.credentials.AwsCredentials'>
<class 'prefect_aws.credentials.AwsCredentials'>

Screenshots

Checklist