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

Fixed SlackAPIFileOperator to upload file and file content #17247

Merged
merged 8 commits into from
Aug 1, 2021
52 changes: 40 additions & 12 deletions airflow/providers/slack/operators/slack.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,15 +161,25 @@ class SlackAPIFileOperator(SlackAPIOperator):

.. code-block:: python

# Send file with filename and filetype
slack = SlackAPIFileOperator(
task_id="slack_file_upload",
dag=dag,
slack_conn_id="slack",
channel="#general",
initial_comment="Hello World!",
filename="hello_world.csv",
filetype="csv",
content="hello,world,csv,file",
filetype="csv"
)

# Send file content
slack = SlackAPIFileOperator(
task_id="slack_file_upload",
dag=dag,
slack_conn_id="slack",
channel="#general",
initial_comment="Hello World!",
content="file content in txt",
)

:param channel: channel in which to sent file on slack name (templated)
Expand All @@ -192,9 +202,9 @@ def __init__(
self,
channel: str = '#general',
initial_comment: str = 'No message has been set!',
filename: str = 'default_name.csv',
filetype: str = 'csv',
content: str = 'default,content,csv,file',
filename: str = None,
filetype: str = None,
content: str = None,
**kwargs,
) -> None:
self.method = 'files.upload'
Expand All @@ -203,13 +213,31 @@ def __init__(
self.filename = filename
self.filetype = filetype
self.content = content
self.file_params = {}
super().__init__(method=self.method, **kwargs)

def construct_api_call_params(self) -> Any:
self.api_params = {
'channels': self.channel,
'content': self.content,
'filename': self.filename,
'filetype': self.filetype,
'initial_comment': self.initial_comment,
}
if self.content is not None:
self.api_params = {
'channels': self.channel,
'content': self.content,
'initial_comment': self.initial_comment,
}
elif self.filename is not None:
self.api_params = {
'channels': self.channel,
'filename': self.filename,
'filetype': self.filetype,
'initial_comment': self.initial_comment,
}
self.file_params = {'file': self.filename}

def execute(self, **kwargs):
"""
The SlackAPIOperator calls will not fail even if the call is not unsuccessful.
It should not prevent a DAG from completing in success
"""
if not self.api_params:
self.construct_api_call_params()
slack = SlackHook(token=self.token, slack_conn_id=self.slack_conn_id)
slack.call(self.method, data=self.api_params, files=self.file_params)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need to pass filename in both files & data ?
api_params has a filename key that contains self.filename
file_params has file key that contains self.filename

28 changes: 23 additions & 5 deletions tests/providers/slack/operators/test_slack.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,21 +188,39 @@ def test_init_with_valid_params(self):
assert slack_api_post_operator.slack_conn_id == test_slack_conn_id

@mock.patch('airflow.providers.slack.operators.slack.SlackHook')
def test_api_call_params_with_default_args(self, mock_hook):
def test_api_call_params_with_content_args(self, mock_hook):
test_slack_conn_id = 'test_slack_conn_id'

slack_api_post_operator = SlackAPIFileOperator(
task_id='slack',
slack_conn_id=test_slack_conn_id,
task_id='slack', slack_conn_id=test_slack_conn_id, content='test-content'
)

slack_api_post_operator.execute()

expected_api_params = {
'channels': '#general',
'initial_comment': 'No message has been set!',
'content': 'test-content',
}
assert expected_api_params == slack_api_post_operator.api_params

@mock.patch('airflow.providers.slack.operators.slack.SlackHook')
def test_api_call_params_with_file_args(self, mock_hook):
test_slack_conn_id = 'test_slack_conn_id'

slack_api_post_operator = SlackAPIFileOperator(
task_id='slack', slack_conn_id=test_slack_conn_id, filename='test.csv', filetype='csv'
)

slack_api_post_operator.execute()

expected_api_params = {
'channels': '#general',
'initial_comment': 'No message has been set!',
'filename': 'default_name.csv',
'filename': 'test.csv',
'filetype': 'csv',
'content': 'default,content,csv,file',
}

expected_file_params = {'file': 'test.csv'}
assert expected_api_params == slack_api_post_operator.api_params
assert expected_file_params == slack_api_post_operator.file_params