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)
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"