Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Mar 17, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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