Skip to content

Commit

Permalink
[AIRFLOW-7079] Remove redundant code for storing template_fields (apa…
Browse files Browse the repository at this point in the history
  • Loading branch information
kaxil committed Mar 17, 2020
1 parent d0cf421 commit de7e934
Show file tree
Hide file tree
Showing 4 changed files with 5 additions and 33 deletions.
12 changes: 1 addition & 11 deletions airflow/providers/google/cloud/operators/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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):
"""
Expand Down
12 changes: 1 addition & 11 deletions airflow/providers/qubole/operators/qubole.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion tests/serialization/test_dag_serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
},
Expand Down
12 changes: 2 additions & 10 deletions tests/test_utils/mock_operators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand All @@ -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):
"""
Expand Down

0 comments on commit de7e934

Please sign in to comment.