Implements SqlToSlackApiFileOperator

airflow/providers/slack/transfers/sql_to_slack.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Iterable, Mapping, Sequence
 
 from pandas import DataFrame
@@ -25,13 +26,59 @@
 from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator
 from airflow.providers.common.sql.hooks.sql import DbApiHook
+from airflow.providers.slack.hooks.slack import SlackHook
 from airflow.providers.slack.hooks.slack_webhook import SlackWebhookHook
+from airflow.providers.slack.utils import parse_filename
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class SqlToSlackOperator(BaseOperator):
+class BaseSqlToSlackOperator(BaseOperator):
+    """
+    Operator implements base sql methods for SQL to Slack Transfer operators.
+
+    :param sql: The SQL query to be executed
+    :param sql_conn_id: reference to a specific DB-API Connection.
+    :param sql_hook_params: Extra config params to be passed to the underlying hook.
+        Should match the desired hook constructor params.
+    :param parameters: The parameters to pass to the SQL query.
+    """
+
+    def __init__(
+        self,
+        *,
+        sql: str,
+        sql_conn_id: str,
+        sql_hook_params: dict | None = None,
+        parameters: Iterable | Mapping | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.sql_conn_id = sql_conn_id
+        self.sql_hook_params = sql_hook_params
+        self.sql = sql
+        self.parameters = parameters
+
+    def _get_hook(self) -> DbApiHook:
+        self.log.debug("Get connection for %s", self.sql_conn_id)
+        conn = BaseHook.get_connection(self.sql_conn_id)
+        hook = conn.get_hook(hook_params=self.sql_hook_params)
+        if not callable(getattr(hook, "get_pandas_df", None)):
+            raise AirflowException(
+                "This hook is not supported. The hook class must have get_pandas_df method."
+            )
+        return hook
+
+    def _get_query_results(self) -> DataFrame:
+        sql_hook = self._get_hook()
+
+        self.log.info("Running SQL query: %s", self.sql)
+        df = sql_hook.get_pandas_df(self.sql, parameters=self.parameters)
+        return df
+
+
+class SqlToSlackOperator(BaseSqlToSlackOperator):
     """
     Executes an SQL statement in a given SQL connection and sends the results to Slack. The results of the
     query are rendered into the 'slack_message' parameter as a Pandas dataframe using a JINJA variable called
@@ -79,12 +126,10 @@ def __init__(
         **kwargs,
     ) -> None:
 
-        super().__init__(**kwargs)
+        super().__init__(
+            sql=sql, sql_conn_id=sql_conn_id, sql_hook_params=sql_hook_params, parameters=parameters, **kwargs
+        )
 
-        self.sql_conn_id = sql_conn_id
-        self.sql_hook_params = sql_hook_params
-        self.sql = sql
-        self.parameters = parameters
         self.slack_conn_id = slack_conn_id
         self.slack_webhook_token = slack_webhook_token
         self.slack_channel = slack_channel
@@ -97,23 +142,6 @@ def __init__(
                 "SqlToSlackOperator requires either a `slack_conn_id` or a `slack_webhook_token` argument"
             )
 
-    def _get_hook(self) -> DbApiHook:
-        self.log.debug("Get connection for %s", self.sql_conn_id)
-        conn = BaseHook.get_connection(self.sql_conn_id)
-        hook = conn.get_hook(hook_params=self.sql_hook_params)
-        if not callable(getattr(hook, "get_pandas_df", None)):
-            raise AirflowException(
-                "This hook is not supported. The hook class must have get_pandas_df method."
-            )
-        return hook
-
-    def _get_query_results(self) -> DataFrame:
-        sql_hook = self._get_hook()
-
-        self.log.info("Running SQL query: %s", self.sql)
-        df = sql_hook.get_pandas_df(self.sql, parameters=self.parameters)
-        return df
-
     def _render_and_send_slack_message(self, context, df) -> None:
         # Put the dataframe into the context and render the JINJA template fields
         context[self.results_df_name] = df
@@ -157,3 +185,115 @@ def execute(self, context: Context) -> None:
         self._render_and_send_slack_message(context, df)
 
         self.log.debug("Finished sending SQL data to Slack")
+
+
+class SqlToSlackApiFileOperator(BaseSqlToSlackOperator):
+    """
+    Executes an SQL statement in a given SQL connection and sends the results to Slack API as file.
+
+    :param sql: The SQL query to be executed
+    :param sql_conn_id: reference to a specific DB-API Connection.
+    :param slack_conn_id: :ref:`Slack API Connection <howto/connection:slack>`.
+    :param slack_filename: Filename for display in slack.
+        Should contain supported extension which referenced to ``SUPPORTED_FILE_FORMATS``.
+        It is also possible to set compression in extension:
+        ``filename.csv.gzip``, ``filename.json.zip``, etc.
+    :param sql_hook_params: Extra config params to be passed to the underlying hook.
+        Should match the desired hook constructor params.
+    :param parameters: The parameters to pass to the SQL query.
+    :param slack_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 slack_initial_comment: The message text introducing the file in specified ``slack_channels``.
+    :param slack_title: Title of file.
+    :param df_kwargs: Keyword arguments forwarded to ``pandas.DataFrame.to_{format}()`` method.
+
+    Example:
+     .. code-block:: python
+
+        SqlToSlackApiFileOperator(
+            task_id="sql_to_slack",
+            sql="SELECT 1 a, 2 b, 3 c",
+            sql_conn_id="sql-connection",
+            slack_conn_id="slack-api-connection",
+            slack_filename="awesome.json.gz",
+            slack_channels="#random,#general",
+            slack_initial_comment="Awesome load to compressed multiline JSON.",
+            df_kwargs={
+                "orient": "records",
+                "lines": True,
+            },
+        )
+    """
+
+    template_fields: Sequence[str] = (
+        "sql",
+        "slack_channels",
+        "slack_filename",
+        "slack_initial_comment",
+        "slack_title",
+    )
+    template_ext: Sequence[str] = (".sql", ".jinja", ".j2")
+    template_fields_renderers = {"sql": "sql", "slack_message": "jinja"}
+
+    SUPPORTED_FILE_FORMATS: Sequence[str] = ("csv", "json", "html")
+
+    def __init__(
+        self,
+        *,
+        sql: str,
+        sql_conn_id: str,
+        sql_hook_params: dict | None = None,
+        parameters: Iterable | Mapping | None = None,
+        slack_conn_id: str,
+        slack_filename: str,
+        slack_channels: str | Sequence[str] | None = None,
+        slack_initial_comment: str | None = None,
+        slack_title: str | None = None,
+        df_kwargs: dict | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            sql=sql, sql_conn_id=sql_conn_id, sql_hook_params=sql_hook_params, parameters=parameters, **kwargs
+        )
+        self.slack_conn_id = slack_conn_id
+        self.slack_filename = slack_filename
+        self.slack_channels = slack_channels
+        self.slack_initial_comment = slack_initial_comment
+        self.slack_title = slack_title
+        self.df_kwargs = df_kwargs or {}
+
+    def execute(self, context: Context) -> None:
+        # Parse file format from filename
+        output_file_format, _ = parse_filename(
+            filename=self.slack_filename,
+            supported_file_formats=self.SUPPORTED_FILE_FORMATS,
+        )
+
+        slack_hook = SlackHook(slack_conn_id=self.slack_conn_id)
+        with NamedTemporaryFile(mode="w+", suffix=f"_{self.slack_filename}") as fp:
+            # tempfile.NamedTemporaryFile used only for create and remove temporary file,
+            # pandas will open file in correct mode itself depend on file type.
+            # So we close file descriptor here for avoid incidentally write anything.
+            fp.close()
+
+            output_file_name = fp.name
+            output_file_format = output_file_format.upper()
+            df_result = self._get_query_results()
+            if output_file_format == "CSV":
+                df_result.to_csv(output_file_name, **self.df_kwargs)
+            elif output_file_format == "JSON":
+                df_result.to_json(output_file_name, **self.df_kwargs)
+            elif output_file_format == "HTML":
+                df_result.to_html(output_file_name, **self.df_kwargs)
+            else:
+                # Not expected that this error happen. This only possible
+                # if SUPPORTED_FILE_FORMATS extended and no actual implementation for specific format.
+                raise AirflowException(f"Unexpected output file format: {output_file_format}")
+
+            slack_hook.send_file(
+                channels=self.slack_channels,
+                file=output_file_name,
+                filename=self.slack_filename,
+                initial_comment=self.slack_initial_comment,
+                title=self.slack_title,
+            )

airflow/providers/slack/utils/__init__.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any
+from typing import Any, Sequence
 
 from airflow.utils.types import NOTSET
 
@@ -77,3 +77,41 @@ def getint(self, field, default: Any = NOTSET) -> Any:
         if value != default:
             value = int(value)
         return value
+
+
+def parse_filename(
+    filename: str, supported_file_formats: Sequence[str], fallback: str | None = None
+) -> tuple[str, str | None]:
+    """
+    Parse filetype and compression from given filename.
+    :param filename: filename to parse.
+    :param supported_file_formats: list of supported file extensions.
+    :param fallback: fallback to given file format.
+    :returns: filetype and compression (if specified)
+    """
+    if not filename:
+        raise ValueError("Expected 'filename' parameter is missing.")
+    if fallback and fallback not in supported_file_formats:
+        raise ValueError(f"Invalid fallback value {fallback!r}, expected one of {supported_file_formats}.")
+
+    parts = filename.rsplit(".", 2)
+    try:
+        if len(parts) == 1:
+            raise ValueError(f"No file extension specified in filename {filename!r}.")
+        if parts[-1] in supported_file_formats:
+            return parts[-1], None
+        elif len(parts) == 2:
+            raise ValueError(
+                f"Unsupported file format {parts[-1]!r}, expected one of {supported_file_formats}."
+            )
+        else:
+            if parts[-2] not in supported_file_formats:
+                raise ValueError(
+                    f"Unsupported file format '{parts[-2]}.{parts[-1]}', "
+                    f"expected one of {supported_file_formats} with compression extension."
+                )
+            return parts[-2], parts[-1]
+    except ValueError as ex:
+        if fallback:
+            return fallback, None
+        raise ex from None