incubator-airflow git commit: [AIRFLOW-464] Add setdefault method to Variable

[sa...@apache.org]

Repository: incubator-airflow
Updated Branches:
  refs/heads/master c6dd4d457 -> dd1f50e59


[AIRFLOW-464] Add setdefault method to Variable

In order to assist with environment migrations, we
added a setdefault method to
the Variable object. This allows default variables
to be created (and then
edited) with less chance for typos/copy+paste
bugs.

Also changed Variable.get ValueError exception to
a KeyError

Variable.setdefault(key, default,
deserialize_json=[True|False]) returns either the
value stored in Variable(key) or sets
Variable(key) = default and returns default.

Dear Airflow Maintainers,

Please accept this PR that addresses the following
issues:
- _https://issues.apache.org/jira/browse/AIRFLOW-
464_

This was changed from adding a create_if_none flag
to Variable.get based on feedback on this PR.

Testing Done:
- Added a test to test/core.py to cover this
functionality

Closes #1765 from btallman/CreateIfNone_feature


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/dd1f50e5
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/dd1f50e5
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/dd1f50e5

Branch: refs/heads/master
Commit: dd1f50e59008e95f457ce47ba877f631d5187530
Parents: c6dd4d4
Author: Benjamin Tallman <bt...@gmail.com>
Authored: Mon Nov 14 23:04:04 2016 -0800
Committer: Siddharth Anand <si...@yahoo.com>
Committed: Mon Nov 14 23:04:21 2016 -0800

----------------------------------------------------------------------
 airflow/models.py | 32 +++++++++++++++++++++++++++++++-
 tests/core.py     | 12 ++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/dd1f50e5/airflow/models.py
----------------------------------------------------------------------
diff --git a/airflow/models.py b/airflow/models.py
index 3509a0c..e64d04f 100755
--- a/airflow/models.py
+++ b/airflow/models.py
@@ -3366,6 +3366,36 @@ class Variable(Base):
                        descriptor=property(cls.get_val, cls.set_val))
 
     @classmethod
+    def setdefault(cls, key, default, deserialize_json=False):
+        """
+        Like a Python builtin dict object, setdefault returns the current value
+        for a key, and if it isn't there, stores the default value and returns it.
+
+        :param key: Dict key for this Variable
+        :type key: String
+        :param: default: Default value to set and return if the variable
+        isn't already in the DB
+        :type: default: Mixed
+        :param: deserialize_json: Store this as a JSON encoded value in the DB
+         and un-encode it when retrieving a value
+        :return: Mixed
+        """
+        default_sentinel = object()
+        obj = Variable.get(key, default_var=default_sentinel, deserialize_json=False)
+        if obj is default_sentinel:
+            if default is not None:
+                Variable.set(key, default, serialize_json=deserialize_json)
+                return default
+            else:
+                raise ValueError('Default Value must be set')
+        else:
+            if deserialize_json:
+                return json.loads(obj.val)
+            else:
+                return obj.val
+
+
+    @classmethod
     @provide_session
     def get(cls, key, default_var=None, deserialize_json=False, session=None):
         obj = session.query(cls).filter(cls.key == key).first()
@@ -3373,7 +3403,7 @@ class Variable(Base):
             if default_var is not None:
                 return default_var
             else:
-                raise ValueError('Variable {} does not exist'.format(key))
+                raise KeyError('Variable {} does not exist'.format(key))
         else:
             if deserialize_json:
                 return json.loads(obj.val)

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/dd1f50e5/tests/core.py
----------------------------------------------------------------------
diff --git a/tests/core.py b/tests/core.py
index effc63d..b52945a 100644
--- a/tests/core.py
+++ b/tests/core.py
@@ -691,6 +691,18 @@ class CoreTest(unittest.TestCase):
                                              default_var=default_value,
                                              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)
+        assert value == Variable.get(key)
+
+    def test_variable_setdefault_round_trip_json(self):
+        key = "tested_var_setdefault_2_id"
+        value = {"city": 'Paris', "Hapiness": True}
+        Variable.setdefault(key, value, deserialize_json=True)
+        assert value == Variable.get(key, deserialize_json=True)
+
     def test_parameterized_config_gen(self):
 
         cfg = configuration.parameterized_config(configuration.DEFAULT_CONFIG)