Thread
#prefect-community
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    Hey. How can I create a custom implementation of PrefectSecretthat accept a Parameteras an argument? is that possible?
    Anna Geller

    Anna Geller

    10 months ago
    could you describe your use case a bit more to understand how would you combine parameter with Secret? is it that you get the secret name as parameter?
    Kevin Kho

    Kevin Kho

    10 months ago
    I was gonna link you to the other thread but I realized it was you also lol. I think you can try
    with Flow(...) as flow:
        x = Parameter("x")
        name = create_secret_name(x)
        s1 = PrefectSecet(name)
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    class IDMappingDBSecret(PrefectSecret):
        environment: str
    
        def __init__(self, name=None, env=None, **kwargs):
            self.environment = env
            super().__init__(name, **kwargs)
    
        def run(self, name: str = None):
            logger = prefect.context.get("logger")
            name: str = name or self.secret_name
            if self.environment == "dev":
                name = f'dev_{name}'
            <http://logger.info|logger.info>(f'Secret {name} is going to be retrieved')
            super().run(name)
    that is the implementation I ’ve done of
    PrefectSecret
    the idea is to change the name of the secret , depending on the value of a parameter that is called
    env
    with Flow('flow_sample', schedule=None) as flow:
        env = Parameter('env', required=True)
        env_secret = IDMappingDBSecret(name='STAGE_ID_MAPPING_DB_URL', env=env)
        print_result()
    this a flow sample I ’ve done to test this behavior
    Kevin Kho

    Kevin Kho

    10 months ago
    I think this might be easier:
    class IDMappingDBSecret(PrefectSecret):
        environment: str
    
        def __init__(self, name=None, env=None, **kwargs):
            self.environment = env
            new_name = self.environment + name
            logger = prefect.context.get("logger")
            <http://logger.info|logger.info>(f'Secret {new_name} is going to be retrieved')
            super().__init__(new_name, **kwargs)
    Anna Geller

    Anna Geller

    10 months ago
    alternatively, you could get the parameter value from context perhaps?
    prefect.context["parameters"]
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    Oh, that could work. thanks, I will test it
    I am getting this
    "'GraphQLResult' object has no attribute 'name'"
    with this code:
    class IDMappingDBSecret(PrefectSecret):
        environment: str
    
        def __init__(self, name=None, **kwargs):
            self.environment = prefect.config.get("environment")
            super().__init__(name=name, **kwargs)
    
        def run(self, name: str = None):
            logger = prefect.context.get("logger")
            <http://logger.info|logger.info>(f"Environment is {self.environment}")
            name: str = name or self.secret_name
            if self.environment == "dev":
                name = f'{self.environment.upper}_{name}'
            <http://logger.info|logger.info>(f'Secret {name} is going to be retrieved')
            super().run(name=name)
    
    @task
    def print_result():
        logger = prefect.context.get("logger")
        <http://logger.info|logger.info>('Result is 6')
    
    
    with Flow('flow_sample', schedule=None) as flow:
        env_secret = IDMappingDBSecret(name='STAGE_ID_MAPPING_DB_URL')
        print_result()
    What I am doing wrong?
    I ’ve decided to try
    prefect.config.get("environment")
    instead of create a new parameter. Is environment associated to the
    label
    provided in the UI (the one that specifies which agent is going to take it)?
    Kevin Kho

    Kevin Kho

    10 months ago
    I think just do more in the init and it will be easier. No need to define the run. This worked on cloud for me:
    from prefect.tasks.secrets import PrefectSecret
    import prefect
    from prefect import Flow, task
    
    class IDMappingDBSecret(PrefectSecret):
    
        def __init__(self, name=None, environment="dev", **kwargs):
            self.environment = environment
            if self.environment == "stage":
                _name = f'{self.environment.upper()}_{name}'
            else:
                _name = name
            super().__init__(name=_name, **kwargs)
    
    @task
    def print_result(secret):
        logger = prefect.context.get("logger")
        <http://logger.info|logger.info>(secret)
        return
    
    with Flow('flow_sample', schedule=None) as flow:
        env_secret = IDMappingDBSecret('ID_MAPPING_DB_URL', environment="stage")
        print_result(env_secret)
    
    flow.run()
    Anna Geller

    Anna Geller

    10 months ago
    @Santiago Gonzalez I think you can combine what Kevin suggested with init and your environment from config like so:
    from prefect.tasks.secrets import PrefectSecret
    import prefect
    from prefect import Flow, task
    
    
    class EnvSecret(PrefectSecret):
        def __init__(self, name=None, **kwargs):
            self.environment = prefect.config.get("environment")
            if self.environment in ["stage", "dev"]:
                _name = f"{self.environment.upper()}_{name}"
            else:
                _name = name
            super().__init__(name=_name, **kwargs)
    
    
    env_secret = EnvSecret("DB_NAME")
    
    
    @task
    def print_result(secret):
        logger = prefect.context.get("logger")
        <http://logger.info|logger.info>(secret)
    
    
    with Flow("flow_sample") as flow:
        secret = env_secret(task_args={"name": "DB_NAME"})
        print_result(secret)
    
    if __name__ == "__main__":
        flow.run()
    Running with value “prod”:
    [2021-11-13 13:11:35+0100] INFO - prefect.TaskRunner | Task 'DB_NAME': Starting task run...
    [2021-11-13 13:11:36+0100] INFO - prefect.TaskRunner | Task 'DB_NAME': Finished task run for task with final state: 'Success'
    [2021-11-13 13:11:36+0100] INFO - prefect.TaskRunner | Task 'print_result': Starting task run...
    [2021-11-13 13:11:36+0100] INFO - prefect.print_result | PROD_DB
    Running with value “dev” in config.toml:
    [2021-11-13 13:11:47+0100] INFO - prefect.TaskRunner | Task 'DEV_DB_NAME': Starting task run...
    [2021-11-13 13:11:47+0100] INFO - prefect.TaskRunner | Task 'DEV_DB_NAME': Finished task run for task with final state: 'Success'
    [2021-11-13 13:11:47+0100] INFO - prefect.TaskRunner | Task 'print_result': Starting task run...
    [2021-11-13 13:11:47+0100] INFO - prefect.print_result | DEV_DB
    Is that what you intended to do?
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    Yes, it is very similar. The label configured in the agent, is the same as environment???
    Thanks @Anna Geller
    Anna Geller

    Anna Geller

    10 months ago
    no, you would need to specifically set it in your config.toml. This would be a custom variable you set.
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    Hey @Anna Geller I realized that this command
    os.environ.get('environment')
    When I execute that command inside of a task, it gets the value configured in the agent. However, when I try to get the value (with the same command) inside of the
    PrefectSecret
    class, it has no value. What do you think could be a way to get the value in the prefect secret class? I leave you a piece of the flow code:
    class IDMappingDBSecret(PrefectSecret):
        environment: str
    
        def __init__(self, name=None, **kwargs):
            env = os.environ.get('environment')
            if env in ["dev"]:
                _name = f"{env.upper()}_{name}"
            else:
                _name = name
            super().__init__(name=_name, **kwargs)
    
    @task
    def print_result(secret):
        env = os.environ.get('environment')
        logger = prefect.context.get("logger")
        <http://logger.info|logger.info>(f"Environment {env}")
        <http://logger.info|logger.info>(f'Result is {secret}')
    
    with Flow('flow_sample', schedule=None) as flow:
        secret = IDMappingDBSecret('DB_URL')
        print_result(secret)
    Thanks in advance
    Kevin Kho

    Kevin Kho

    10 months ago
    Looking at this again, I think pickle based storage might serialize the env variable. I guess you need to do
    os.environ
    in the
    run
    method. Not sure if that would work though. I think the easier approach is to use script-based storage so that the Flow is not serialized. See this and this
    Anna Geller

    Anna Geller

    10 months ago
    @Santiago Gonzalez there must have been some mix-up here - we were talking about retrieving this value from config.toml rather than from env variable? Can you have a look at the previous posts?
    elf.environment = prefect.config.get("environment")
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    I tried with the code that you provided and it didn’t work out. I will try in the run method then, or with the other solution.
    Anna Geller

    Anna Geller

    10 months ago
    I’m 100% sure that the method provided here will work with a local agent, as long as you set the variable in the config.toml: https://prefect-community.slack.com/archives/CL09KU1K7/p1636805732056300?thread_ts=1636749897.042500&amp;cid=CL09KU1K7 But feel free to experiment and see what works best for you
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    Oh. Ok, the thing is that I am setting the variable in the EC2 Agent like this:
    screen -dmS prefect-screen python3.7 -m prefect agent docker start --token {TOKEN}  --label dev --env environment=dev
    Anna Geller

    Anna Geller

    10 months ago
    sure, try this and make sure that this environment variable also exist in the environment from which you register your flow and this could work, too. btw, token is deprecated in favour of API keys, so it should be --key YOUR_KEY instead of --token
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    thanks, I will try then.
    This is how I have config.tml
    environment = "dev"
    [cloud]
       [cloud.agent]
       name = "agent"
       labels = ["dev"]
    
       # Set to `DEBUG` for verbose logging
       level = "INFO"
    With this piece of code:
    class IDMappingDBSecret(PrefectSecret):
        environment: str
    
        def __init__(self, name=None, **kwargs):
            env = prefect.config.get("environment")
            if env in ["dev"]:
                _name = f"{env.upper()}_{name}"
            else:
                _name = name
            super().__init__(name=_name, **kwargs)
    
    @task
    def print_result(secret):
        logger = prefect.context.get("logger")
        <http://logger.info|logger.info>(f'Result is {secret}')
    
    with Flow('flow_sample', schedule=None) as flow:
        secret = IDMappingDBSecret('DB_URL')
        print_result(secret)
    And yet, it is not working as expected. Is something wrong there?
    Kevin Kho

    Kevin Kho

    10 months ago
    What is the issue you run into?
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    That the value is not being retrieved using
    prefect.config.get('environment')
    Kevin Kho

    Kevin Kho

    10 months ago
    Will try this in a bit
    Anna Geller

    Anna Geller

    10 months ago
    can you try this?
    from prefect.tasks.secrets import PrefectSecret
    import prefect
    from prefect import Flow, task
    
    
    class IDMappingDBSecret(PrefectSecret):
        def __init__(self, name=None, **kwargs):
            env = prefect.config.get("environment")
            if env in ["dev"]:
                _name = f"{env.upper()}_{name}"
            else:
                _name = name
            super().__init__(name=_name, **kwargs)
    
    
    env_secret = IDMappingDBSecret("DB_NAME")
    
    @task
    def print_result(secret):
        logger = prefect.context.get("logger")
        <http://logger.info|logger.info>('Result is %s', secret)
    
    
    with Flow('flow_sample') as flow:
        secret = env_secret(task_args={"name": "DB_NAME"})
        print_result(secret)
    
    if __name__ == "__main__":
        flow.run()
    @Santiago Gonzalez to give you some explanation: initializing the task and then calling it within the Flow block with a custom task name passed via task_args ensures that your task name remains the same in the flow structure, regardless of whether your secret name is DEV_DB or PROD_DB. As long as this
    environment = "dev"
    is set in the config.toml, it should work with the DEV_DB secret. You would need to set it in the environment from which you start the agent
    Santiago Gonzalez

    Santiago Gonzalez

    10 months ago
    @Anna Geller this is the thing: I am running the flow inside a EC2 Docker Agent (not a LocalAgent), and in this context, the environment variable is not being getting. I ’ve been trying to get this working , setting the variable in config.toml (as an argument in agent docker start as well) but it is not working, still getting
    None
    value as a result when I do
    prefect.config.get('environment')
    .
    I figured it out. I got it working using
    os.environ['environment']
    in the run method. thanks