Python airflow.models.Variable.get() Examples

def test_variables_set_different_types(self):
        """Test storage of various data types"""
        # Set a dict
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'dict', '{"foo": "oops"}']))
        # Set a list
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'list', '["oops"]']))
        # Set str
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'str', 'hello string']))
        # Set int
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'int', '42']))
        # Set float
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'float', '42.0']))
        # Set true
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'true', 'true']))
        # Set false
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'false', 'false']))
        # Set none
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'null', 'null']))

        # Export and then import
        variable_command.variables_export(self.parser.parse_args([
            'variables', 'export', 'variables_types.json']))
        variable_command.variables_import(self.parser.parse_args([
            'variables', 'import', 'variables_types.json']))

        # Assert value
        self.assertEqual({'foo': 'oops'}, Variable.get('dict', deserialize_json=True))
        self.assertEqual(['oops'], Variable.get('list', deserialize_json=True))
        self.assertEqual('hello string', Variable.get('str'))  # cannot json.loads(str)
        self.assertEqual(42, Variable.get('int', deserialize_json=True))
        self.assertEqual(42.0, Variable.get('float', deserialize_json=True))
        self.assertEqual(True, Variable.get('true', deserialize_json=True))
        self.assertEqual(False, Variable.get('false', deserialize_json=True))
        self.assertEqual(None, Variable.get('null', deserialize_json=True))

        os.remove('variables_types.json') 

def render_template(**context):
    """ Render HTML template using questions metadata from S3 bucket """

    hook = S3Hook(aws_conn_id="s3_connection")
    file_content = hook.read_key(
        key=S3_FILE_NAME, bucket_name=Variable.get("S3_BUCKET")
    )
    questions = json.loads(file_content)

    root = os.path.dirname(os.path.abspath(__file__))
    env = Environment(loader=FileSystemLoader(root))
    template = env.get_template("email_template.html")
    html_content = template.render(questions=questions)

    # Push rendered HTML as a string to the Airflow metadata database
    # to make it available for the next task

    task_instance = context["task_instance"]
    task_instance.xcom_push(key="html_content", value=html_content) 

def create_slack_message(state):
        """
        :param state: State of dagrun such as 'Success' or 'Failure' etc.
        :return: Standard body content for the slack message
        """
        jinja_confs = {
            'dagid': '{{ dag.dag_id }}',
            'date': "{{ macros.ds_format(ts, '%Y-%m-%dT%H:%M:%S', '%x %I:%M:%S %p') }}",
            'url_date': "{{ macros.ds_format(ts, '%Y-%m-%dT%H:%M:%S', '%Y-%m-%d+%H%%3A%M%%3A%S') }}",
            'taskid': '{{ task.task_id }}',
            'owner': "{{ conf.get('operators', 'default_owner') }}",
            'url': "{{ conf.get('webserver', 'base_url') }}"
        }
        if state.lower() == 'failure':
            msg = 'DAG *{dagid}* for {date} has *failed* task {taskid} on the {owner} instance. Current task failures: ' \
                '<{url}/admin/taskinstance/?flt0_dag_id_equals={dagid}&flt1_state_equals=failed' \
                '&flt4_execution_date_equals={url_date}| Task Instances.>'
        if state.lower() == 'success':
            msg = 'DAG *{dagid}* has successfully processed {date} on the {owner} instance. <{url}/admin/airflow/gantt?' \
                  'dag_id={dagid}&execution_date={url_date}|Gantt chart.>'
        return msg.format(**jinja_confs) 

def test_variable_set_get_round_trip(self):
        Variable.set("tested_var_set_id", "Monday morning breakfast")
        self.assertEqual("Monday morning breakfast", Variable.get("tested_var_set_id")) 

def variables_get(args):
    """Displays variable by a given name"""
    try:
        if args.default is None:
            var = Variable.get(
                args.key,
                deserialize_json=args.json
            )
            print(var)
        else:
            var = Variable.get(
                args.key,
                deserialize_json=args.json,
                default_var=args.default
            )
            print(var)
    except (ValueError, KeyError) as e:
        print(str(e), file=sys.stderr)
        sys.exit(1) 

def create_subject_message(state):
        """
        :param state: State of dagrun such as 'Success' or 'Failure' etc.
        :return: Standard subject message for the email
        """
        jinja_confs = {
            'state': state,
            'owner': '{{ dag.owner }}',
            'dagid': '{{ dag.dag_id }}',
            'date': "{{ macros.ds_format(ts, '%Y-%m-%dT%H:%M:%S', '%x %I:%M:%S %p') }}",
            'instance': "{{ conf.get('operators', 'default_owner') }}"
        }
        return 'Airflow {state} of {owner} DAG {dagid} for date {date} on the {instance} instance.'.format(**jinja_confs) 

def test_backend_fallback_to_default_var(self):
        """
        Test if a default_var is defined and no backend has the Variable,
        the value returned is default_var
        """
        variable_value = Variable.get(key="test_var", default_var="new")
        self.assertEqual("new", variable_value) 

def create_dag(customer):
    """
    Takes a cust parameters dict, uses that to override default args creates DAG object
    
    Returns: DAG() Object
    """
    default_args = {
        'owner': 'airflow',
        'depends_on_past': False,
        'email': 'xyz@xyz.com',
        'retries': 1,
        'retry_delay': timedelta(minutes=5),
        'start_date': datetime(2017, 1, 1, 0, 0),
        'end_date': None
    }

    # This allows DAG parameters to be passed in from the Variable if a customer needs something specific overridden in their DAG
    # Consider how email being passed in from the customer object overrides email in the resulting replaced_args object
    replaced_args = {k: default_args[k] if customer.get(
        k, None) is None else customer[k] for k in default_args}


    dag_id = '{base_name}_{id}'.format(
        base_name='load_clickstream_data', id=customer['customer_id'])

    return DAG(dag_id=dag_id, default_args=replaced_args, schedule_interval=customer['schedule_interval']) 

def requires_authentication(function):
    @wraps(function)
    def decorated(*args, **kwargs):
        PUBLIC_KEY = "jwt_backend_public_key"
        ALGORITHM = "jwt_backend_algorithm"
        try:
            json_data = {k: v for k, v in request.get_json(force=True).copy().items() if k != "token"}
            decoded_token = jwt.decode(request.get_json(force=True)["token"], Variable.get(PUBLIC_KEY), algorithms=Variable.get(ALGORITHM))
            assert (json_data == decoded_token)
        except Exception:
            return Response("Failed to verify data", 403)
        return function(*args, **kwargs)
    return decorated 

def sign_with_jwt(data, private_key=None, algorithm=None):
    try:
        data = jwt.encode(payload=data,
                          key=private_key or Variable.get(PRIVATE_KEY),
                          algorithm=algorithm or Variable.get(ALGORITHM)).decode("utf-8")
    except Exception as e:
        logger.warning("Failed to sign data with JWT key: {}".format(e))
    return data 

def read_var(var_name, var_prefix=None, required=False, **kwargs):
    full_var_name = f'{var_prefix}{var_name}' if var_prefix is not None else var_name
    var = Variable.get(full_var_name, '')
    var = var if var != '' else None
    if var is None:
        var = kwargs.get(var_name)
    if required and var is None:
        raise ValueError(f'{full_var_name} variable is required')
    return var 

def write_questions_to_s3():
    hook = S3Hook(aws_conn_id="s3_connection")
    hook.load_string(
        string_data=filter_questions(),
        key=S3_FILE_NAME,
        bucket_name=Variable.get("S3_BUCKET"),
        replace=True,
    ) 

def call_stack_overflow_api() -> dict:
    """ Get first 100 questions created two days ago sorted by user votes """

    stack_overflow_question_url = Variable.get("STACK_OVERFLOW_QUESTION_URL")

    today = datetime.now()
    three_days_ago = today - timedelta(days=7)
    two_days_ago = today - timedelta(days=5)

    payload = {
        "fromdate": int(datetime.timestamp(three_days_ago)),
        "todate": int(datetime.timestamp(two_days_ago)),
        "sort": "votes",
        "site": "stackoverflow",
        "order": "desc",
        "tagged": Variable.get("TAG"),
        "client_id": Variable.get("STACK_OVERFLOW_CLIENT_ID"),
        "client_secret": Variable.get("STACK_OVERFLOW_CLIENT_SECRET"),
        "key": Variable.get("STACK_OVERFLOW_KEY"),
    }

    response = requests.get(stack_overflow_question_url, params=payload)

    for question in response.json().get("items", []):
        yield {
            "question_id": question["question_id"],
            "title": question["title"],
            "is_answered": question["is_answered"],
            "link": question["link"],
            "owner_reputation": question["owner"].get("reputation", 0),
            "score": question["score"],
            "tags": question["tags"],
        } 

def get_variable(variable_key: str) -> Response:
    """
    Get a variables by key
    """
    try:
        var = Variable.get(variable_key)
    except KeyError:
        raise NotFound("Variable not found")
    return variable_schema.dump({"key": variable_key, "val": var}) 

def test_variables_isolation(self):
        """Test isolation of variables"""
        tmp1 = tempfile.NamedTemporaryFile(delete=True)
        tmp2 = tempfile.NamedTemporaryFile(delete=True)

        # First export
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'foo', '{"foo":"bar"}']))
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'bar', 'original']))
        variable_command.variables_export(self.parser.parse_args([
            'variables', 'export', tmp1.name]))

        first_exp = open(tmp1.name, 'r')

        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'bar', 'updated']))
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'foo', '{"foo":"oops"}']))
        variable_command.variables_delete(self.parser.parse_args([
            'variables', 'delete', 'foo']))
        variable_command.variables_import(self.parser.parse_args([
            'variables', 'import', tmp1.name]))

        self.assertEqual('original', Variable.get('bar'))
        self.assertEqual('{\n  "foo": "bar"\n}', Variable.get('foo'))

        # Second export
        variable_command.variables_export(self.parser.parse_args([
            'variables', 'export', tmp2.name]))

        second_exp = open(tmp2.name, 'r')
        self.assertEqual(first_exp.read(), second_exp.read())

        # Clean up files
        second_exp.close()
        first_exp.close() 

def test_get_variable_missing_variable(self):
        with self.assertRaises(SystemExit):
            variable_command.variables_get(self.parser.parse_args([
                'variables', 'get', 'no-existing-VAR'])) 

def test_get_variable_default_value(self):
        with redirect_stdout(io.StringIO()) as stdout:
            variable_command.variables_get(self.parser.parse_args([
                'variables', 'get', 'baz', '--default', 'bar']))
            self.assertEqual("bar\n", stdout.getvalue()) 

def test_variables_get(self):
        Variable.set('foo', {'foo': 'bar'}, serialize_json=True)

        with redirect_stdout(io.StringIO()) as stdout:
            variable_command.variables_get(self.parser.parse_args([
                'variables', 'get', 'foo']))
            self.assertEqual('{\n  "foo": "bar"\n}\n', stdout.getvalue()) 

def test_variables_set(self):
        """Test variable_set command"""
        variable_command.variables_set(self.parser.parse_args([
            'variables', 'set', 'foo', 'bar']))
        self.assertIsNotNone(Variable.get("foo"))
        self.assertRaises(KeyError, Variable.get, "foo1") 

def test_variable_setdefault_existing_json(self):
        key = "tested_var_setdefault_2_id"
        value = {"city": 'Paris', "Happiness": True}
        Variable.set(key, value, serialize_json=True)
        val = Variable.setdefault(key, value, deserialize_json=True)
        # Check the returned value, and the stored value are handled correctly.
        self.assertEqual(value, val)
        self.assertEqual(value, Variable.get(key, deserialize_json=True)) 

def test_variable_setdefault_round_trip_json(self):
        key = "tested_var_setdefault_2_id"
        value = {"city": 'Paris', "Happiness": True}
        Variable.setdefault(key, value, deserialize_json=True)
        self.assertEqual(value, Variable.get(key, deserialize_json=True)) 

def test_variable_setdefault_round_trip(self):
        key = "tested_var_setdefault_1_id"
        value = "Monday morning breakfast in Paris"
        Variable.setdefault(key, value)
        self.assertEqual(value, Variable.get(key)) 

def test_get_non_existing_var_should_not_deserialize_json_default(self):
        default_value = "}{ this is a non JSON default }{"
        self.assertEqual(default_value, Variable.get("thisIdDoesNotExist",
                                                     default_var=default_value,
                                                     deserialize_json=True)) 

def test_get_non_existing_var_with_none_default_should_return_none(self):
        self.assertIsNone(Variable.get("thisIdDoesNotExist", default_var=None)) 

def test_get_non_existing_var_should_return_default(self):
        default_value = "some default val"
        self.assertEqual(default_value, Variable.get("thisIdDoesNotExist",
                                                     default_var=default_value)) 

def test_variable_set_existing_value_to_blank(self):
        test_value = 'Some value'
        test_key = 'test_key'
        Variable.set(test_key, test_value)
        Variable.set(test_key, '')
        self.assertEqual('', Variable.get('test_key')) 

def test_variable_set_get_round_trip_json(self):
        value = {"a": 17, "b": 47}
        Variable.set("tested_var_set_id", value, serialize_json=True)
        self.assertEqual(value, Variable.get("tested_var_set_id", deserialize_json=True))