[AIRFLOW-7079] Remove redundant code for storing template_fields (#7750)

kaxil · web-flow · commit de7e934ca3f2 · 2020-03-17T19:03:22.000Z
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
@@ -23,7 +23,7 @@
 
 import json
 import warnings
-from typing import Any, Dict, FrozenSet, Iterable, List, Optional, SupportsAbs, Union
+from typing import Any, Dict, Iterable, List, Optional, SupportsAbs, Union
 
 import attr
 from googleapiclient.errors import HttpError
@@ -456,9 +456,6 @@ class BigQueryExecuteQueryOperator(BaseOperator):
     template_ext = ('.sql', )
     ui_color = '#e4f0e8'
 
-    # The _serialized_fields are lazily loaded when get_serialized_fields() method is called
-    __serialized_fields: Optional[FrozenSet[str]] = None
-
     @property
     def operator_extra_links(self):
         """
@@ -592,13 +589,6 @@ def on_kill(self):
             self.log.info('Cancelling running query')
             self.hook.cancel_query()
 
-    @classmethod
-    def get_serialized_fields(cls):
-        """Serialized BigQueryOperator contain exactly these fields."""
-        if not cls.__serialized_fields:
-            cls.__serialized_fields = frozenset(super().get_serialized_fields() | {"sql"})
-        return cls.__serialized_fields
-
 
 class BigQueryCreateEmptyTableOperator(BaseOperator):
     """
diff --git a/airflow/providers/qubole/operators/qubole.py b/airflow/providers/qubole/operators/qubole.py
@@ -17,7 +17,7 @@
 # under the License.
 """Qubole operator"""
 import re
-from typing import FrozenSet, Iterable, Optional
+from typing import Iterable
 
 from airflow.hooks.base_hook import BaseHook
 from airflow.models import BaseOperator, BaseOperatorLink
@@ -181,9 +181,6 @@ class QuboleOperator(BaseOperator):
         QDSLink(),
     )
 
-    # The _serialized_fields are lazily loaded when get_serialized_fields() method is called
-    __serialized_fields: Optional[FrozenSet[str]] = None
-
     @apply_defaults
     def __init__(self, qubole_conn_id="qubole_default", *args, **kwargs):
         self.args = args
@@ -243,10 +240,3 @@ def __setattr__(self, name, value):
             self.kwargs[name] = value
         else:
             object.__setattr__(self, name, value)
-
-    @classmethod
-    def get_serialized_fields(cls):
-        """Serialized QuboleOperator contain exactly these fields."""
-        if not cls.__serialized_fields:
-            cls.__serialized_fields = frozenset(super().get_serialized_fields() | {"qubole_conn_id"})
-        return cls.__serialized_fields
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
@@ -81,7 +81,7 @@
                 "_operator_extra_links": [{"tests.test_utils.mock_operators.CustomOpLink": {}}],
                 "ui_color": "#fff",
                 "ui_fgcolor": "#000",
-                "template_fields": [],
+                "template_fields": ['bash_command'],
                 "_task_type": "CustomOperator",
                 "_task_module": "tests.test_utils.mock_operators",
             },
diff --git a/tests/test_utils/mock_operators.py b/tests/test_utils/mock_operators.py
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import FrozenSet, NamedTuple, Optional
+from typing import NamedTuple
 
 import attr
 
@@ -103,8 +103,7 @@ def get_link(self, operator, dttm):
 
 class CustomOperator(BaseOperator):
 
-    # The _serialized_fields are lazily loaded when get_serialized_fields() method is called
-    __serialized_fields: Optional[FrozenSet[str]] = None
+    template_fields = ['bash_command']
 
     @property
     def operator_extra_links(self):
@@ -128,13 +127,6 @@ def execute(self, context):
         self.log.info("Hello World!")
         context['task_instance'].xcom_push(key='search_query', value="dummy_value")
 
-    @classmethod
-    def get_serialized_fields(cls):
-        """Stringified CustomOperator contain exactly these fields."""
-        if not cls.__serialized_fields:
-            cls.__serialized_fields = frozenset(super().get_serialized_fields() | {"bash_command"})
-        return cls.__serialized_fields
-
 
 class GoogleLink(BaseOperatorLink):
     """
