diff --git a/airflow/providers/slack/hooks/slack.py b/airflow/providers/slack/hooks/slack.py index cb12132eafd72..5af1e1e5734a2 100644 --- a/airflow/providers/slack/hooks/slack.py +++ b/airflow/providers/slack/hooks/slack.py @@ -18,7 +18,8 @@ import json import warnings -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union from slack_sdk import WebClient from slack_sdk.errors import SlackApiError @@ -213,6 +214,58 @@ def call(self, api_method: str, **kwargs) -> "SlackResponse": """ return self.client.api_call(api_method, **kwargs) + def send_file( + self, + *, + channels: Optional[Union[str, Sequence[str]]] = None, + file: Optional[Union[str, Path]] = None, + content: Optional[str] = None, + filename: Optional[str] = None, + filetype: Optional[str] = None, + initial_comment: Optional[str] = None, + title: Optional[str] = None, + ) -> "SlackResponse": + """ + Create or upload an existing file. + + :param channels: Comma-separated list of channel names or IDs where the file will be shared. + If omitting this parameter, then file will send to workspace. + :param file: Path to file which need to be sent. + :param content: File contents. If omitting this parameter, you must provide a file. + :param filename: Displayed filename. + :param filetype: A file type identifier. + :param initial_comment: The message text introducing the file in specified ``channels``. + :param title: Title of file. + + .. seealso:: + - `Slack API files.upload method `_ + - `File types `_ + """ + if not ((not file) ^ (not content)): + raise ValueError("Either `file` or `content` must be provided, not both.") + elif file: + file = Path(file) + with open(file, "rb") as fp: + if not filename: + filename = file.name + return self.client.files_upload( + file=fp, + filename=filename, + filetype=filetype, + initial_comment=initial_comment, + title=title, + channels=channels, + ) + + return self.client.files_upload( + content=content, + filename=filename, + filetype=filetype, + initial_comment=initial_comment, + title=title, + channels=channels, + ) + def test_connection(self): """Tests the Slack API connection. diff --git a/airflow/providers/slack/operators/slack.py b/airflow/providers/slack/operators/slack.py index 0c4ec53df2b85..0a1087bb99ce0 100644 --- a/airflow/providers/slack/operators/slack.py +++ b/airflow/providers/slack/operators/slack.py @@ -16,10 +16,13 @@ # specific language governing permissions and limitations # under the License. import json -from typing import Any, Dict, List, Optional, Sequence +import warnings +from typing import Any, Dict, List, Optional, Sequence, Union +from airflow.compat.functools import cached_property from airflow.models import BaseOperator from airflow.providers.slack.hooks.slack import SlackHook +from airflow.utils.log.secrets_masker import mask_secret class SlackAPIOperator(BaseOperator): @@ -47,13 +50,19 @@ def __init__( **kwargs, ) -> None: super().__init__(**kwargs) - - self.token = token # type: Optional[str] - self.slack_conn_id = slack_conn_id # type: Optional[str] + if token: + mask_secret(token) + self.token = token + self.slack_conn_id = slack_conn_id self.method = method self.api_params = api_params + @cached_property + def hook(self) -> SlackHook: + """Slack Hook.""" + return SlackHook(token=self.token, slack_conn_id=self.slack_conn_id) + def construct_api_call_params(self) -> Any: """ Used by the execute function. Allows templating on the source fields @@ -70,14 +79,9 @@ def construct_api_call_params(self) -> Any: ) 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, json=self.api_params) + self.hook.call(self.method, json=self.api_params) class SlackAPIPostOperator(SlackAPIOperator): @@ -144,7 +148,7 @@ def construct_api_call_params(self) -> Any: class SlackAPIFileOperator(SlackAPIOperator): """ - Send a file to a slack channel + Send a file to a slack channels Examples: .. code-block:: python @@ -154,7 +158,7 @@ class SlackAPIFileOperator(SlackAPIOperator): task_id="slack_file_upload_1", dag=dag, slack_conn_id="slack", - channel="#general", + channels="#general,#random", initial_comment="Hello World!", filename="/files/dags/test.txt", filetype="txt", @@ -165,63 +169,67 @@ class SlackAPIFileOperator(SlackAPIOperator): task_id="slack_file_upload_2", dag=dag, slack_conn_id="slack", - channel="#general", + channels="#general", initial_comment="Hello World!", content="file content in txt", ) - :param channel: channel in which to sent file on slack name (templated) + :param channels: Comma-separated list of channel names or IDs where the file will be shared. + If set this argument to None, then file will send to associated workspace. (templated) :param initial_comment: message to send to slack. (templated) :param filename: name of the file (templated) - :param filetype: slack filetype. (templated) - - see https://api.slack.com/types/file + :param filetype: slack filetype. (templated) See: https://api.slack.com/types/file#file_types :param content: file content. (templated) + :param title: title of file. (templated) + :param channel: (deprecated) channel in which to sent file on slack name """ - template_fields: Sequence[str] = ('channel', 'initial_comment', 'filename', 'filetype', 'content') + template_fields: Sequence[str] = ( + 'channels', + 'initial_comment', + 'filename', + 'filetype', + 'content', + 'title', + ) ui_color = '#44BEDF' def __init__( self, - channel: str = '#general', - initial_comment: str = 'No message has been set!', + channels: Optional[Union[str, Sequence[str]]] = None, + initial_comment: Optional[str] = None, filename: Optional[str] = None, filetype: Optional[str] = None, content: Optional[str] = None, + title: Optional[str] = None, + channel: Optional[str] = None, **kwargs, ) -> None: - self.method = 'files.upload' - self.channel = channel + if channel: + warnings.warn( + "Argument `channel` is deprecated and will removed in a future releases. " + "Please use `channels` instead.", + DeprecationWarning, + stacklevel=2, + ) + if channels: + raise ValueError(f"Cannot set both arguments: channel={channel!r} and channels={channels!r}.") + channels = channel + + self.channels = channels self.initial_comment = initial_comment self.filename = filename self.filetype = filetype self.content = content - self.file_params: Dict = {} - super().__init__(method=self.method, **kwargs) + self.title = title + super().__init__(method="files.upload", **kwargs) 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 - """ - slack = SlackHook(token=self.token, slack_conn_id=self.slack_conn_id) - - # If file content is passed. - if self.content is not None: - self.api_params = { - 'channels': self.channel, - 'content': self.content, - 'initial_comment': self.initial_comment, - } - slack.call(self.method, data=self.api_params) - # If file name is passed. - elif self.filename is not None: - self.api_params = { - 'channels': self.channel, - 'filename': self.filename, - 'filetype': self.filetype, - 'initial_comment': self.initial_comment, - } - with open(self.filename, "rb") as file_handle: - slack.call(self.method, data=self.api_params, files={'file': file_handle}) - file_handle.close() + self.hook.send_file( + channels=self.channels, + # For historical reason SlackAPIFileOperator use filename as reference to file + file=self.filename, + content=self.content, + initial_comment=self.initial_comment, + title=self.title, + ) diff --git a/tests/providers/slack/hooks/test_slack.py b/tests/providers/slack/hooks/test_slack.py index 9981ed331adf0..db7e02afd8f5a 100644 --- a/tests/providers/slack/hooks/test_slack.py +++ b/tests/providers/slack/hooks/test_slack.py @@ -366,3 +366,99 @@ def test_hook_connection_failed(self, mock_webclient_cls, response_data): conn_test = hook.test_connection() mock_webclient_call.assert_called_once_with("auth.test") assert not conn_test[0] + + @pytest.mark.parametrize("file,content", [(None, None), ("", ""), ("foo.bar", "test-content")]) + def test_send_file_wrong_parameters(self, file, content): + hook = SlackHook(slack_conn_id=SLACK_API_DEFAULT_CONN_ID) + error_message = r"Either `file` or `content` must be provided, not both\." + with pytest.raises(ValueError, match=error_message): + hook.send_file(file=file, content=content) + + @mock.patch('airflow.providers.slack.hooks.slack.WebClient') + @pytest.mark.parametrize("initial_comment", [None, "test comment"]) + @pytest.mark.parametrize("title", [None, "test title"]) + @pytest.mark.parametrize("filetype", [None, "auto"]) + @pytest.mark.parametrize("channels", [None, "#random", "#random,#general", ("#random", "#general")]) + def test_send_file_path( + self, mock_webclient_cls, tmp_path_factory, initial_comment, title, filetype, channels + ): + """Test send file by providing filepath.""" + mock_files_upload = mock.MagicMock() + mock_webclient_cls.return_value.files_upload = mock_files_upload + + tmp = tmp_path_factory.mktemp("test_send_file_path") + file = tmp / "test.json" + file.write_bytes(b'{"foo": "bar"}') + + hook = SlackHook(slack_conn_id=SLACK_API_DEFAULT_CONN_ID) + hook.send_file( + channels=channels, + file=file, + filename="filename.mock", + initial_comment=initial_comment, + title=title, + filetype=filetype, + ) + + mock_files_upload.assert_called_once_with( + channels=channels, + file=mock.ANY, # Validate file properties later + filename="filename.mock", + initial_comment=initial_comment, + title=title, + filetype=filetype, + ) + + # Validate file properties + mock_file = mock_files_upload.call_args[1]["file"] + assert mock_file.mode == "rb" + assert mock_file.name == str(file) + + @mock.patch('airflow.providers.slack.hooks.slack.WebClient') + @pytest.mark.parametrize("filename", ["test.json", "1.parquet.snappy"]) + def test_send_file_path_set_filename(self, mock_webclient_cls, tmp_path_factory, filename): + """Test set filename in send_file method if it not set.""" + mock_files_upload = mock.MagicMock() + mock_webclient_cls.return_value.files_upload = mock_files_upload + + tmp = tmp_path_factory.mktemp("test_send_file_path_set_filename") + file = tmp / filename + file.write_bytes(b'{"foo": "bar"}') + + hook = SlackHook(slack_conn_id=SLACK_API_DEFAULT_CONN_ID) + hook.send_file(file=file) + + assert mock_files_upload.call_count == 1 + call_args = mock_files_upload.call_args[1] + assert "filename" in call_args + assert call_args["filename"] == filename + + @mock.patch('airflow.providers.slack.hooks.slack.WebClient') + @pytest.mark.parametrize("initial_comment", [None, "test comment"]) + @pytest.mark.parametrize("title", [None, "test title"]) + @pytest.mark.parametrize("filetype", [None, "auto"]) + @pytest.mark.parametrize("filename", [None, "foo.bar"]) + @pytest.mark.parametrize("channels", [None, "#random", "#random,#general", ("#random", "#general")]) + def test_send_file_content( + self, mock_webclient_cls, initial_comment, title, filetype, channels, filename + ): + """Test send file by providing content.""" + mock_files_upload = mock.MagicMock() + mock_webclient_cls.return_value.files_upload = mock_files_upload + hook = SlackHook(slack_conn_id=SLACK_API_DEFAULT_CONN_ID) + hook.send_file( + channels=channels, + content='{"foo": "bar"}', + filename=filename, + initial_comment=initial_comment, + title=title, + filetype=filetype, + ) + mock_files_upload.assert_called_once_with( + channels=channels, + content='{"foo": "bar"}', + filename=filename, + initial_comment=initial_comment, + title=title, + filetype=filetype, + ) diff --git a/tests/providers/slack/operators/test.csv b/tests/providers/slack/operators/test.csv deleted file mode 100644 index 9daeafb9864cf..0000000000000 --- a/tests/providers/slack/operators/test.csv +++ /dev/null @@ -1 +0,0 @@ -test diff --git a/tests/providers/slack/operators/test_slack.py b/tests/providers/slack/operators/test_slack.py index 495dc7f6bae64..d184093dba392 100644 --- a/tests/providers/slack/operators/test_slack.py +++ b/tests/providers/slack/operators/test_slack.py @@ -16,15 +16,64 @@ # specific language governing permissions and limitations # under the License. import json -import unittest from unittest import mock from unittest.mock import MagicMock -from airflow.providers.slack.operators.slack import SlackAPIFileOperator, SlackAPIPostOperator +import pytest +from airflow.models import Connection +from airflow.providers.slack.operators.slack import ( + SlackAPIFileOperator, + SlackAPIOperator, + SlackAPIPostOperator, +) -class TestSlackAPIPostOperator(unittest.TestCase): - def setUp(self): +SLACK_API_TEST_CONNECTION_ID = "test_slack_conn_id" + + +@pytest.fixture(scope="module", autouse=True) +def slack_api_connections(): + """Create tests connections.""" + connections = [ + Connection( + conn_id=SLACK_API_TEST_CONNECTION_ID, + conn_type="slack", + password="xoxb-1234567890123-09876543210987-AbCdEfGhIjKlMnOpQrStUvWx", + ), + ] + conn_uris = {f"AIRFLOW_CONN_{c.conn_id.upper()}": c.get_uri() for c in connections} + + with mock.patch.dict("os.environ", values=conn_uris): + yield + + +class TestSlackAPIOperator: + @mock.patch("airflow.providers.slack.operators.slack.mask_secret") + def test_mask_token(self, mock_mask_secret): + SlackAPIOperator(task_id="test-mask-token", token="super-secret-token") + mock_mask_secret.assert_called_once_with("super-secret-token") + + @mock.patch("airflow.providers.slack.operators.slack.SlackHook") + @pytest.mark.parametrize( + "token,conn_id", + [ + ("token", SLACK_API_TEST_CONNECTION_ID), + ("token", None), + (None, SLACK_API_TEST_CONNECTION_ID), + ], + ) + def test_hook(self, mock_slack_hook_cls, token, conn_id): + mock_slack_hook = mock_slack_hook_cls.return_value + op = SlackAPIOperator(task_id="test-mask-token", token=token, slack_conn_id=conn_id) + hook = op.hook + assert hook == mock_slack_hook + assert hook is op.hook + mock_slack_hook_cls.assert_called_once_with(token=token, slack_conn_id=conn_id) + + +class TestSlackAPIPostOperator: + @pytest.fixture(autouse=True) + def setup(self): self.test_username = 'test_username' self.test_channel = '#test_slack_channel' self.test_text = 'test_text' @@ -91,7 +140,6 @@ def __construct_operator(self, test_token, test_slack_conn_id, test_api_params=N def test_init_with_valid_params(self): test_token = 'test_token' - test_slack_conn_id = 'test_slack_conn_id' slack_api_post_operator = self.__construct_operator(test_token, None, self.test_api_params) assert slack_api_post_operator.token == test_token @@ -105,18 +153,16 @@ def test_init_with_valid_params(self): assert slack_api_post_operator.attachments == self.test_attachments assert slack_api_post_operator.blocks == self.test_blocks - slack_api_post_operator = self.__construct_operator(None, test_slack_conn_id) + slack_api_post_operator = self.__construct_operator(None, SLACK_API_TEST_CONNECTION_ID) assert slack_api_post_operator.token is None - assert slack_api_post_operator.slack_conn_id == test_slack_conn_id + assert slack_api_post_operator.slack_conn_id == SLACK_API_TEST_CONNECTION_ID @mock.patch('airflow.providers.slack.operators.slack.SlackHook') def test_api_call_params_with_default_args(self, mock_hook): - test_slack_conn_id = 'test_slack_conn_id' - slack_api_post_operator = SlackAPIPostOperator( task_id='slack', username=self.test_username, - slack_conn_id=test_slack_conn_id, + slack_conn_id=SLACK_API_TEST_CONNECTION_ID, ) slack_api_post_operator.execute(context=MagicMock()) @@ -135,32 +181,24 @@ def test_api_call_params_with_default_args(self, mock_hook): assert expected_api_params == slack_api_post_operator.api_params -class TestSlackAPIFileOperator(unittest.TestCase): - def setUp(self): +class TestSlackAPIFileOperator: + @pytest.fixture(autouse=True) + def setup(self): self.test_username = 'test_username' self.test_channel = '#test_slack_channel' self.test_initial_comment = 'test text file test_filename.txt' self.filename = 'test_filename.txt' self.test_filetype = 'text' self.test_content = 'This is a test text file!' - self.test_api_params = {'key': 'value'} - self.expected_method = 'files.upload' - self.expected_api_params = { - 'channel': self.test_channel, - 'initial_comment': self.test_initial_comment, - 'file': self.filename, - 'filetype': self.test_filetype, - 'content': self.test_content, - } def __construct_operator(self, test_token, test_slack_conn_id, test_api_params=None): return SlackAPIFileOperator( task_id='slack', token=test_token, slack_conn_id=test_slack_conn_id, - channel=self.test_channel, + channels=self.test_channel, initial_comment=self.test_initial_comment, filename=self.filename, filetype=self.test_filetype, @@ -170,63 +208,84 @@ def __construct_operator(self, test_token, test_slack_conn_id, test_api_params=N def test_init_with_valid_params(self): test_token = 'test_token' - test_slack_conn_id = 'test_slack_conn_id' slack_api_post_operator = self.__construct_operator(test_token, None, self.test_api_params) assert slack_api_post_operator.token == test_token assert slack_api_post_operator.slack_conn_id is None assert slack_api_post_operator.method == self.expected_method assert slack_api_post_operator.initial_comment == self.test_initial_comment - assert slack_api_post_operator.channel == self.test_channel + assert slack_api_post_operator.channels == self.test_channel assert slack_api_post_operator.api_params == self.test_api_params assert slack_api_post_operator.filename == self.filename assert slack_api_post_operator.filetype == self.test_filetype assert slack_api_post_operator.content == self.test_content - slack_api_post_operator = self.__construct_operator(None, test_slack_conn_id) + slack_api_post_operator = self.__construct_operator(None, SLACK_API_TEST_CONNECTION_ID) assert slack_api_post_operator.token is None - assert slack_api_post_operator.slack_conn_id == test_slack_conn_id + assert slack_api_post_operator.slack_conn_id == SLACK_API_TEST_CONNECTION_ID - @mock.patch('airflow.providers.slack.operators.slack.SlackHook') - def test_api_call_params_with_content_args(self, mock_hook): - test_slack_conn_id = 'test_slack_conn_id' + @mock.patch('airflow.providers.slack.operators.slack.SlackHook.send_file') + @pytest.mark.parametrize("initial_comment", [None, "foo-bar"]) + @pytest.mark.parametrize("title", [None, "Spam Egg"]) + def test_api_call_params_with_content_args(self, mock_send_file, initial_comment, title): + SlackAPIFileOperator( + task_id='slack', + slack_conn_id=SLACK_API_TEST_CONNECTION_ID, + content='test-content', + channels='#test-channel', + initial_comment=initial_comment, + title=title, + ).execute(context=MagicMock()) - slack_api_post_operator = SlackAPIFileOperator( - task_id='slack', slack_conn_id=test_slack_conn_id, content='test-content' + mock_send_file.assert_called_once_with( + channels='#test-channel', + content='test-content', + file=None, + initial_comment=initial_comment, + title=title, ) - slack_api_post_operator.execute(context=MagicMock()) - - 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' - - import os - - # Look for your absolute directory path - absolute_path = os.path.dirname(os.path.abspath(__file__)) - # Or: file_path = os.path.join(absolute_path, 'folder', 'my_file.py') - file_path = absolute_path + '/test.csv' - - print(f"full path ${file_path}") + @mock.patch('airflow.providers.slack.operators.slack.SlackHook.send_file') + @pytest.mark.parametrize("initial_comment", [None, "foo-bar"]) + @pytest.mark.parametrize("title", [None, "Spam Egg"]) + def test_api_call_params_with_file_args(self, mock_send_file, initial_comment, title): + SlackAPIFileOperator( + task_id='slack', + slack_conn_id=SLACK_API_TEST_CONNECTION_ID, + channels='C1234567890', + filename='/dev/null', + initial_comment=initial_comment, + title=title, + ).execute(context=MagicMock()) - slack_api_post_operator = SlackAPIFileOperator( - task_id='slack', slack_conn_id=test_slack_conn_id, filename=file_path, filetype='csv' + mock_send_file.assert_called_once_with( + channels='C1234567890', + content=None, + file='/dev/null', + initial_comment=initial_comment, + title=title, ) - slack_api_post_operator.execute(context=MagicMock()) + def test_channel_deprecated(self): + warning_message = ( + r"Argument `channel` is deprecated and will removed in a future releases\. " + r"Please use `channels` instead\." + ) + with pytest.warns(DeprecationWarning, match=warning_message): + op = SlackAPIFileOperator( + task_id='slack', + slack_conn_id=SLACK_API_TEST_CONNECTION_ID, + channel="#random", + channels=None, + ) + assert op.channels == "#random" - expected_api_params = { - 'channels': '#general', - 'initial_comment': 'No message has been set!', - 'filename': file_path, - 'filetype': 'csv', - } - assert expected_api_params == slack_api_post_operator.api_params + def test_both_channel_and_channels_set(self): + error_message = r"Cannot set both arguments: channel=.* and channels=.*\." + with pytest.raises(ValueError, match=error_message): + SlackAPIFileOperator( + task_id='slack', + slack_conn_id=SLACK_API_TEST_CONNECTION_ID, + channel="#random", + channels="#general", + )