Add influxdb operator (#19356)

airflow/providers/influxdb/example_dags/example_influxdb_query.py

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from datetime import datetime
+
+from airflow.models.dag import DAG
+from airflow.providers.influxdb.operators.influxdb import InfluxDBOperator
+
+dag = DAG(
+    'example_influxdb_operator',
+    start_date=datetime(2021, 1, 1),
+    tags=['example'],
+    catchup=False,
+)
+
+# [START howto_operator_influxdb]
+
+query_influxdb_task = InfluxDBOperator(
+    influxdb_conn_id='influxdb_conn_id',
+    task_id='query_influxdb',
+    sql='from(bucket:"test-influx") |> range(start: -10m, stop: {{ds}})',
+    dag=dag,
+)
+
+# [END howto_operator_influxdb]

airflow/providers/influxdb/operators/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/influxdb/operators/influxdb.py

@@ -0,0 +1,55 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from typing import Dict
+
+from airflow.models import BaseOperator
+from airflow.providers.influxdb.hooks.influxdb import InfluxDBHook
+
+
+class InfluxDBOperator(BaseOperator):
+    """
+    Executes sql code in a specific InfluxDB database
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:InfluxDBOperator`
+
+    :param sql: the sql code to be executed. Can receive a str representing a
+        sql statement
+    :type sql: str
+    :param influxdb_conn_id: Reference to :ref:`Influxdb connection id <howto/connection:influxdb>`.
+    :type influxdb_conn_id: str
+    """
+
+    template_fields = ['sql']
+
+    def __init__(
+        self,
+        *,
+        sql: str,
+        influxdb_conn_id: str = 'influxdb_default',
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.influxdb_conn_id = influxdb_conn_id
+        self.sql = sql
+
+    def execute(self, context: Dict) -> None:
+        self.log.info('Executing: %s', self.sql)
+        self.hook = InfluxDBHook(conn_id=self.influxdb_conn_id)
+        self.hook.query(self.sql)

airflow/providers/influxdb/provider.yaml

@@ -32,6 +32,11 @@ hooks:
     python-modules:
       - airflow.providers.influxdb.hooks.influxdb
 
+operators:
+  - integration-name: Influxdb
+    python-modules:
+      - airflow.providers.influxdb.operators.influxdb
+
 connection-types:
   - hook-class-name: airflow.providers.influxdb.hooks.influxdb.InfluxDBHook
     connection-type: influxdb

docs/apache-airflow-providers-influxdb/index.rst

@@ -27,6 +27,7 @@ Content
     :caption: Guides
 
     Connection types <connections/influxdb>
+    Operators <operators/index>
 
 .. toctree::
     :maxdepth: 1

docs/apache-airflow-providers-influxdb/operators/index.rst

@@ -0,0 +1,33 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+
+
+.. _howto/operator:InfluxDBOperator:
+
+InfluxDBOperator
+=================
+
+Use the :class:`~airflow.providers.influxdb.operators.InfluxDBOperator` to execute
+SQL commands in a `InfluxDB <https://www.influxdata.com/>`__ database.
+
+An example of running the query using the operator:
+
+.. exampleinclude:: /../../airflow/providers/influxdb/example_dags/example_influxdb_query.py
+    :language: python
+    :start-after: [START howto_operator_influxdb]
+    :end-before: [END howto_operator_influxdb]

tests/providers/influxdb/operators/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

tests/providers/influxdb/operators/test_influxdb.py

@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import unittest
+from unittest import mock
+
+from airflow.providers.influxdb.operators.influxdb import InfluxDBOperator
+
+
+class TestInfluxDBOperator(unittest.TestCase):
+    @mock.patch('airflow.providers.influxdb.operators.influxdb.InfluxDBHook')
+    def test_influxdb_operator_test(self, mock_hook):
+
+        sql = """from(bucket:"test") |> range(start: -10m)"""
+        op = InfluxDBOperator(task_id='basic_influxdb', sql=sql, influxdb_conn_id='influxdb_default')
+        op.execute(mock.MagicMock())
+        mock_hook.assert_called_once_with(conn_id='influxdb_default')
+        mock_hook.return_value.query.assert_called_once_with(sql)