pythonpydanticpydantic-settings

Pydantic CLIApp/CLISubCommand with Env Vars


I am currently building a CLI using Pydantic. One of the sub commands has 2 parameters I would like to load via an Env variable. I was able to load the environment variables on its own when I instantiated the class directly but now that it is a CliSubCommand. I am having issues loading the two parameters via Env variables. I also was able to get the CLI working when AWSStsGcp class was inheriting from BaseModel by providing via the command line. Pydantic is throwing an error saying the two AWS flags are required. They are located in the aws_sts_gcp.py model -> aws_access_key and aws_secret_key.

I also don't want it to be a mandatory global config, since it is currently only relevant for one command.

root_tool.py


    from pydantic import BaseModel
    from pydantic_settings import CliSubCommand, CliApp
    from python_tools.gcp.models.aws_sts_gcp import AwsStsGcp
    
    class DummyCommand(BaseModel):
        project_id: str
    
        def cli_cmd(self) -> None:
            print(f'This is a dummy command.{self.project_id}"')
        
    class Tool(BaseModel):
        aws_sts: CliSubCommand[AwsStsGcp]
        dummy: CliSubCommand[DummyCommand]
    
        def cli_cmd(self) -> None:
            CliApp.run_subcommand(self)

aws_sts_gcp.py

    from pydantic import SecretStr
    from pydantic_settings import BaseSettings, SettingsConfigDict
    
    from python_tools.gcp.aws_gcs_transfer import create_aws_transfer_job 
    
    class AwsStsGcp(BaseSettings, cli_parse_args=True):
        model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8')
        destination_bucket: str
        src_bucket: str
        manifest_path: str | None = None
        aws_access_key: SecretStr
        aws_secret_key: SecretStr
        tranfer_name: str
        project_id: str
    
        def cli_cmd(self) -> None:
            create_aws_transfer_job(self)

cli_runner.py

from python_tools.gcp.models.root_tool import Tool
from pydantic_settings import CliApp

CliApp.run(Tool)

aws_gcs_transfer.py

from google.cloud.storage_transfer_v1 import (
    StorageTransferServiceClient,
    TransferJob,
    TransferSpec,
    TransferManifest,
    AwsS3Data,
    AwsAccessKey,
    GcsData,
    RunTransferJobRequest,
    CreateTransferJobRequest
)
#from python_tools.gcp.models.aws_sts_gcp import AwsStsGcp
from python_tools.consts import timestr
from python_tools.logging import logger

import time

def create_aws_transfer_job(transfer_details) -> None:
    s3_config = None

    transfer_manifest = None 

    client = StorageTransferServiceClient()
    s3_config = AwsS3Data(
        bucket_name=transfer_details.src_bucket,
        aws_access_key=AwsAccessKey(
            access_key_id=transfer_details.aws_access_key.get_secret_value(), secret_access_key=transfer_details.aws_secret_key.get_secret_value()
        )
    )

    gcs_dest = GcsData(bucket_name=transfer_details.destination_bucket)

    if transfer_details.manifest_path is not None:
        transfer_manifest = TransferManifest(location=transfer_details.manifest_path)
    
    sts_spec = TransferSpec(gcs_data_sink=gcs_dest, aws_s3_data_source=s3_config, transfer_manifest=transfer_manifest)
    timestamp = time.strftime(timestr)
    name = f"transferJobs/{transfer_details.tranfer_name}-{timestamp}"
    description = "Automated STS Job created from Python Tools."
    sts_job = TransferJob(
        project_id=transfer_details.project_id,
        name=name,
        description=description,
        transfer_spec=sts_spec,
        status=TransferJob.Status.ENABLED,
    )
    job_request = CreateTransferJobRequest(transfer_job=sts_job)

    logger.info(f"Starting Transfer Job for Job ID: {name}")
    transfer_request = RunTransferJobRequest(project_id=transfer_details.project_id, job_name=name)

    client.create_transfer_job(request=job_request)
    client.run_transfer_job(request=transfer_request)

.env

AWS_ACCESS_KEY = "test"
AWS_SECRET_KEY = "test"

I have also tried using BaseSettings at the root like this.

from pydantic import BaseModel
from pydantic_settings import CliSubCommand, CliApp, BaseSettings, SettingsConfigDict

from python_tools.gcp.models.aws_sts_gcp import AwsStsGcp

class DummyCommand(BaseModel):
    project_id: str

    def cli_cmd(self) -> None:
        print(f'This is a dummy command.{self.project_id}"')
    
class Tool(BaseSettings, cli_parse_args=True):
    model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8', extra='ignore')

    aws_sts: CliSubCommand[AwsStsGcp]
    dummy: CliSubCommand[DummyCommand]

    def cli_cmd(self) -> None:
        CliApp.run_subcommand(self)

from pydantic import SecretStr, BaseModel, Field  
from python_tools.gcp.aws_gcs_transfer import create_aws_transfer_job 

class AwsStsGcp(BaseModel):
    destination_bucket: str
    src_bucket: str
    manifest_path: str | None = None
    aws_access_key: SecretStr = Field(alias='AWS_ACCESS_KEY', env="AWS_ACCESS_KEY")
    aws_secret_key: SecretStr = Field(alias='AWS_SECRET_KEY',  env="AWS_SECRET_KEY")
    tranfer_name: str
    project_id: str

    def cli_cmd(self) -> None:
        create_aws_transfer_job(self)

Solution

  • I realized I needed to use the env_nested_delimiter to set variables in nested models. And update the Env variable name to

    AWS_STS__AWS_ACCESS_KEY = "test"
    AWS_STS__AWS_SECRET_KEY = "test"