You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ASF GitHub Bot (JIRA)" <ji...@apache.org> on 2018/08/07 17:22:04 UTC
[jira] [Commented] (AIRFLOW-2732) Split hooks and operators out
from core Airflow
[ https://issues.apache.org/jira/browse/AIRFLOW-2732?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel&focusedCommentId=16572001#comment-16572001 ]
ASF GitHub Bot commented on AIRFLOW-2732:
-----------------------------------------
tedmiston closed pull request #3587: [WIP][AIRFLOW-2732] Remove hooks and operators from core
URL: https://github.com/apache/incubator-airflow/pull/3587
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/airflow/contrib/example_dags/example_databricks_operator.py b/airflow/contrib/example_dags/example_databricks_operator.py
deleted file mode 100644
index bc827d465b..0000000000
--- a/airflow/contrib/example_dags/example_databricks_operator.py
+++ /dev/null
@@ -1,87 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow
-
-from airflow import DAG
-from airflow.contrib.operators.databricks_operator import DatabricksSubmitRunOperator
-
-# This is an example DAG which uses the DatabricksSubmitRunOperator.
-# In this example, we create two tasks which execute sequentially.
-# The first task is to run a notebook at the workspace path "/test"
-# and the second task is to run a JAR uploaded to DBFS. Both,
-# tasks use new clusters.
-#
-# Because we have set a downstream dependency on the notebook task,
-# the spark jar task will NOT run until the notebook task completes
-# successfully.
-#
-# The definition of a succesful run is if the run has a result_state of "SUCCESS".
-# For more information about the state of a run refer to
-# https://docs.databricks.com/api/latest/jobs.html#runstate
-
-args = {
- 'owner': 'airflow',
- 'email': ['airflow@example.com'],
- 'depends_on_past': False,
- 'start_date': airflow.utils.dates.days_ago(2)
-}
-
-dag = DAG(
- dag_id='example_databricks_operator', default_args=args,
- schedule_interval='@daily')
-
-new_cluster = {
- 'spark_version': '2.1.0-db3-scala2.11',
- 'node_type_id': 'r3.xlarge',
- 'aws_attributes': {
- 'availability': 'ON_DEMAND'
- },
- 'num_workers': 8
-}
-
-notebook_task_params = {
- 'new_cluster': new_cluster,
- 'notebook_task': {
- 'notebook_path': '/Users/airflow@example.com/PrepareData',
- },
-}
-# Example of using the JSON parameter to initialize the operator.
-notebook_task = DatabricksSubmitRunOperator(
- task_id='notebook_task',
- dag=dag,
- json=notebook_task_params)
-
-# Example of using the named parameters of DatabricksSubmitRunOperator
-# to initialize the operator.
-spark_jar_task = DatabricksSubmitRunOperator(
- task_id='spark_jar_task',
- dag=dag,
- new_cluster=new_cluster,
- spark_jar_task={
- 'main_class_name': 'com.example.ProcessData'
- },
- libraries=[
- {
- 'jar': 'dbfs:/lib/etl-0.1.jar'
- }
- ]
-)
-
-notebook_task.set_downstream(spark_jar_task)
diff --git a/airflow/contrib/example_dags/example_emr_job_flow_automatic_steps.py b/airflow/contrib/example_dags/example_emr_job_flow_automatic_steps.py
deleted file mode 100644
index 098a7a04ea..0000000000
--- a/airflow/contrib/example_dags/example_emr_job_flow_automatic_steps.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 timedelta
-import airflow
-from airflow import DAG
-from airflow.contrib.operators.emr_create_job_flow_operator \
- import EmrCreateJobFlowOperator
-from airflow.contrib.sensors.emr_job_flow_sensor import EmrJobFlowSensor
-
-DEFAULT_ARGS = {
- 'owner': 'airflow',
- 'depends_on_past': False,
- 'start_date': airflow.utils.dates.days_ago(2),
- 'email': ['airflow@example.com'],
- 'email_on_failure': False,
- 'email_on_retry': False
-}
-
-SPARK_TEST_STEPS = [
- {
- 'Name': 'calculate_pi',
- 'ActionOnFailure': 'CONTINUE',
- 'HadoopJarStep': {
- 'Jar': 'command-runner.jar',
- 'Args': [
- '/usr/lib/spark/bin/run-example',
- 'SparkPi',
- '10'
- ]
- }
- }
-]
-
-JOB_FLOW_OVERRIDES = {
- 'Name': 'PiCalc',
- 'Steps': SPARK_TEST_STEPS
-}
-
-dag = DAG(
- 'emr_job_flow_automatic_steps_dag',
- default_args=DEFAULT_ARGS,
- dagrun_timeout=timedelta(hours=2),
- schedule_interval='0 3 * * *'
-)
-
-job_flow_creator = EmrCreateJobFlowOperator(
- task_id='create_job_flow',
- job_flow_overrides=JOB_FLOW_OVERRIDES,
- aws_conn_id='aws_default',
- emr_conn_id='emr_default',
- dag=dag
-)
-
-job_sensor = EmrJobFlowSensor(
- task_id='check_job_flow',
- job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
- aws_conn_id='aws_default',
- dag=dag
-)
-
-job_flow_creator.set_downstream(job_sensor)
diff --git a/airflow/contrib/example_dags/example_emr_job_flow_manual_steps.py b/airflow/contrib/example_dags/example_emr_job_flow_manual_steps.py
deleted file mode 100644
index 48a178a2a9..0000000000
--- a/airflow/contrib/example_dags/example_emr_job_flow_manual_steps.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 timedelta
-
-import airflow
-from airflow import DAG
-from airflow.contrib.operators.emr_create_job_flow_operator \
- import EmrCreateJobFlowOperator
-from airflow.contrib.operators.emr_add_steps_operator \
- import EmrAddStepsOperator
-from airflow.contrib.sensors.emr_step_sensor import EmrStepSensor
-from airflow.contrib.operators.emr_terminate_job_flow_operator \
- import EmrTerminateJobFlowOperator
-
-DEFAULT_ARGS = {
- 'owner': 'airflow',
- 'depends_on_past': False,
- 'start_date': airflow.utils.dates.days_ago(2),
- 'email': ['airflow@example.com'],
- 'email_on_failure': False,
- 'email_on_retry': False
-}
-
-SPARK_TEST_STEPS = [
- {
- 'Name': 'calculate_pi',
- 'ActionOnFailure': 'CONTINUE',
- 'HadoopJarStep': {
- 'Jar': 'command-runner.jar',
- 'Args': [
- '/usr/lib/spark/bin/run-example',
- 'SparkPi',
- '10'
- ]
- }
- }
-]
-
-JOB_FLOW_OVERRIDES = {
- 'Name': 'PiCalc',
- 'KeepJobFlowAliveWhenNoSteps': True
-}
-
-dag = DAG(
- 'emr_job_flow_manual_steps_dag',
- default_args=DEFAULT_ARGS,
- dagrun_timeout=timedelta(hours=2),
- schedule_interval='0 3 * * *'
-)
-
-cluster_creator = EmrCreateJobFlowOperator(
- task_id='create_job_flow',
- job_flow_overrides=JOB_FLOW_OVERRIDES,
- aws_conn_id='aws_default',
- emr_conn_id='emr_default',
- dag=dag
-)
-
-step_adder = EmrAddStepsOperator(
- task_id='add_steps',
- job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
- aws_conn_id='aws_default',
- steps=SPARK_TEST_STEPS,
- dag=dag
-)
-
-step_checker = EmrStepSensor(
- task_id='watch_step',
- job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
- step_id="{{ task_instance.xcom_pull('add_steps', key='return_value')[0] }}",
- aws_conn_id='aws_default',
- dag=dag
-)
-
-cluster_remover = EmrTerminateJobFlowOperator(
- task_id='remove_cluster',
- job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
- aws_conn_id='aws_default',
- dag=dag
-)
-
-cluster_creator.set_downstream(step_adder)
-step_adder.set_downstream(step_checker)
-step_checker.set_downstream(cluster_remover)
diff --git a/airflow/contrib/example_dags/example_jenkins_job_trigger_operator.py.notexecutable b/airflow/contrib/example_dags/example_jenkins_job_trigger_operator.py.notexecutable
deleted file mode 100644
index 2d8906b0b3..0000000000
--- a/airflow/contrib/example_dags/example_jenkins_job_trigger_operator.py.notexecutable
+++ /dev/null
@@ -1,82 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow import DAG
-from airflow.contrib.operators.jenkins_job_trigger_operator import JenkinsJobTriggerOperator
-from airflow.operators.python_operator import PythonOperator
-from airflow.contrib.hooks.jenkins_hook import JenkinsHook
-
-from six.moves.urllib.request import Request
-
-import jenkins
-import datetime
-
-
-datetime_start_date = datetime(2017, 6, 1)
-default_args = {
- "owner": "airflow",
- "start_date": datetime_start_date,
- "retries": 1,
- "retry_delay": timedelta(minutes=5),
- "depends_on_past": False,
- "concurrency": 8,
- "max_active_runs": 8
-
-}
-
-
-dag = DAG("test_jenkins", default_args=default_args, schedule_interval=None)
-#This DAG shouldn't be executed and is only here to provide example of how to use the JenkinsJobTriggerOperator
-#(it requires a jenkins server to be executed)
-
-job_trigger = JenkinsJobTriggerOperator(
- dag=dag,
- task_id="trigger_job",
- job_name="generate-merlin-config",
- parameters={"first_parameter":"a_value", "second_parameter":"18"},
- #parameters="resources/paremeter.json", You can also pass a path to a json file containing your param
- jenkins_connection_id="your_jenkins_connection" #The connection must be configured first
- )
-
-def grabArtifactFromJenkins(**context):
- """
- Grab an artifact from the previous job
- The python-jenkins library doesn't expose a method for that
- But it's totally possible to build manually the request for that
- """
- hook = JenkinsHook("your_jenkins_connection")
- jenkins_server = hook.get_jenkins_server()
- url = context['task_instance'].xcom_pull(task_ids='trigger_job')
- #The JenkinsJobTriggerOperator store the job url in the xcom variable corresponding to the task
- #You can then use it to access things or to get the job number
- #This url looks like : http://jenkins_url/job/job_name/job_number/
- url = url + "artifact/myartifact.xml" #Or any other artifact name
- request = Request(url)
- response = jenkins_server.jenkins_open(request)
- return response #We store the artifact content in a xcom variable for later use
-
-artifact_grabber = PythonOperator(
- task_id='artifact_grabber',
- provide_context=True,
- python_callable=grabArtifactFromJenkins,
- dag=dag)
-
-artifact_grabber.set_upstream(job_trigger)
diff --git a/airflow/contrib/example_dags/example_pubsub_flow.py b/airflow/contrib/example_dags/example_pubsub_flow.py
deleted file mode 100644
index 6c13ec1daf..0000000000
--- a/airflow/contrib/example_dags/example_pubsub_flow.py
+++ /dev/null
@@ -1,86 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-
-"""
-This example DAG demonstrates how the PubSub*Operators and PubSubPullSensor
-can be used to trigger dependant tasks upon receipt of a Pub/Sub message.
-
-NOTE: project_id must be updated to a GCP project ID accessible with the
- Google Default Credentials on the machine running the workflow
-"""
-from __future__ import unicode_literals
-from base64 import b64encode
-
-import datetime
-
-from airflow import DAG
-from airflow.operators.bash_operator import BashOperator
-from airflow.contrib.operators.pubsub_operator import (
- PubSubTopicCreateOperator, PubSubSubscriptionCreateOperator,
- PubSubPublishOperator, PubSubTopicDeleteOperator,
- PubSubSubscriptionDeleteOperator
-)
-from airflow.contrib.sensors.pubsub_sensor import PubSubPullSensor
-from airflow.utils import dates
-
-project = 'your-project-id' # Change this to your own GCP project_id
-topic = 'example-topic' # Cloud Pub/Sub topic
-subscription = 'subscription-to-example-topic' # Cloud Pub/Sub subscription
-# Sample messages to push/pull
-messages = [
- {'data': b64encode(b'Hello World')},
- {'data': b64encode(b'Another message')},
- {'data': b64encode(b'A final message')}
-]
-
-default_args = {
- 'owner': 'airflow',
- 'depends_on_past': False,
- 'start_date': dates.days_ago(2),
- 'email': ['airflow@example.com'],
- 'email_on_failure': False,
- 'email_on_retry': False,
- 'project': project,
- 'topic': topic,
- 'subscription': subscription,
-}
-
-
-echo_template = '''
-{% for m in task_instance.xcom_pull(task_ids='pull-messages') %}
- echo "AckID: {{ m.get('ackId') }}, Base64-Encoded: {{ m.get('message') }}"
-{% endfor %}
-'''
-
-with DAG('pubsub-end-to-end', default_args=default_args,
- schedule_interval=datetime.timedelta(days=1)) as dag:
- t1 = PubSubTopicCreateOperator(task_id='create-topic')
- t2 = PubSubSubscriptionCreateOperator(
- task_id='create-subscription', topic_project=project,
- subscription=subscription)
- t3 = PubSubPublishOperator(
- task_id='publish-messages', messages=messages)
- t4 = PubSubPullSensor(task_id='pull-messages', ack_messages=True)
- t5 = BashOperator(task_id='echo-pulled-messages',
- bash_command=echo_template)
- t6 = PubSubSubscriptionDeleteOperator(task_id='delete-subscription')
- t7 = PubSubTopicDeleteOperator(task_id='delete-topic')
-
- t1 >> t2 >> t3
- t2 >> t4 >> t5 >> t6 >> t7
diff --git a/airflow/contrib/example_dags/example_qubole_operator.py b/airflow/contrib/example_dags/example_qubole_operator.py
deleted file mode 100644
index 826a50af99..0000000000
--- a/airflow/contrib/example_dags/example_qubole_operator.py
+++ /dev/null
@@ -1,228 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-
-"""
-This is only an example DAG to highlight usage of QuboleOperator in various scenarios,
-some of these tasks may or may not work based on your Qubole account setup.
-
-Run a shell command from Qubole Analyze against your Airflow cluster with following to
-trigger it manually `airflow trigger_dag example_qubole_operator`.
-
-*Note: Make sure that connection `qubole_default` is properly set before running this
-example. Also be aware that it might spin up clusters to run these examples.*
-"""
-
-import airflow
-from airflow import DAG
-from airflow.operators.dummy_operator import DummyOperator
-from airflow.operators.python_operator import PythonOperator, BranchPythonOperator
-from airflow.contrib.operators.qubole_operator import QuboleOperator
-import filecmp
-import random
-
-default_args = {
- 'owner': 'airflow',
- 'depends_on_past': False,
- 'start_date': airflow.utils.dates.days_ago(2),
- 'email': ['airflow@example.com'],
- 'email_on_failure': False,
- 'email_on_retry': False
-}
-
-dag = DAG('example_qubole_operator', default_args=default_args, schedule_interval=None)
-
-dag.doc_md = __doc__
-
-
-def compare_result(ds, **kwargs):
- ti = kwargs['ti']
- r1 = t1.get_results(ti)
- r2 = t2.get_results(ti)
- return filecmp.cmp(r1, r2)
-
-
-t1 = QuboleOperator(
- task_id='hive_show_table',
- command_type='hivecmd',
- query='show tables',
- cluster_label='{{ params.cluster_label }}',
- fetch_logs=True,
- # If `fetch_logs`=true, will fetch qubole command logs and concatenate
- # them into corresponding airflow task logs
- tags='aiflow_example_run',
- # To attach tags to qubole command, auto attach 3 tags - dag_id, task_id, run_id
- qubole_conn_id='qubole_default',
- # Connection id to submit commands inside QDS, if not set "qubole_default" is used
- dag=dag,
- params={
- 'cluster_label': 'default',
- }
-)
-
-t2 = QuboleOperator(
- task_id='hive_s3_location',
- command_type="hivecmd",
- script_location="s3n://public-qubole/qbol-library/scripts/show_table.hql",
- notfiy=True,
- tags=['tag1', 'tag2'],
- # If the script at s3 location has any qubole specific macros to be replaced
- # macros='[{"date": "{{ ds }}"}, {"name" : "abc"}]',
- trigger_rule="all_done",
- dag=dag)
-
-t3 = PythonOperator(
- task_id='compare_result',
- provide_context=True,
- python_callable=compare_result,
- trigger_rule="all_done",
- dag=dag)
-
-t3.set_upstream(t1)
-t3.set_upstream(t2)
-
-options = ['hadoop_jar_cmd', 'presto_cmd', 'db_query', 'spark_cmd']
-
-branching = BranchPythonOperator(
- task_id='branching',
- python_callable=lambda: random.choice(options),
- dag=dag)
-branching.set_upstream(t3)
-
-join = DummyOperator(
- task_id='join',
- trigger_rule='one_success',
- dag=dag
-)
-
-t4 = QuboleOperator(
- task_id='hadoop_jar_cmd',
- command_type='hadoopcmd',
- sub_command='jar s3://paid-qubole/HadoopAPIExamples/'
- 'jars/hadoop-0.20.1-dev-streaming.jar '
- '-mapper wc '
- '-numReduceTasks 0 -input s3://paid-qubole/HadoopAPITests/'
- 'data/3.tsv -output '
- 's3://paid-qubole/HadoopAPITests/data/3_wc',
- cluster_label='{{ params.cluster_label }}',
- fetch_logs=True,
- dag=dag,
- params={
- 'cluster_label': 'default',
- }
-)
-
-t5 = QuboleOperator(
- task_id='pig_cmd',
- command_type="pigcmd",
- script_location="s3://public-qubole/qbol-library/scripts/script1-hadoop-s3-small.pig",
- parameters="key1=value1 key2=value2",
- trigger_rule="all_done",
- dag=dag)
-
-t4.set_upstream(branching)
-t5.set_upstream(t4)
-t5.set_downstream(join)
-
-t6 = QuboleOperator(
- task_id='presto_cmd',
- command_type='prestocmd',
- query='show tables',
- dag=dag)
-
-t7 = QuboleOperator(
- task_id='shell_cmd',
- command_type="shellcmd",
- script_location="s3://public-qubole/qbol-library/scripts/shellx.sh",
- parameters="param1 param2",
- trigger_rule="all_done",
- dag=dag)
-
-t6.set_upstream(branching)
-t7.set_upstream(t6)
-t7.set_downstream(join)
-
-t8 = QuboleOperator(
- task_id='db_query',
- command_type='dbtapquerycmd',
- query='show tables',
- db_tap_id=2064,
- dag=dag)
-
-t9 = QuboleOperator(
- task_id='db_export',
- command_type='dbexportcmd',
- mode=1,
- hive_table='default_qubole_airline_origin_destination',
- db_table='exported_airline_origin_destination',
- partition_spec='dt=20110104-02',
- dbtap_id=2064,
- trigger_rule="all_done",
- dag=dag)
-
-t8.set_upstream(branching)
-t9.set_upstream(t8)
-t9.set_downstream(join)
-
-t10 = QuboleOperator(
- task_id='db_import',
- command_type='dbimportcmd',
- mode=1,
- hive_table='default_qubole_airline_origin_destination',
- db_table='exported_airline_origin_destination',
- where_clause='id < 10',
- db_parallelism=2,
- dbtap_id=2064,
- trigger_rule="all_done",
- dag=dag)
-
-prog = '''
-import scala.math.random
-
-import org.apache.spark._
-
-/** Computes an approximation to pi */
-object SparkPi {
- def main(args: Array[String]) {
- val conf = new SparkConf().setAppName("Spark Pi")
- val spark = new SparkContext(conf)
- val slices = if (args.length > 0) args(0).toInt else 2
- val n = math.min(100000L * slices, Int.MaxValue).toInt // avoid overflow
- val count = spark.parallelize(1 until n, slices).map { i =>
- val x = random * 2 - 1
- val y = random * 2 - 1
- if (x*x + y*y < 1) 1 else 0
- }.reduce(_ + _)
- println("Pi is roughly " + 4.0 * count / n)
- spark.stop()
- }
-}
-'''
-
-t11 = QuboleOperator(
- task_id='spark_cmd',
- command_type="sparkcmd",
- program=prog,
- language='scala',
- arguments='--class SparkPi',
- tags='aiflow_example_run',
- dag=dag)
-
-t11.set_upstream(branching)
-t11.set_downstream(t10)
-t10.set_downstream(join)
diff --git a/airflow/contrib/example_dags/example_qubole_sensor.py b/airflow/contrib/example_dags/example_qubole_sensor.py
deleted file mode 100644
index 3922d3eca9..0000000000
--- a/airflow/contrib/example_dags/example_qubole_sensor.py
+++ /dev/null
@@ -1,79 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-
-"""
-This is only an example DAG to highlight usage of QuboleSensor in various scenarios,
-some of these tasks may or may not work based on your QDS account setup.
-
-Run a shell command from Qubole Analyze against your Airflow cluster with following to
-trigger it manually `airflow trigger_dag example_qubole_sensor`.
-
-*Note: Make sure that connection `qubole_default` is properly set before running
-this example.*
-"""
-
-from airflow import DAG
-from airflow.contrib.sensors.qubole_sensor import QuboleFileSensor, QubolePartitionSensor
-from airflow.utils import dates
-
-default_args = {
- 'owner': 'airflow',
- 'depends_on_past': False,
- 'start_date': dates.days_ago(2),
- 'email': ['airflow@example.com'],
- 'email_on_failure': False,
- 'email_on_retry': False
-}
-
-dag = DAG('example_qubole_sensor', default_args=default_args, schedule_interval=None)
-
-dag.doc_md = __doc__
-
-t1 = QuboleFileSensor(
- task_id='check_s3_file',
- qubole_conn_id='qubole_default',
- poke_interval=60,
- timeout=600,
- data={
- "files":
- [
- "s3://paid-qubole/HadoopAPIExamples/jars/hadoop-0.20.1-dev-streaming.jar",
- "s3://paid-qubole/HadoopAPITests/data/{{ ds.split('-')[2] }}.tsv"
- ] # will check for availability of all the files in array
- },
- dag=dag
-)
-
-t2 = QubolePartitionSensor(
- task_id='check_hive_partition',
- poke_interval=10,
- timeout=60,
- data={"schema": "default",
- "table": "my_partitioned_table",
- "columns": [
- {"column": "month", "values":
- ["{{ ds.split('-')[1] }}"]},
- {"column": "day", "values":
- ["{{ ds.split('-')[2] }}", "{{ yesterday_ds.split('-')[2] }}"]}
- ] # will check for partitions like [month=12/day=12,month=12/day=13]
- },
- dag=dag
-)
-
-t1.set_downstream(t2)
diff --git a/airflow/contrib/example_dags/example_twitter_README.md b/airflow/contrib/example_dags/example_twitter_README.md
deleted file mode 100644
index 319eac39f6..0000000000
--- a/airflow/contrib/example_dags/example_twitter_README.md
+++ /dev/null
@@ -1,36 +0,0 @@
-# Example Twitter DAG
-
-***Introduction:*** This example dag depicts a typical ETL process and is a perfect use case automation scenario for Airflow. Please note that the main scripts associated with the tasks are returning None. The purpose of this DAG is to demonstrate how to write a functional DAG within Airflow.
-
-**Background:** Twitter is a social networking platform that enables users to send or broadcast short messages (140 Characters). A user has a user ID, i.e. JohnDoe, which is also known as a Twitter Handle. A short message, or tweet, can either be sent directed at another user using the @ symbol (i.e. @JohnDoe) or can be broadcast with a hashtag # followed by the topic name. *As most of the data on twitter is public, and twitter provides a generous API to retrieve these data, Twitter is the so called Gold Mine for Text Mining based data analytic.* This example DAG was driven out of our real use case, where we have used the SEARCH API from twitter to retrieve tweets from yesterday. The DAG is scheduled to run each day, and therefore works in an ETL fashion.
-
-***Overview:*** At first, we need tasks that will get the tweets of our interest and save them on the hard-disk. Then, we need subsequent tasks that will clean and analyze the tweets. Then we want to store these files into HDFS, and load them into a Data Warehousing platform like Hive or HBase. The main reason we have selected Hive here is because it gives us a familiar SQL like interface, and makes our life of writing different queries a lot easier. Finally, the DAG needs to store a summarized result to a traditional database, i.e. MySQL or PostgreSQL, which is used by a reporting or business intelligence application. In other words, we basically want to achieve the following steps:
-
-1. Fetch Tweets
-2. Clean Tweets
-3. Analyze Tweets
-4. Put Tweets to HDFS
-5. Load data to Hive
-6. Save Summary to MySQL
-
-***Screenshot:***
-<img src="http://i.imgur.com/rRpSO12.png" width="99%"/>
-
-***Example Structure:*** In this example dag, we are collecting tweets for four users account or twitter handle. Each twitter handle has two channels, incoming tweets and outgoing tweets. Hence, in this example, by running the fetch_tweet task, we should have eight output files. For better management, each of the eight output files should be saved with the yesterday's date (we are collecting tweets from yesterday), i.e. toTwitter_A_2016-03-21.csv. We are using three kind of operators: PythonOperator, BashOperator, and HiveOperator. However, for this example only the Python scripts are stored externally. Hence this example DAG only has the following directory structure:
-
-The python functions here are just placeholders. In case you are interested to actually make this DAG fully functional, first start with filling out the scripts as separate files and importing them into the DAG with absolute or relative import. My approach was to store the retrieved data in memory using Pandas dataframe first, and then use the built in method to save the CSV file on hard-disk.
-The eight different CSV files are then put into eight different folders within HDFS. Each of the newly inserted files are then loaded into eight different external hive tables. Hive tables can be external or internal. In this case, we are inserting the data right into the table, and so we are making our tables internal. Each file is inserted into the respected Hive table named after the twitter channel, i.e. toTwitter_A or fromTwitter_A. It is also important to note that when we created the tables, we facilitated for partitioning by date using the variable dt and declared comma as the row deliminator. The partitioning is very handy and ensures our query execution time remains constant even with growing volume of data.
-As most probably these folders and hive tables doesn't exist in your system, you will get an error for these tasks within the DAG. If you rebuild a function DAG from this example, make sure those folders and hive tables exists. When you create the table, keep the consideration of table partitioning and declaring comma as the row deliminator in your mind. Furthermore, you may also need to skip headers on each read and ensure that the user under which you have Airflow running has the right permission access. Below is a sample HQL snippet on creating such table:
-```
-CREATE TABLE toTwitter_A(id BIGINT, id_str STRING
- created_at STRING, text STRING)
- PARTITIONED BY (dt STRING)
- ROW FORMAT DELIMITED FIELDS TERMINATED BY ','
- STORED AS TEXTFILE;
- alter table toTwitter_A SET serdeproperties ('skip.header.line.count' = '1');
-```
-When you review the code for the DAG, you will notice that these tasks are generated using for loop. These two for loops could be combined into one loop. However, in most cases, you will be running different analysis on your incoming incoming and outgoing tweets, and hence they are kept separated in this example.
-Final step is a running the broker script, brokerapi.py, which will run queries in Hive and store the summarized data to MySQL in our case. To connect to Hive, pyhs2 library is extremely useful and easy to use. To insert data into MySQL from Python, sqlalchemy is also a good one to use.
-I hope you find this tutorial useful. If you have question feel free to ask me on [Twitter](https://twitter.com/EkhtiarSyed) or via the live Airflow chatroom room in [Gitter](https://gitter.im/airbnb/airflow).<p>
--Ekhtiar Syed
-Last Update: 8-April-2016
diff --git a/airflow/contrib/example_dags/example_twitter_dag.py b/airflow/contrib/example_dags/example_twitter_dag.py
deleted file mode 100644
index c2c8f4dd07..0000000000
--- a/airflow/contrib/example_dags/example_twitter_dag.py
+++ /dev/null
@@ -1,190 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-# --------------------------------------------------------------------------------
-# Written By: Ekhtiar Syed
-# Last Update: 8th April 2016
-# Caveat: This Dag will not run because of missing scripts.
-# The purpose of this is to give you a sample of a real world example DAG!
-# --------------------------------------------------------------------------------
-
-# --------------------------------------------------------------------------------
-# Load The Dependencies
-# --------------------------------------------------------------------------------
-
-import airflow
-from airflow import DAG
-from airflow.operators.bash_operator import BashOperator
-from airflow.operators.python_operator import PythonOperator
-from airflow.operators.hive_operator import HiveOperator
-from datetime import date, timedelta
-
-# --------------------------------------------------------------------------------
-# Create a few placeholder scripts. In practice these would be different python
-# script files, which are imported in this section with absolute or relative imports
-# --------------------------------------------------------------------------------
-
-
-def fetchtweets():
- return None
-
-
-def cleantweets():
- return None
-
-
-def analyzetweets():
- return None
-
-
-def transfertodb():
- return None
-
-
-# --------------------------------------------------------------------------------
-# set default arguments
-# --------------------------------------------------------------------------------
-
-default_args = {
- 'owner': 'Ekhtiar',
- 'depends_on_past': False,
- 'start_date': airflow.utils.dates.days_ago(5),
- 'email': ['airflow@example.com'],
- 'email_on_failure': False,
- 'email_on_retry': False,
- 'retries': 1,
- 'retry_delay': timedelta(minutes=5),
- # 'queue': 'bash_queue',
- # 'pool': 'backfill',
- # 'priority_weight': 10,
- # 'end_date': datetime(2016, 1, 1),
-}
-
-dag = DAG(
- 'example_twitter_dag', default_args=default_args,
- schedule_interval="@daily")
-
-# --------------------------------------------------------------------------------
-# This task should call Twitter API and retrieve tweets from yesterday from and to
-# for the four twitter users (Twitter_A,..,Twitter_D) There should be eight csv
-# output files generated by this task and naming convention
-# is direction(from or to)_twitterHandle_date.csv
-# --------------------------------------------------------------------------------
-
-fetch_tweets = PythonOperator(
- task_id='fetch_tweets',
- python_callable=fetchtweets,
- dag=dag)
-
-# --------------------------------------------------------------------------------
-# Clean the eight files. In this step you can get rid of or cherry pick columns
-# and different parts of the text
-# --------------------------------------------------------------------------------
-
-clean_tweets = PythonOperator(
- task_id='clean_tweets',
- python_callable=cleantweets,
- dag=dag)
-
-clean_tweets.set_upstream(fetch_tweets)
-
-# --------------------------------------------------------------------------------
-# In this section you can use a script to analyze the twitter data. Could simply
-# be a sentiment analysis through algorithms like bag of words or something more
-# complicated. You can also take a look at Web Services to do such tasks
-# --------------------------------------------------------------------------------
-
-analyze_tweets = PythonOperator(
- task_id='analyze_tweets',
- python_callable=analyzetweets,
- dag=dag)
-
-analyze_tweets.set_upstream(clean_tweets)
-
-# --------------------------------------------------------------------------------
-# Although this is the last task, we need to declare it before the next tasks as we
-# will use set_downstream This task will extract summary from Hive data and store
-# it to MySQL
-# --------------------------------------------------------------------------------
-
-hive_to_mysql = PythonOperator(
- task_id='hive_to_mysql',
- python_callable=transfertodb,
- dag=dag)
-
-# --------------------------------------------------------------------------------
-# The following tasks are generated using for loop. The first task puts the eight
-# csv files to HDFS. The second task loads these files from HDFS to respected Hive
-# tables. These two for loops could be combined into one loop. However, in most cases,
-# you will be running different analysis on your incoming incoming and outgoing tweets,
-# and hence they are kept separated in this example.
-# --------------------------------------------------------------------------------
-
-from_channels = ['fromTwitter_A', 'fromTwitter_B', 'fromTwitter_C', 'fromTwitter_D']
-to_channels = ['toTwitter_A', 'toTwitter_B', 'toTwitter_C', 'toTwitter_D']
-yesterday = date.today() - timedelta(days=1)
-dt = yesterday.strftime("%Y-%m-%d")
-# define where you want to store the tweets csv file in your local directory
-local_dir = "/tmp/"
-# define the location where you want to store in HDFS
-hdfs_dir = " /tmp/"
-
-for channel in to_channels:
-
- file_name = "to_" + channel + "_" + yesterday.strftime("%Y-%m-%d") + ".csv"
-
- load_to_hdfs = BashOperator(
- task_id="put_" + channel + "_to_hdfs",
- bash_command="HADOOP_USER_NAME=hdfs hadoop fs -put -f " +
- local_dir + file_name +
- hdfs_dir + channel + "/",
- dag=dag)
-
- load_to_hdfs.set_upstream(analyze_tweets)
-
- load_to_hive = HiveOperator(
- task_id="load_" + channel + "_to_hive",
- hql="LOAD DATA INPATH '" +
- hdfs_dir + channel + "/" + file_name + "' "
- "INTO TABLE " + channel + " "
- "PARTITION(dt='" + dt + "')",
- dag=dag)
- load_to_hive.set_upstream(load_to_hdfs)
- load_to_hive.set_downstream(hive_to_mysql)
-
-for channel in from_channels:
- file_name = "from_" + channel + "_" + yesterday.strftime("%Y-%m-%d") + ".csv"
- load_to_hdfs = BashOperator(
- task_id="put_" + channel + "_to_hdfs",
- bash_command="HADOOP_USER_NAME=hdfs hadoop fs -put -f " +
- local_dir + file_name +
- hdfs_dir + channel + "/",
- dag=dag)
-
- load_to_hdfs.set_upstream(analyze_tweets)
-
- load_to_hive = HiveOperator(
- task_id="load_" + channel + "_to_hive",
- hql="LOAD DATA INPATH '" +
- hdfs_dir + channel + "/" + file_name + "' "
- "INTO TABLE " + channel + " "
- "PARTITION(dt='" + dt + "')",
- dag=dag)
-
- load_to_hive.set_upstream(load_to_hdfs)
- load_to_hive.set_downstream(hive_to_mysql)
diff --git a/airflow/contrib/example_dags/example_winrm_operator.py b/airflow/contrib/example_dags/example_winrm_operator.py
deleted file mode 100644
index 195bf5d98d..0000000000
--- a/airflow/contrib/example_dags/example_winrm_operator.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-# --------------------------------------------------------------------------------
-# Written By: Ekhtiar Syed
-# Last Update: 8th April 2016
-# Caveat: This Dag will not run because of missing scripts.
-# The purpose of this is to give you a sample of a real world example DAG!
-# --------------------------------------------------------------------------------
-
-# --------------------------------------------------------------------------------
-# Load The Dependencies
-# --------------------------------------------------------------------------------
-import airflow
-from airflow.operators.dummy_operator import DummyOperator
-from airflow.models import DAG
-from datetime import timedelta
-
-from airflow.contrib.hooks import WinRMHook
-from airflow.contrib.operators.winrm_operator import WinRMOperator
-
-
-args = {
- 'owner': 'airflow',
- 'start_date': airflow.utils.dates.days_ago(2)
-}
-
-dag = DAG(
- dag_id='POC_winrm_parallel', default_args=args,
- schedule_interval='0 0 * * *',
- dagrun_timeout=timedelta(minutes=60))
-
-cmd = 'ls -l'
-run_this_last = DummyOperator(task_id='run_this_last', dag=dag)
-
-winRMHook = WinRMHook(ssh_conn_id='ssh_POC1')
-
-t1 = WinRMOperator(
- task_id="wintask1",
- command='ls -altr',
- winrm_hook=winRMHook,
- dag=dag)
-
-t2 = WinRMOperator(
- task_id="wintask2",
- command='sleep 60',
- winrm_hook=winRMHook,
- dag=dag)
-
-t3 = WinRMOperator(
- task_id="wintask3",
- command='echo \'luke test\' ',
- winrm_hook=winRMHook,
- dag=dag)
-
-t1.set_downstream(run_this_last)
-t2.set_downstream(run_this_last)
-t3.set_downstream(run_this_last)
diff --git a/airflow/contrib/hooks/__init__.py b/airflow/contrib/hooks/__init__.py
deleted file mode 100644
index 95f9892d89..0000000000
--- a/airflow/contrib/hooks/__init__.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-#
-
-
-# Contrib hooks are not imported by default. They should be accessed
-# directly: from airflow.contrib.hooks.hook_module import Hook
-
-
-import sys
-import os as _os
-
-# ------------------------------------------------------------------------
-#
-# #TODO #FIXME Airflow 2.0
-#
-# Old import machinary below.
-#
-# This is deprecated but should be kept until Airflow 2.0
-# for compatibility.
-#
-# ------------------------------------------------------------------------
-_hooks = {
- 'docker_hook': ['DockerHook'],
- 'ftp_hook': ['FTPHook'],
- 'ftps_hook': ['FTPSHook'],
- 'vertica_hook': ['VerticaHook'],
- 'ssh_hook': ['SSHHook'],
- 'winrm_hook': ['WinRMHook'],
- 'sftp_hook': ['SFTPHook'],
- 'bigquery_hook': ['BigQueryHook'],
- 'qubole_hook': ['QuboleHook'],
- 'gcs_hook': ['GoogleCloudStorageHook'],
- 'datastore_hook': ['DatastoreHook'],
- 'gcp_cloudml_hook': ['CloudMLHook'],
- 'redshift_hook': ['RedshiftHook'],
- 'gcp_dataproc_hook': ['DataProcHook'],
- 'gcp_dataflow_hook': ['DataFlowHook'],
- 'spark_submit_operator': ['SparkSubmitOperator'],
- 'cloudant_hook': ['CloudantHook'],
- 'fs_hook': ['FSHook'],
- 'wasb_hook': ['WasbHook'],
- 'gcp_pubsub_hook': ['PubSubHook'],
- 'jenkins_hook': ['JenkinsHook'],
- 'aws_dynamodb_hook': ['AwsDynamoDBHook'],
- 'azure_data_lake_hook': ['AzureDataLakeHook'],
- 'azure_fileshare_hook': ['AzureFileShareHook'],
-}
-
-
-if not _os.environ.get('AIRFLOW_USE_NEW_IMPORTS', False):
- from airflow.utils.helpers import AirflowImporter
-
- airflow_importer = AirflowImporter(sys.modules[__name__], _hooks)
diff --git a/airflow/contrib/hooks/aws_dynamodb_hook.py b/airflow/contrib/hooks/aws_dynamodb_hook.py
deleted file mode 100644
index 69324c6dba..0000000000
--- a/airflow/contrib/hooks/aws_dynamodb_hook.py
+++ /dev/null
@@ -1,69 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.exceptions import AirflowException
-from airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class AwsDynamoDBHook(AwsHook):
- """
- Interact with AWS DynamoDB.
-
- :param table_keys: partition key and sort key
- :type table_keys: list
- :param table_name: target DynamoDB table
- :type table_name: str
- :param region_name: aws region name (example: us-east-1)
- :type region_name: str
- """
-
- def __init__(self,
- table_keys=None,
- table_name=None,
- region_name=None,
- *args, **kwargs):
- self.table_keys = table_keys
- self.table_name = table_name
- self.region_name = region_name
- super(AwsDynamoDBHook, self).__init__(*args, **kwargs)
-
- def get_conn(self):
- self.conn = self.get_resource_type('dynamodb', self.region_name)
- return self.conn
-
- def write_batch_data(self, items):
- """
- Write batch items to dynamodb table with provisioned throughout capacity.
- """
-
- dynamodb_conn = self.get_conn()
-
- try:
- table = dynamodb_conn.Table(self.table_name)
-
- with table.batch_writer(overwrite_by_pkeys=self.table_keys) as batch:
- for item in items:
- batch.put_item(Item=item)
- return True
- except Exception as general_error:
- raise AirflowException(
- 'Failed to insert items in dynamodb, error: {error}'.format(
- error=str(general_error)
- )
- )
diff --git a/airflow/contrib/hooks/aws_firehose_hook.py b/airflow/contrib/hooks/aws_firehose_hook.py
deleted file mode 100644
index b273e22b54..0000000000
--- a/airflow/contrib/hooks/aws_firehose_hook.py
+++ /dev/null
@@ -1,56 +0,0 @@
-#
-# 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 airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class AwsFirehoseHook(AwsHook):
- """
- Interact with AWS Kinesis Firehose.
- :param delivery_stream: Name of the delivery stream
- :type delivery_stream: str
- :param region_name: AWS region name (example: us-east-1)
- :type region_name: str
- """
-
- def __init__(self, delivery_stream, region_name=None, *args, **kwargs):
- self.delivery_stream = delivery_stream
- self.region_name = region_name
- super(AwsFirehoseHook, self).__init__(*args, **kwargs)
-
- def get_conn(self):
- """
- Returns AwsHook connection object.
- """
-
- self.conn = self.get_client_type('firehose', self.region_name)
- return self.conn
-
- def put_records(self, records):
- """
- Write batch records to Kinesis Firehose
- """
-
- firehose_conn = self.get_conn()
-
- response = firehose_conn.put_record_batch(
- DeliveryStreamName=self.delivery_stream,
- Records=records
- )
-
- return response
diff --git a/airflow/contrib/hooks/aws_hook.py b/airflow/contrib/hooks/aws_hook.py
deleted file mode 100644
index c712d2de79..0000000000
--- a/airflow/contrib/hooks/aws_hook.py
+++ /dev/null
@@ -1,177 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 boto3
-import configparser
-import logging
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-
-
-def _parse_s3_config(config_file_name, config_format='boto', profile=None):
- """
- Parses a config file for s3 credentials. Can currently
- parse boto, s3cmd.conf and AWS SDK config formats
-
- :param config_file_name: path to the config file
- :type config_file_name: str
- :param config_format: config type. One of "boto", "s3cmd" or "aws".
- Defaults to "boto"
- :type config_format: str
- :param profile: profile name in AWS type config file
- :type profile: str
- """
- config = configparser.ConfigParser()
- if config.read(config_file_name): # pragma: no cover
- sections = config.sections()
- else:
- raise AirflowException("Couldn't read {0}".format(config_file_name))
- # Setting option names depending on file format
- if config_format is None:
- config_format = 'boto'
- conf_format = config_format.lower()
- if conf_format == 'boto': # pragma: no cover
- if profile is not None and 'profile ' + profile in sections:
- cred_section = 'profile ' + profile
- else:
- cred_section = 'Credentials'
- elif conf_format == 'aws' and profile is not None:
- cred_section = profile
- else:
- cred_section = 'default'
- # Option names
- if conf_format in ('boto', 'aws'): # pragma: no cover
- key_id_option = 'aws_access_key_id'
- secret_key_option = 'aws_secret_access_key'
- # security_token_option = 'aws_security_token'
- else:
- key_id_option = 'access_key'
- secret_key_option = 'secret_key'
- # Actual Parsing
- if cred_section not in sections:
- raise AirflowException("This config file format is not recognized")
- else:
- try:
- access_key = config.get(cred_section, key_id_option)
- secret_key = config.get(cred_section, secret_key_option)
- except:
- logging.warning("Option Error in parsing s3 config file")
- raise
- return access_key, secret_key
-
-
-class AwsHook(BaseHook):
- """
- Interact with AWS.
- This class is a thin wrapper around the boto3 python library.
- """
-
- def __init__(self, aws_conn_id='aws_default'):
- self.aws_conn_id = aws_conn_id
-
- def _get_credentials(self, region_name):
- aws_access_key_id = None
- aws_secret_access_key = None
- aws_session_token = None
- endpoint_url = None
-
- if self.aws_conn_id:
- try:
- connection_object = self.get_connection(self.aws_conn_id)
- if connection_object.login:
- aws_access_key_id = connection_object.login
- aws_secret_access_key = connection_object.password
-
- elif 'aws_secret_access_key' in connection_object.extra_dejson:
- aws_access_key_id = connection_object.extra_dejson[
- 'aws_access_key_id']
- aws_secret_access_key = connection_object.extra_dejson[
- 'aws_secret_access_key']
-
- elif 's3_config_file' in connection_object.extra_dejson:
- aws_access_key_id, aws_secret_access_key = \
- _parse_s3_config(
- connection_object.extra_dejson['s3_config_file'],
- connection_object.extra_dejson.get('s3_config_format'))
-
- if region_name is None:
- region_name = connection_object.extra_dejson.get('region_name')
-
- role_arn = connection_object.extra_dejson.get('role_arn')
- aws_account_id = connection_object.extra_dejson.get('aws_account_id')
- aws_iam_role = connection_object.extra_dejson.get('aws_iam_role')
-
- if role_arn is None and aws_account_id is not None and \
- aws_iam_role is not None:
- role_arn = "arn:aws:iam::" + aws_account_id + ":role/" + aws_iam_role
-
- if role_arn is not None:
- sts_session = boto3.session.Session(
- aws_access_key_id=aws_access_key_id,
- aws_secret_access_key=aws_secret_access_key,
- region_name=region_name)
-
- sts_client = sts_session.client('sts')
- sts_response = sts_client.assume_role(
- RoleArn=role_arn,
- RoleSessionName='Airflow_' + self.aws_conn_id)
- aws_access_key_id = sts_response['Credentials']['AccessKeyId']
- aws_secret_access_key = sts_response['Credentials']['SecretAccessKey']
- aws_session_token = sts_response['Credentials']['SessionToken']
-
- endpoint_url = connection_object.extra_dejson.get('host')
-
- except AirflowException:
- # No connection found: fallback on boto3 credential strategy
- # http://boto3.readthedocs.io/en/latest/guide/configuration.html
- pass
-
- return boto3.session.Session(
- aws_access_key_id=aws_access_key_id,
- aws_secret_access_key=aws_secret_access_key,
- aws_session_token=aws_session_token,
- region_name=region_name), endpoint_url
-
- def get_client_type(self, client_type, region_name=None):
- session, endpoint_url = self._get_credentials(region_name)
-
- return session.client(client_type, endpoint_url=endpoint_url)
-
- def get_resource_type(self, resource_type, region_name=None):
- session, endpoint_url = self._get_credentials(region_name)
-
- return session.resource(resource_type, endpoint_url=endpoint_url)
-
- def get_session(self, region_name=None):
- """Get the underlying boto3.session."""
- session, _ = self._get_credentials(region_name)
- return session
-
- def get_credentials(self, region_name=None):
- """Get the underlying `botocore.Credentials` object.
-
- This contains the attributes: access_key, secret_key and token.
- """
- session, _ = self._get_credentials(region_name)
- # Credentials are refreshable, so accessing your access key / secret key
- # separately can lead to a race condition.
- # See https://stackoverflow.com/a/36291428/8283373
- return session.get_credentials().get_frozen_credentials()
diff --git a/airflow/contrib/hooks/aws_lambda_hook.py b/airflow/contrib/hooks/aws_lambda_hook.py
deleted file mode 100644
index 57cf716fda..0000000000
--- a/airflow/contrib/hooks/aws_lambda_hook.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class AwsLambdaHook(AwsHook):
- """
- Interact with AWS Lambda
-
- :param function_name: AWS Lambda Function Name
- :type function_name: str
- :param region_name: AWS Region Name (example: us-west-2)
- :type region_name: str
- :param log_type: Tail Invocation Request
- :type log_type: str
- :param qualifier: AWS Lambda Function Version or Alias Name
- :type qualifier: str
- :param invocation_type: AWS Lambda Invocation Type (RequestResponse, Event etc)
- :type invocation_type: str
- """
-
- def __init__(self, function_name, region_name=None,
- log_type='None', qualifier='$LATEST',
- invocation_type='RequestResponse', *args, **kwargs):
- self.function_name = function_name
- self.region_name = region_name
- self.log_type = log_type
- self.invocation_type = invocation_type
- self.qualifier = qualifier
- super(AwsLambdaHook, self).__init__(*args, **kwargs)
-
- def get_conn(self):
- self.conn = self.get_client_type('lambda', self.region_name)
- return self.conn
-
- def invoke_lambda(self, payload):
- """
- Invoke Lambda Function
- """
-
- awslambda_conn = self.get_conn()
-
- response = awslambda_conn.invoke(
- FunctionName=self.function_name,
- InvocationType=self.invocation_type,
- LogType=self.log_type,
- Payload=payload,
- Qualifier=self.qualifier
- )
-
- return response
diff --git a/airflow/contrib/hooks/azure_data_lake_hook.py b/airflow/contrib/hooks/azure_data_lake_hook.py
deleted file mode 100644
index 1a02d78f27..0000000000
--- a/airflow/contrib/hooks/azure_data_lake_hook.py
+++ /dev/null
@@ -1,141 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.hooks.base_hook import BaseHook
-from azure.datalake.store import core, lib, multithread
-
-
-class AzureDataLakeHook(BaseHook):
- """
- Interacts with Azure Data Lake.
-
- Client ID and client secret should be in user and password parameters.
- Tenant and account name should be extra field as
- {"tenant": "<TENANT>", "account_name": "ACCOUNT_NAME"}.
-
- :param azure_data_lake_conn_id: Reference to the Azure Data Lake connection.
- :type azure_data_lake_conn_id: str
- """
-
- def __init__(self, azure_data_lake_conn_id='azure_data_lake_default'):
- self.conn_id = azure_data_lake_conn_id
- self.connection = self.get_conn()
-
- def get_conn(self):
- """Return a AzureDLFileSystem object."""
- conn = self.get_connection(self.conn_id)
- service_options = conn.extra_dejson
- self.account_name = service_options.get('account_name')
-
- adlCreds = lib.auth(tenant_id=service_options.get('tenant'),
- client_secret=conn.password,
- client_id=conn.login)
- adlsFileSystemClient = core.AzureDLFileSystem(adlCreds,
- store_name=self.account_name)
- adlsFileSystemClient.connect()
- return adlsFileSystemClient
-
- def check_for_file(self, file_path):
- """
- Check if a file exists on Azure Data Lake.
-
- :param file_path: Path and name of the file.
- :type file_path: str
- :return: True if the file exists, False otherwise.
- :rtype bool
- """
- try:
- files = self.connection.glob(file_path, details=False, invalidate_cache=True)
- return len(files) == 1
- except FileNotFoundError:
- return False
-
- def upload_file(self, local_path, remote_path, nthreads=64, overwrite=True,
- buffersize=4194304, blocksize=4194304):
- """
- Upload a file to Azure Data Lake.
-
- :param local_path: local path. Can be single file, directory (in which case,
- upload recursively) or glob pattern. Recursive glob patterns using `**`
- are not supported.
- :type local_path: str
- :param remote_path: Remote path to upload to; if multiple files, this is the
- dircetory root to write within.
- :type remote_path: str
- :param nthreads: Number of threads to use. If None, uses the number of cores.
- :type nthreads: int
- :param overwrite: Whether to forcibly overwrite existing files/directories.
- If False and remote path is a directory, will quit regardless if any files
- would be overwritten or not. If True, only matching filenames are actually
- overwritten.
- :type overwrite: bool
- :param buffersize: int [2**22]
- Number of bytes for internal buffer. This block cannot be bigger than
- a chunk and cannot be smaller than a block.
- :type buffersize: int
- :param blocksize: int [2**22]
- Number of bytes for a block. Within each chunk, we write a smaller
- block for each API call. This block cannot be bigger than a chunk.
- :type blocksize: int
- """
- multithread.ADLUploader(self.connection,
- lpath=local_path,
- rpath=remote_path,
- nthreads=nthreads,
- overwrite=overwrite,
- buffersize=buffersize,
- blocksize=blocksize)
-
- def download_file(self, local_path, remote_path, nthreads=64, overwrite=True,
- buffersize=4194304, blocksize=4194304):
- """
- Download a file from Azure Blob Storage.
-
- :param local_path: local path. If downloading a single file, will write to this
- specific file, unless it is an existing directory, in which case a file is
- created within it. If downloading multiple files, this is the root
- directory to write within. Will create directories as required.
- :type local_path: str
- :param remote_path: remote path/globstring to use to find remote files.
- Recursive glob patterns using `**` are not supported.
- :type remote_path: str
- :param nthreads: Number of threads to use. If None, uses the number of cores.
- :type nthreads: int
- :param overwrite: Whether to forcibly overwrite existing files/directories.
- If False and remote path is a directory, will quit regardless if any files
- would be overwritten or not. If True, only matching filenames are actually
- overwritten.
- :type overwrite: bool
- :param buffersize: int [2**22]
- Number of bytes for internal buffer. This block cannot be bigger than
- a chunk and cannot be smaller than a block.
- :type buffersize: int
- :param blocksize: int [2**22]
- Number of bytes for a block. Within each chunk, we write a smaller
- block for each API call. This block cannot be bigger than a chunk.
- :type blocksize: int
- """
- multithread.ADLDownloader(self.connection,
- lpath=local_path,
- rpath=remote_path,
- nthreads=nthreads,
- overwrite=overwrite,
- buffersize=buffersize,
- blocksize=blocksize)
diff --git a/airflow/contrib/hooks/azure_fileshare_hook.py b/airflow/contrib/hooks/azure_fileshare_hook.py
deleted file mode 100644
index edabc17293..0000000000
--- a/airflow/contrib/hooks/azure_fileshare_hook.py
+++ /dev/null
@@ -1,212 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.hooks.base_hook import BaseHook
-from azure.storage.file import FileService
-
-
-class AzureFileShareHook(BaseHook):
- """
- Interacts with Azure FileShare Storage.
-
- Additional options passed in the 'extra' field of the connection will be
- passed to the `FileService()` constructor.
-
- :param wasb_conn_id: Reference to the wasb connection.
- :type wasb_conn_id: str
- """
-
- def __init__(self, wasb_conn_id='wasb_default'):
- self.conn_id = wasb_conn_id
- self.connection = self.get_conn()
-
- def get_conn(self):
- """Return the FileService object."""
- conn = self.get_connection(self.conn_id)
- service_options = conn.extra_dejson
- return FileService(account_name=conn.login,
- account_key=conn.password, **service_options)
-
- def check_for_directory(self, share_name, directory_name, **kwargs):
- """
- Check if a directory exists on Azure File Share.
-
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.exists()` takes.
- :type kwargs: object
- :return: True if the file exists, False otherwise.
- :rtype bool
- """
- return self.connection.exists(share_name, directory_name,
- **kwargs)
-
- def check_for_file(self, share_name, directory_name, file_name, **kwargs):
- """
- Check if a file exists on Azure File Share.
-
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param file_name: Name of the file.
- :type file_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.exists()` takes.
- :type kwargs: object
- :return: True if the file exists, False otherwise.
- :rtype bool
- """
- return self.connection.exists(share_name, directory_name,
- file_name, **kwargs)
-
- def list_directories_and_files(self, share_name, directory_name=None, **kwargs):
- """
- Return the list of directories and files stored on a Azure File Share.
-
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.list_directories_and_files()` takes.
- :type kwargs: object
- :return: A list of files and directories
- :rtype list
- """
- return self.connection.list_directories_and_files(share_name,
- directory_name,
- **kwargs)
-
- def create_directory(self, share_name, directory_name, **kwargs):
- """
- Create a new direcotry on a Azure File Share.
-
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.create_directory()` takes.
- :type kwargs: object
- :return: A list of files and directories
- :rtype list
- """
- return self.connection.create_directory(share_name, directory_name, **kwargs)
-
- def get_file(self, file_path, share_name, directory_name, file_name, **kwargs):
- """
- Download a file from Azure File Share.
-
- :param file_path: Where to store the file.
- :type file_path: str
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param file_name: Name of the file.
- :type file_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.get_file_to_path()` takes.
- :type kwargs: object
- """
- self.connection.get_file_to_path(share_name, directory_name,
- file_name, file_path, **kwargs)
-
- def get_file_to_stream(self, stream, share_name, directory_name, file_name, **kwargs):
- """
- Download a file from Azure File Share.
-
- :param stream: A filehandle to store the file to.
- :type stream: file-like object
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param file_name: Name of the file.
- :type file_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.get_file_to_stream()` takes.
- :type kwargs: object
- """
- self.connection.get_file_to_stream(share_name, directory_name,
- file_name, stream, **kwargs)
-
- def load_file(self, file_path, share_name, directory_name, file_name, **kwargs):
- """
- Upload a file to Azure File Share.
-
- :param file_path: Path to the file to load.
- :type file_path: str
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param file_name: Name of the file.
- :type file_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.create_file_from_path()` takes.
- :type kwargs: object
- """
- self.connection.create_file_from_path(share_name, directory_name,
- file_name, file_path, **kwargs)
-
- def load_string(self, string_data, share_name, directory_name, file_name, **kwargs):
- """
- Upload a string to Azure File Share.
-
- :param string_data: String to load.
- :type string_data: str
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param file_name: Name of the file.
- :type file_name: str
- :param kwargs: Optional keyword arguments that
- `FileService.create_file_from_text()` takes.
- :type kwargs: object
- """
- self.connection.create_file_from_text(share_name, directory_name,
- file_name, string_data, **kwargs)
-
- def load_stream(self, stream, share_name, directory_name, file_name, count, **kwargs):
- """
- Upload a stream to Azure File Share.
-
- :param stream: Opened file/stream to upload as the file content.
- :type stream: file-like
- :param share_name: Name of the share.
- :type share_name: str
- :param directory_name: Name of the directory.
- :type directory_name: str
- :param file_name: Name of the file.
- :type file_name: str
- :param count: Size of the stream in bytes
- :type count: int
- :param kwargs: Optional keyword arguments that
- `FileService.create_file_from_stream()` takes.
- :type kwargs: object
- """
- self.connection.create_file_from_stream(share_name, directory_name,
- file_name, stream, count, **kwargs)
diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py
deleted file mode 100644
index b452f51c22..0000000000
--- a/airflow/contrib/hooks/bigquery_hook.py
+++ /dev/null
@@ -1,1534 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-#
-"""
-This module contains a BigQuery Hook, as well as a very basic PEP 249
-implementation for BigQuery.
-"""
-
-import time
-from builtins import range
-
-from past.builtins import basestring
-
-from airflow import AirflowException
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-from airflow.hooks.dbapi_hook import DbApiHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-from apiclient.discovery import HttpError, build
-from googleapiclient import errors
-from pandas_gbq.gbq import \
- _check_google_client_version as gbq_check_google_client_version
-from pandas_gbq import read_gbq
-from pandas_gbq.gbq import \
- _test_google_api_imports as gbq_test_google_api_imports
-from pandas_gbq.gbq import GbqConnector
-
-
-class BigQueryHook(GoogleCloudBaseHook, DbApiHook, LoggingMixin):
- """
- Interact with BigQuery. This hook uses the Google Cloud Platform
- connection.
- """
- conn_name_attr = 'bigquery_conn_id'
-
- def __init__(self,
- bigquery_conn_id='bigquery_default',
- delegate_to=None,
- use_legacy_sql=True):
- super(BigQueryHook, self).__init__(
- gcp_conn_id=bigquery_conn_id, delegate_to=delegate_to)
- self.use_legacy_sql = use_legacy_sql
-
- def get_conn(self):
- """
- Returns a BigQuery PEP 249 connection object.
- """
- service = self.get_service()
- project = self._get_field('project')
- return BigQueryConnection(
- service=service,
- project_id=project,
- use_legacy_sql=self.use_legacy_sql)
-
- def get_service(self):
- """
- Returns a BigQuery service object.
- """
- http_authorized = self._authorize()
- return build(
- 'bigquery', 'v2', http=http_authorized, cache_discovery=False)
-
- def insert_rows(self, table, rows, target_fields=None, commit_every=1000):
- """
- Insertion is currently unsupported. Theoretically, you could use
- BigQuery's streaming API to insert rows into a table, but this hasn't
- been implemented.
- """
- raise NotImplementedError()
-
- def get_pandas_df(self, sql, parameters=None, dialect=None):
- """
- Returns a Pandas DataFrame for the results produced by a BigQuery
- query. The DbApiHook method must be overridden because Pandas
- doesn't support PEP 249 connections, except for SQLite. See:
-
- https://github.com/pydata/pandas/blob/master/pandas/io/sql.py#L447
- https://github.com/pydata/pandas/issues/6900
-
- :param sql: The BigQuery SQL to execute.
- :type sql: string
- :param parameters: The parameters to render the SQL query with (not
- used, leave to override superclass method)
- :type parameters: mapping or iterable
- :param dialect: Dialect of BigQuery SQL – legacy SQL or standard SQL
- defaults to use `self.use_legacy_sql` if not specified
- :type dialect: string in {'legacy', 'standard'}
- """
- if dialect is None:
- dialect = 'legacy' if self.use_legacy_sql else 'standard'
-
- return read_gbq(sql,
- project_id=self._get_field('project'),
- dialect=dialect,
- verbose=False)
-
- def table_exists(self, project_id, dataset_id, table_id):
- """
- Checks for the existence of a table in Google BigQuery.
-
- :param project_id: The Google cloud project in which to look for the
- table. The connection supplied to the hook must provide access to
- the specified project.
- :type project_id: string
- :param dataset_id: The name of the dataset in which to look for the
- table.
- :type dataset_id: string
- :param table_id: The name of the table to check the existence of.
- :type table_id: string
- """
- service = self.get_service()
- try:
- service.tables().get(
- projectId=project_id, datasetId=dataset_id,
- tableId=table_id).execute()
- return True
- except errors.HttpError as e:
- if e.resp['status'] == '404':
- return False
- raise
-
-
-class BigQueryPandasConnector(GbqConnector):
- """
- This connector behaves identically to GbqConnector (from Pandas), except
- that it allows the service to be injected, and disables a call to
- self.get_credentials(). This allows Airflow to use BigQuery with Pandas
- without forcing a three legged OAuth connection. Instead, we can inject
- service account credentials into the binding.
- """
-
- def __init__(self,
- project_id,
- service,
- reauth=False,
- verbose=False,
- dialect='legacy'):
- super(BigQueryPandasConnector, self).__init__(project_id)
- gbq_check_google_client_version()
- gbq_test_google_api_imports()
- self.project_id = project_id
- self.reauth = reauth
- self.service = service
- self.verbose = verbose
- self.dialect = dialect
-
-
-class BigQueryConnection(object):
- """
- BigQuery does not have a notion of a persistent connection. Thus, these
- objects are small stateless factories for cursors, which do all the real
- work.
- """
-
- def __init__(self, *args, **kwargs):
- self._args = args
- self._kwargs = kwargs
-
- def close(self):
- """ BigQueryConnection does not have anything to close. """
- pass
-
- def commit(self):
- """ BigQueryConnection does not support transactions. """
- pass
-
- def cursor(self):
- """ Return a new :py:class:`Cursor` object using the connection. """
- return BigQueryCursor(*self._args, **self._kwargs)
-
- def rollback(self):
- raise NotImplementedError(
- "BigQueryConnection does not have transactions")
-
-
-class BigQueryBaseCursor(LoggingMixin):
- """
- The BigQuery base cursor contains helper methods to execute queries against
- BigQuery. The methods can be used directly by operators, in cases where a
- PEP 249 cursor isn't needed.
- """
-
- def __init__(self, service, project_id, use_legacy_sql=True):
- self.service = service
- self.project_id = project_id
- self.use_legacy_sql = use_legacy_sql
- self.running_job_id = None
-
- def create_empty_table(self,
- project_id,
- dataset_id,
- table_id,
- schema_fields=None,
- time_partitioning={}
- ):
- """
- Creates a new, empty table in the dataset.
-
- :param project_id: The project to create the table into.
- :type project_id: str
- :param dataset_id: The dataset to create the table into.
- :type dataset_id: str
- :param table_id: The Name of the table to be created.
- :type table_id: str
- :param schema_fields: If set, the schema field list as defined here:
- https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schema
-
- **Example**: ::
-
- schema_fields=[{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"},
- {"name": "salary", "type": "INTEGER", "mode": "NULLABLE"}]
-
- :type schema_fields: list
- :param time_partitioning: configure optional time partitioning fields i.e.
- partition by field, type and expiration as per API specifications.
-
- .. seealso::
- https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#timePartitioning
- :type time_partitioning: dict
-
- :return:
- """
- project_id = project_id if project_id is not None else self.project_id
-
- table_resource = {
- 'tableReference': {
- 'tableId': table_id
- }
- }
-
- if schema_fields:
- table_resource['schema'] = {'fields': schema_fields}
-
- if time_partitioning:
- table_resource['timePartitioning'] = time_partitioning
-
- self.log.info('Creating Table %s:%s.%s',
- project_id, dataset_id, table_id)
-
- try:
- self.service.tables().insert(
- projectId=project_id,
- datasetId=dataset_id,
- body=table_resource).execute()
-
- self.log.info('Table created successfully: %s:%s.%s',
- project_id, dataset_id, table_id)
-
- except HttpError as err:
- raise AirflowException(
- 'BigQuery job failed. Error was: {}'.format(err.content)
- )
-
- def create_external_table(self,
- external_project_dataset_table,
- schema_fields,
- source_uris,
- source_format='CSV',
- autodetect=False,
- compression='NONE',
- ignore_unknown_values=False,
- max_bad_records=0,
- skip_leading_rows=0,
- field_delimiter=',',
- quote_character=None,
- allow_quoted_newlines=False,
- allow_jagged_rows=False,
- src_fmt_configs={}
- ):
- """
- Creates a new external table in the dataset with the data in Google
- Cloud Storage. See here:
-
- https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource
-
- for more details about these parameters.
-
- :param external_project_dataset_table:
- The dotted (<project>.|<project>:)<dataset>.<table>($<partition>) BigQuery
- table name to create external table.
- If <project> is not included, project will be the
- project defined in the connection json.
- :type external_project_dataset_table: string
- :param schema_fields: The schema field list as defined here:
- https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource
- :type schema_fields: list
- :param source_uris: The source Google Cloud
- Storage URI (e.g. gs://some-bucket/some-file.txt). A single wild
- per-object name can be used.
- :type source_uris: list
- :param source_format: File format to export.
- :type source_format: string
- :param autodetect: Try to detect schema and format options automatically.
- Any option specified explicitly will be honored.
- :type autodetect: bool
- :param compression: [Optional] The compression type of the data source.
- Possible values include GZIP and NONE.
- The default value is NONE.
- This setting is ignored for Google Cloud Bigtable,
- Google Cloud Datastore backups and Avro formats.
- :type compression: string
- :param ignore_unknown_values: [Optional] Indicates if BigQuery should allow
- extra values that are not represented in the table schema.
- If true, the extra values are ignored. If false, records with extra columns
- are treated as bad records, and if there are too many bad records, an
- invalid error is returned in the job result.
- :type ignore_unknown_values: bool
- :param max_bad_records: The maximum number of bad records that BigQuery can
- ignore when running the job.
- :type max_bad_records: int
- :param skip_leading_rows: Number of rows to skip when loading from a CSV.
- :type skip_leading_rows: int
- :param field_delimiter: The delimiter to use when loading from a CSV.
- :type field_delimiter: string
- :param quote_character: The value that is used to quote data sections in a CSV
- file.
- :type quote_character: string
- :param allow_quoted_newlines: Whether to allow quoted newlines (true) or not
- (false).
- :type allow_quoted_newlines: boolean
- :param allow_jagged_rows: Accept rows that are missing trailing optional columns.
- The missing values are treated as nulls. If false, records with missing
- trailing columns are treated as bad records, and if there are too many bad
- records, an invalid error is returned in the job result. Only applicable when
- soure_format is CSV.
- :type allow_jagged_rows: bool
- :param src_fmt_configs: configure optional fields specific to the source format
- :type src_fmt_configs: dict
- """
-
- project_id, dataset_id, external_table_id = \
- _split_tablename(table_input=external_project_dataset_table,
- default_project_id=self.project_id,
- var_name='external_project_dataset_table')
-
- # bigquery only allows certain source formats
- # we check to make sure the passed source format is valid
- # if it's not, we raise a ValueError
- # Refer to this link for more details:
- # https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#externalDataConfiguration.sourceFormat
-
- source_format = source_format.upper()
- allowed_formats = [
- "CSV", "NEWLINE_DELIMITED_JSON", "AVRO", "GOOGLE_SHEETS",
- "DATASTORE_BACKUP", "PARQUET"
- ]
- if source_format not in allowed_formats:
- raise ValueError("{0} is not a valid source format. "
- "Please use one of the following types: {1}"
- .format(source_format, allowed_formats))
-
- compression = compression.upper()
- allowed_compressions = ['NONE', 'GZIP']
- if compression not in allowed_compressions:
- raise ValueError("{0} is not a valid compression format. "
- "Please use one of the following types: {1}"
- .format(compression, allowed_compressions))
-
- table_resource = {
- 'externalDataConfiguration': {
- 'autodetect': autodetect,
- 'sourceFormat': source_format,
- 'sourceUris': source_uris,
- 'compression': compression,
- 'ignoreUnknownValues': ignore_unknown_values
- },
- 'tableReference': {
- 'projectId': project_id,
- 'datasetId': dataset_id,
- 'tableId': external_table_id,
- }
- }
-
- if schema_fields:
- table_resource['externalDataConfiguration'].update({
- 'schema': {
- 'fields': schema_fields
- }
- })
-
- self.log.info('Creating external table: %s', external_project_dataset_table)
-
- if max_bad_records:
- table_resource['externalDataConfiguration']['maxBadRecords'] = max_bad_records
-
- # if following fields are not specified in src_fmt_configs,
- # honor the top-level params for backward-compatibility
- if 'skipLeadingRows' not in src_fmt_configs:
- src_fmt_configs['skipLeadingRows'] = skip_leading_rows
- if 'fieldDelimiter' not in src_fmt_configs:
- src_fmt_configs['fieldDelimiter'] = field_delimiter
- if 'quote_character' not in src_fmt_configs:
- src_fmt_configs['quote'] = quote_character
- if 'allowQuotedNewlines' not in src_fmt_configs:
- src_fmt_configs['allowQuotedNewlines'] = allow_quoted_newlines
- if 'allowJaggedRows' not in src_fmt_configs:
- src_fmt_configs['allowJaggedRows'] = allow_jagged_rows
-
- src_fmt_to_param_mapping = {
- 'CSV': 'csvOptions',
- 'GOOGLE_SHEETS': 'googleSheetsOptions'
- }
-
- src_fmt_to_configs_mapping = {
- 'csvOptions': [
- 'allowJaggedRows', 'allowQuotedNewlines',
- 'fieldDelimiter', 'skipLeadingRows',
- 'quote'
- ],
- 'googleSheetsOptions': ['skipLeadingRows']
- }
-
- if source_format in src_fmt_to_param_mapping.keys():
-
- valid_configs = src_fmt_to_configs_mapping[
- src_fmt_to_param_mapping[source_format]
- ]
-
- src_fmt_configs = {
- k: v
- for k, v in src_fmt_configs.items() if k in valid_configs
- }
-
- table_resource['externalDataConfiguration'][src_fmt_to_param_mapping[
- source_format]] = src_fmt_configs
-
- try:
- self.service.tables().insert(
- projectId=project_id,
- datasetId=dataset_id,
- body=table_resource
- ).execute()
-
- self.log.info('External table created successfully: %s',
- external_project_dataset_table)
-
- except HttpError as err:
- raise Exception(
- 'BigQuery job failed. Error was: {}'.format(err.content)
- )
-
- def run_query(self,
- bql=None,
- sql=None,
- destination_dataset_table=False,
- write_disposition='WRITE_EMPTY',
- allow_large_results=False,
- flatten_results=False,
- udf_config=False,
- use_legacy_sql=None,
- maximum_billing_tier=None,
- maximum_bytes_billed=None,
- create_disposition='CREATE_IF_NEEDED',
- query_params=None,
- schema_update_options=(),
- priority='INTERACTIVE',
- time_partitioning={}):
- """
- Executes a BigQuery SQL query. Optionally persists results in a BigQuery
- table. See here:
-
- https://cloud.google.com/bigquery/docs/reference/v2/jobs
-
- For more details about these parameters.
-
- :param bql: (Deprecated. Use `sql` parameter instead) The BigQuery SQL
- to execute.
- :type bql: string
- :param sql: The BigQuery SQL to execute.
- :type sql: string
- :param destination_dataset_table: The dotted <dataset>.<table>
- BigQuery table to save the query results.
- :type destination_dataset_table: string
- :param write_disposition: What to do if the table already exists in
- BigQuery.
- :type write_disposition: string
- :param allow_large_results: Whether to allow large results.
- :type allow_large_results: boolean
- :param flatten_results: If true and query uses legacy SQL dialect, flattens
- all nested and repeated fields in the query results. ``allowLargeResults``
- must be true if this is set to false. For standard SQL queries, this
- flag is ignored and results are never flattened.
- :type flatten_results: boolean
- :param udf_config: The User Defined Function configuration for the query.
- See https://cloud.google.com/bigquery/user-defined-functions for details.
- :param use_legacy_sql: Whether to use legacy SQL (true) or standard SQL (false).
- If `None`, defaults to `self.use_legacy_sql`.
- :type use_legacy_sql: boolean
- :type udf_config: list
- :param maximum_billing_tier: Positive integer that serves as a
- multiplier of the basic price.
- :type maximum_billing_tier: integer
- :param maximum_bytes_billed: Limits the bytes billed for this job.
- Queries that will have bytes billed beyond this limit will fail
- (without incurring a charge). If unspecified, this will be
- set to your project default.
- :type maximum_bytes_billed: float
- :param create_disposition: Specifies whether the job is allowed to
- create new tables.
- :type create_disposition: string
- :param query_params a dictionary containing query parameter types and
- values, passed to BigQuery
- :type query_params: dict
- :param schema_update_options: Allows the schema of the desitination
- table to be updated as a side effect of the query job.
- :type schema_update_options: tuple
- :param priority: Specifies a priority for the query.
- Possible values include INTERACTIVE and BATCH.
- The default value is INTERACTIVE.
- :type priority: string
- :param time_partitioning: configure optional time partitioning fields i.e.
- partition by field, type and
- expiration as per API specifications. Note that 'field' is not available in
- conjunction with dataset.table$partition.
- :type time_partitioning: dict
-
- """
-
- # TODO remove `bql` in Airflow 2.0 - Jira: [AIRFLOW-2513]
- sql = bql if sql is None else sql
-
- if bql:
- import warnings
- warnings.warn('Deprecated parameter `bql` used in '
- '`BigQueryBaseCursor.run_query` '
- 'Use `sql` parameter instead to pass the sql to be '
- 'executed. `bql` parameter is deprecated and '
- 'will be removed in a future version of '
- 'Airflow.',
- category=DeprecationWarning)
-
- if sql is None:
- raise TypeError('`BigQueryBaseCursor.run_query` missing 1 required '
- 'positional argument: `sql`')
-
- # BigQuery also allows you to define how you want a table's schema to change
- # as a side effect of a query job
- # for more details:
- # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query.schemaUpdateOptions
- allowed_schema_update_options = [
- 'ALLOW_FIELD_ADDITION', "ALLOW_FIELD_RELAXATION"
- ]
- if not set(allowed_schema_update_options).issuperset(
- set(schema_update_options)):
- raise ValueError(
- "{0} contains invalid schema update options. "
- "Please only use one or more of the following options: {1}"
- .format(schema_update_options, allowed_schema_update_options))
-
- if use_legacy_sql is None:
- use_legacy_sql = self.use_legacy_sql
-
- configuration = {
- 'query': {
- 'query': sql,
- 'useLegacySql': use_legacy_sql,
- 'maximumBillingTier': maximum_billing_tier,
- 'maximumBytesBilled': maximum_bytes_billed,
- 'priority': priority
- }
- }
-
- if destination_dataset_table:
- assert '.' in destination_dataset_table, (
- 'Expected destination_dataset_table in the format of '
- '<dataset>.<table>. Got: {}').format(destination_dataset_table)
- destination_project, destination_dataset, destination_table = \
- _split_tablename(table_input=destination_dataset_table,
- default_project_id=self.project_id)
- configuration['query'].update({
- 'allowLargeResults': allow_large_results,
- 'flattenResults': flatten_results,
- 'writeDisposition': write_disposition,
- 'createDisposition': create_disposition,
- 'destinationTable': {
- 'projectId': destination_project,
- 'datasetId': destination_dataset,
- 'tableId': destination_table,
- }
- })
- if udf_config:
- assert isinstance(udf_config, list)
- configuration['query'].update({
- 'userDefinedFunctionResources': udf_config
- })
-
- if query_params:
- if self.use_legacy_sql:
- raise ValueError("Query paramaters are not allowed when using "
- "legacy SQL")
- else:
- configuration['query']['queryParameters'] = query_params
-
- time_partitioning = _cleanse_time_partitioning(
- destination_dataset_table,
- time_partitioning
- )
- if time_partitioning:
- configuration['query'].update({
- 'timePartitioning': time_partitioning
- })
-
- if schema_update_options:
- if write_disposition not in ["WRITE_APPEND", "WRITE_TRUNCATE"]:
- raise ValueError("schema_update_options is only "
- "allowed if write_disposition is "
- "'WRITE_APPEND' or 'WRITE_TRUNCATE'.")
- else:
- self.log.info(
- "Adding experimental "
- "'schemaUpdateOptions': {0}".format(schema_update_options))
- configuration['query'][
- 'schemaUpdateOptions'] = schema_update_options
-
- return self.run_with_configuration(configuration)
-
- def run_extract( # noqa
- self,
- source_project_dataset_table,
- destination_cloud_storage_uris,
- compression='NONE',
- export_format='CSV',
- field_delimiter=',',
- print_header=True):
- """
- Executes a BigQuery extract command to copy data from BigQuery to
- Google Cloud Storage. See here:
-
- https://cloud.google.com/bigquery/docs/reference/v2/jobs
-
- For more details about these parameters.
-
- :param source_project_dataset_table: The dotted <dataset>.<table>
- BigQuery table to use as the source data.
- :type source_project_dataset_table: string
- :param destination_cloud_storage_uris: The destination Google Cloud
- Storage URI (e.g. gs://some-bucket/some-file.txt). Follows
- convention defined here:
- https://cloud.google.com/bigquery/exporting-data-from-bigquery#exportingmultiple
- :type destination_cloud_storage_uris: list
- :param compression: Type of compression to use.
- :type compression: string
- :param export_format: File format to export.
- :type export_format: string
- :param field_delimiter: The delimiter to use when extracting to a CSV.
- :type field_delimiter: string
- :param print_header: Whether to print a header for a CSV file extract.
- :type print_header: boolean
- """
-
- source_project, source_dataset, source_table = \
- _split_tablename(table_input=source_project_dataset_table,
- default_project_id=self.project_id,
- var_name='source_project_dataset_table')
-
- configuration = {
- 'extract': {
- 'sourceTable': {
- 'projectId': source_project,
- 'datasetId': source_dataset,
- 'tableId': source_table,
- },
- 'compression': compression,
- 'destinationUris': destination_cloud_storage_uris,
- 'destinationFormat': export_format,
- }
- }
-
- if export_format == 'CSV':
- # Only set fieldDelimiter and printHeader fields if using CSV.
- # Google does not like it if you set these fields for other export
- # formats.
- configuration['extract']['fieldDelimiter'] = field_delimiter
- configuration['extract']['printHeader'] = print_header
-
- return self.run_with_configuration(configuration)
-
- def run_copy(self,
- source_project_dataset_tables,
- destination_project_dataset_table,
- write_disposition='WRITE_EMPTY',
- create_disposition='CREATE_IF_NEEDED'):
- """
- Executes a BigQuery copy command to copy data from one BigQuery table
- to another. See here:
-
- https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.copy
-
- For more details about these parameters.
-
- :param source_project_dataset_tables: One or more dotted
- (project:|project.)<dataset>.<table>
- BigQuery tables to use as the source data. Use a list if there are
- multiple source tables.
- If <project> is not included, project will be the project defined
- in the connection json.
- :type source_project_dataset_tables: list|string
- :param destination_project_dataset_table: The destination BigQuery
- table. Format is: (project:|project.)<dataset>.<table>
- :type destination_project_dataset_table: string
- :param write_disposition: The write disposition if the table already exists.
- :type write_disposition: string
- :param create_disposition: The create disposition if the table doesn't exist.
- :type create_disposition: string
- """
- source_project_dataset_tables = ([
- source_project_dataset_tables
- ] if not isinstance(source_project_dataset_tables, list) else
- source_project_dataset_tables)
-
- source_project_dataset_tables_fixup = []
- for source_project_dataset_table in source_project_dataset_tables:
- source_project, source_dataset, source_table = \
- _split_tablename(table_input=source_project_dataset_table,
- default_project_id=self.project_id,
- var_name='source_project_dataset_table')
- source_project_dataset_tables_fixup.append({
- 'projectId':
- source_project,
- 'datasetId':
- source_dataset,
- 'tableId':
- source_table
- })
-
- destination_project, destination_dataset, destination_table = \
- _split_tablename(table_input=destination_project_dataset_table,
- default_project_id=self.project_id)
- configuration = {
- 'copy': {
- 'createDisposition': create_disposition,
- 'writeDisposition': write_disposition,
- 'sourceTables': source_project_dataset_tables_fixup,
- 'destinationTable': {
- 'projectId': destination_project,
- 'datasetId': destination_dataset,
- 'tableId': destination_table
- }
- }
- }
-
- return self.run_with_configuration(configuration)
-
- def run_load(self,
- destination_project_dataset_table,
- schema_fields,
- source_uris,
- source_format='CSV',
- create_disposition='CREATE_IF_NEEDED',
- skip_leading_rows=0,
- write_disposition='WRITE_EMPTY',
- field_delimiter=',',
- max_bad_records=0,
- quote_character=None,
- ignore_unknown_values=False,
- allow_quoted_newlines=False,
- allow_jagged_rows=False,
- schema_update_options=(),
- src_fmt_configs={},
- time_partitioning={}):
- """
- Executes a BigQuery load command to load data from Google Cloud Storage
- to BigQuery. See here:
-
- https://cloud.google.com/bigquery/docs/reference/v2/jobs
-
- For more details about these parameters.
-
- :param destination_project_dataset_table:
- The dotted (<project>.|<project>:)<dataset>.<table>($<partition>) BigQuery
- table to load data into. If <project> is not included, project will be the
- project defined in the connection json. If a partition is specified the
- operator will automatically append the data, create a new partition or create
- a new DAY partitioned table.
- :type destination_project_dataset_table: string
- :param schema_fields: The schema field list as defined here:
- https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.load
- :type schema_fields: list
- :param source_uris: The source Google Cloud
- Storage URI (e.g. gs://some-bucket/some-file.txt). A single wild
- per-object name can be used.
- :type source_uris: list
- :param source_format: File format to export.
- :type source_format: string
- :param create_disposition: The create disposition if the table doesn't exist.
- :type create_disposition: string
- :param skip_leading_rows: Number of rows to skip when loading from a CSV.
- :type skip_leading_rows: int
- :param write_disposition: The write disposition if the table already exists.
- :type write_disposition: string
- :param field_delimiter: The delimiter to use when loading from a CSV.
- :type field_delimiter: string
- :param max_bad_records: The maximum number of bad records that BigQuery can
- ignore when running the job.
- :type max_bad_records: int
- :param quote_character: The value that is used to quote data sections in a CSV
- file.
- :type quote_character: string
- :param ignore_unknown_values: [Optional] Indicates if BigQuery should allow
- extra values that are not represented in the table schema.
- If true, the extra values are ignored. If false, records with extra columns
- are treated as bad records, and if there are too many bad records, an
- invalid error is returned in the job result.
- :type ignore_unknown_values: bool
- :param allow_quoted_newlines: Whether to allow quoted newlines (true) or not
- (false).
- :type allow_quoted_newlines: boolean
- :param allow_jagged_rows: Accept rows that are missing trailing optional columns.
- The missing values are treated as nulls. If false, records with missing
- trailing columns are treated as bad records, and if there are too many bad
- records, an invalid error is returned in the job result. Only applicable when
- soure_format is CSV.
- :type allow_jagged_rows: bool
- :param schema_update_options: Allows the schema of the desitination
- table to be updated as a side effect of the load job.
- :type schema_update_options: tuple
- :param src_fmt_configs: configure optional fields specific to the source format
- :type src_fmt_configs: dict
- :param time_partitioning: configure optional time partitioning fields i.e.
- partition by field, type and
- expiration as per API specifications. Note that 'field' is not available in
- conjunction with dataset.table$partition.
- :type time_partitioning: dict
- """
-
- # bigquery only allows certain source formats
- # we check to make sure the passed source format is valid
- # if it's not, we raise a ValueError
- # Refer to this link for more details:
- # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query.tableDefinitions.(key).sourceFormat
- source_format = source_format.upper()
- allowed_formats = [
- "CSV", "NEWLINE_DELIMITED_JSON", "AVRO", "GOOGLE_SHEETS",
- "DATASTORE_BACKUP", "PARQUET"
- ]
- if source_format not in allowed_formats:
- raise ValueError("{0} is not a valid source format. "
- "Please use one of the following types: {1}"
- .format(source_format, allowed_formats))
-
- # bigquery also allows you to define how you want a table's schema to change
- # as a side effect of a load
- # for more details:
- # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schemaUpdateOptions
- allowed_schema_update_options = [
- 'ALLOW_FIELD_ADDITION', "ALLOW_FIELD_RELAXATION"
- ]
- if not set(allowed_schema_update_options).issuperset(
- set(schema_update_options)):
- raise ValueError(
- "{0} contains invalid schema update options. "
- "Please only use one or more of the following options: {1}"
- .format(schema_update_options, allowed_schema_update_options))
-
- destination_project, destination_dataset, destination_table = \
- _split_tablename(table_input=destination_project_dataset_table,
- default_project_id=self.project_id,
- var_name='destination_project_dataset_table')
-
- configuration = {
- 'load': {
- 'createDisposition': create_disposition,
- 'destinationTable': {
- 'projectId': destination_project,
- 'datasetId': destination_dataset,
- 'tableId': destination_table,
- },
- 'sourceFormat': source_format,
- 'sourceUris': source_uris,
- 'writeDisposition': write_disposition,
- 'ignoreUnknownValues': ignore_unknown_values
- }
- }
-
- time_partitioning = _cleanse_time_partitioning(
- destination_project_dataset_table,
- time_partitioning
- )
- if time_partitioning:
- configuration['load'].update({
- 'timePartitioning': time_partitioning
- })
-
- if schema_fields:
- configuration['load']['schema'] = {'fields': schema_fields}
-
- if schema_update_options:
- if write_disposition not in ["WRITE_APPEND", "WRITE_TRUNCATE"]:
- raise ValueError("schema_update_options is only "
- "allowed if write_disposition is "
- "'WRITE_APPEND' or 'WRITE_TRUNCATE'.")
- else:
- self.log.info(
- "Adding experimental "
- "'schemaUpdateOptions': {0}".format(schema_update_options))
- configuration['load'][
- 'schemaUpdateOptions'] = schema_update_options
-
- if max_bad_records:
- configuration['load']['maxBadRecords'] = max_bad_records
-
- # if following fields are not specified in src_fmt_configs,
- # honor the top-level params for backward-compatibility
- if 'skipLeadingRows' not in src_fmt_configs:
- src_fmt_configs['skipLeadingRows'] = skip_leading_rows
- if 'fieldDelimiter' not in src_fmt_configs:
- src_fmt_configs['fieldDelimiter'] = field_delimiter
- if 'ignoreUnknownValues' not in src_fmt_configs:
- src_fmt_configs['ignoreUnknownValues'] = ignore_unknown_values
- if quote_character is not None:
- src_fmt_configs['quote'] = quote_character
- if allow_quoted_newlines:
- src_fmt_configs['allowQuotedNewlines'] = allow_quoted_newlines
-
- src_fmt_to_configs_mapping = {
- 'CSV': [
- 'allowJaggedRows', 'allowQuotedNewlines', 'autodetect',
- 'fieldDelimiter', 'skipLeadingRows', 'ignoreUnknownValues',
- 'nullMarker', 'quote'
- ],
- 'DATASTORE_BACKUP': ['projectionFields'],
- 'NEWLINE_DELIMITED_JSON': ['autodetect', 'ignoreUnknownValues'],
- 'PARQUET': ['autodetect', 'ignoreUnknownValues'],
- 'AVRO': [],
- }
- valid_configs = src_fmt_to_configs_mapping[source_format]
- src_fmt_configs = {
- k: v
- for k, v in src_fmt_configs.items() if k in valid_configs
- }
- configuration['load'].update(src_fmt_configs)
-
- if allow_jagged_rows:
- configuration['load']['allowJaggedRows'] = allow_jagged_rows
-
- return self.run_with_configuration(configuration)
-
- def run_with_configuration(self, configuration):
- """
- Executes a BigQuery SQL query. See here:
-
- https://cloud.google.com/bigquery/docs/reference/v2/jobs
-
- For more details about the configuration parameter.
-
- :param configuration: The configuration parameter maps directly to
- BigQuery's configuration field in the job object. See
- https://cloud.google.com/bigquery/docs/reference/v2/jobs for
- details.
- """
- jobs = self.service.jobs()
- job_data = {'configuration': configuration}
-
- # Send query and wait for reply.
- query_reply = jobs \
- .insert(projectId=self.project_id, body=job_data) \
- .execute()
- self.running_job_id = query_reply['jobReference']['jobId']
-
- # Wait for query to finish.
- keep_polling_job = True
- while (keep_polling_job):
- try:
- job = jobs.get(
- projectId=self.project_id,
- jobId=self.running_job_id).execute()
- if (job['status']['state'] == 'DONE'):
- keep_polling_job = False
- # Check if job had errors.
- if 'errorResult' in job['status']:
- raise Exception(
- 'BigQuery job failed. Final error was: {}. The job was: {}'.
- format(job['status']['errorResult'], job))
- else:
- self.log.info('Waiting for job to complete : %s, %s',
- self.project_id, self.running_job_id)
- time.sleep(5)
-
- except HttpError as err:
- if err.resp.status in [500, 503]:
- self.log.info(
- '%s: Retryable error, waiting for job to complete: %s',
- err.resp.status, self.running_job_id)
- time.sleep(5)
- else:
- raise Exception(
- 'BigQuery job status check failed. Final error was: %s',
- err.resp.status)
-
- return self.running_job_id
-
- def poll_job_complete(self, job_id):
- jobs = self.service.jobs()
- try:
- job = jobs.get(projectId=self.project_id, jobId=job_id).execute()
- if (job['status']['state'] == 'DONE'):
- return True
- except HttpError as err:
- if err.resp.status in [500, 503]:
- self.log.info(
- '%s: Retryable error while polling job with id %s',
- err.resp.status, job_id)
- else:
- raise Exception(
- 'BigQuery job status check failed. Final error was: %s',
- err.resp.status)
- return False
-
- def cancel_query(self):
- """
- Cancel all started queries that have not yet completed
- """
- jobs = self.service.jobs()
- if (self.running_job_id and
- not self.poll_job_complete(self.running_job_id)):
- self.log.info('Attempting to cancel job : %s, %s', self.project_id,
- self.running_job_id)
- jobs.cancel(
- projectId=self.project_id,
- jobId=self.running_job_id).execute()
- else:
- self.log.info('No running BigQuery jobs to cancel.')
- return
-
- # Wait for all the calls to cancel to finish
- max_polling_attempts = 12
- polling_attempts = 0
-
- job_complete = False
- while (polling_attempts < max_polling_attempts and not job_complete):
- polling_attempts = polling_attempts + 1
- job_complete = self.poll_job_complete(self.running_job_id)
- if (job_complete):
- self.log.info('Job successfully canceled: %s, %s',
- self.project_id, self.running_job_id)
- elif (polling_attempts == max_polling_attempts):
- self.log.info(
- "Stopping polling due to timeout. Job with id %s "
- "has not completed cancel and may or may not finish.",
- self.running_job_id)
- else:
- self.log.info('Waiting for canceled job with id %s to finish.',
- self.running_job_id)
- time.sleep(5)
-
- def get_schema(self, dataset_id, table_id):
- """
- Get the schema for a given datset.table.
- see https://cloud.google.com/bigquery/docs/reference/v2/tables#resource
-
- :param dataset_id: the dataset ID of the requested table
- :param table_id: the table ID of the requested table
- :return: a table schema
- """
- tables_resource = self.service.tables() \
- .get(projectId=self.project_id, datasetId=dataset_id, tableId=table_id) \
- .execute()
- return tables_resource['schema']
-
- def get_tabledata(self, dataset_id, table_id,
- max_results=None, selected_fields=None, page_token=None,
- start_index=None):
- """
- Get the data of a given dataset.table and optionally with selected columns.
- see https://cloud.google.com/bigquery/docs/reference/v2/tabledata/list
-
- :param dataset_id: the dataset ID of the requested table.
- :param table_id: the table ID of the requested table.
- :param max_results: the maximum results to return.
- :param selected_fields: List of fields to return (comma-separated). If
- unspecified, all fields are returned.
- :param page_token: page token, returned from a previous call,
- identifying the result set.
- :param start_index: zero based index of the starting row to read.
- :return: map containing the requested rows.
- """
- optional_params = {}
- if max_results:
- optional_params['maxResults'] = max_results
- if selected_fields:
- optional_params['selectedFields'] = selected_fields
- if page_token:
- optional_params['pageToken'] = page_token
- if start_index:
- optional_params['startIndex'] = start_index
- return (self.service.tabledata().list(
- projectId=self.project_id,
- datasetId=dataset_id,
- tableId=table_id,
- **optional_params).execute())
-
- def run_table_delete(self, deletion_dataset_table,
- ignore_if_missing=False):
- """
- Delete an existing table from the dataset;
- If the table does not exist, return an error unless ignore_if_missing
- is set to True.
-
- :param deletion_dataset_table: A dotted
- (<project>.|<project>:)<dataset>.<table> that indicates which table
- will be deleted.
- :type deletion_dataset_table: str
- :param ignore_if_missing: if True, then return success even if the
- requested table does not exist.
- :type ignore_if_missing: boolean
- :return:
- """
-
- assert '.' in deletion_dataset_table, (
- 'Expected deletion_dataset_table in the format of '
- '<dataset>.<table>. Got: {}').format(deletion_dataset_table)
- deletion_project, deletion_dataset, deletion_table = \
- _split_tablename(table_input=deletion_dataset_table,
- default_project_id=self.project_id)
-
- try:
- self.service.tables() \
- .delete(projectId=deletion_project,
- datasetId=deletion_dataset,
- tableId=deletion_table) \
- .execute()
- self.log.info('Deleted table %s:%s.%s.', deletion_project,
- deletion_dataset, deletion_table)
- except HttpError:
- if not ignore_if_missing:
- raise Exception('Table deletion failed. Table does not exist.')
- else:
- self.log.info('Table does not exist. Skipping.')
-
- def run_table_upsert(self, dataset_id, table_resource, project_id=None):
- """
- creates a new, empty table in the dataset;
- If the table already exists, update the existing table.
- Since BigQuery does not natively allow table upserts, this is not an
- atomic operation.
-
- :param dataset_id: the dataset to upsert the table into.
- :type dataset_id: str
- :param table_resource: a table resource. see
- https://cloud.google.com/bigquery/docs/reference/v2/tables#resource
- :type table_resource: dict
- :param project_id: the project to upsert the table into. If None,
- project will be self.project_id.
- :return:
- """
- # check to see if the table exists
- table_id = table_resource['tableReference']['tableId']
- project_id = project_id if project_id is not None else self.project_id
- tables_list_resp = self.service.tables().list(
- projectId=project_id, datasetId=dataset_id).execute()
- while True:
- for table in tables_list_resp.get('tables', []):
- if table['tableReference']['tableId'] == table_id:
- # found the table, do update
- self.log.info('Table %s:%s.%s exists, updating.',
- project_id, dataset_id, table_id)
- return self.service.tables().update(
- projectId=project_id,
- datasetId=dataset_id,
- tableId=table_id,
- body=table_resource).execute()
- # If there is a next page, we need to check the next page.
- if 'nextPageToken' in tables_list_resp:
- tables_list_resp = self.service.tables()\
- .list(projectId=project_id,
- datasetId=dataset_id,
- pageToken=tables_list_resp['nextPageToken'])\
- .execute()
- # If there is no next page, then the table doesn't exist.
- else:
- # do insert
- self.log.info('Table %s:%s.%s does not exist. creating.',
- project_id, dataset_id, table_id)
- return self.service.tables().insert(
- projectId=project_id,
- datasetId=dataset_id,
- body=table_resource).execute()
-
- def run_grant_dataset_view_access(self,
- source_dataset,
- view_dataset,
- view_table,
- source_project=None,
- view_project=None):
- """
- Grant authorized view access of a dataset to a view table.
- If this view has already been granted access to the dataset, do nothing.
- This method is not atomic. Running it may clobber a simultaneous update.
-
- :param source_dataset: the source dataset
- :type source_dataset: str
- :param view_dataset: the dataset that the view is in
- :type view_dataset: str
- :param view_table: the table of the view
- :type view_table: str
- :param source_project: the project of the source dataset. If None,
- self.project_id will be used.
- :type source_project: str
- :param view_project: the project that the view is in. If None,
- self.project_id will be used.
- :type view_project: str
- :return: the datasets resource of the source dataset.
- """
-
- # Apply default values to projects
- source_project = source_project if source_project else self.project_id
- view_project = view_project if view_project else self.project_id
-
- # we don't want to clobber any existing accesses, so we have to get
- # info on the dataset before we can add view access
- source_dataset_resource = self.service.datasets().get(
- projectId=source_project, datasetId=source_dataset).execute()
- access = source_dataset_resource[
- 'access'] if 'access' in source_dataset_resource else []
- view_access = {
- 'view': {
- 'projectId': view_project,
- 'datasetId': view_dataset,
- 'tableId': view_table
- }
- }
- # check to see if the view we want to add already exists.
- if view_access not in access:
- self.log.info(
- 'Granting table %s:%s.%s authorized view access to %s:%s dataset.',
- view_project, view_dataset, view_table, source_project,
- source_dataset)
- access.append(view_access)
- return self.service.datasets().patch(
- projectId=source_project,
- datasetId=source_dataset,
- body={
- 'access': access
- }).execute()
- else:
- # if view is already in access, do nothing.
- self.log.info(
- 'Table %s:%s.%s already has authorized view access to %s:%s dataset.',
- view_project, view_dataset, view_table, source_project,
- source_dataset)
- return source_dataset_resource
-
-
-class BigQueryCursor(BigQueryBaseCursor):
- """
- A very basic BigQuery PEP 249 cursor implementation. The PyHive PEP 249
- implementation was used as a reference:
-
- https://github.com/dropbox/PyHive/blob/master/pyhive/presto.py
- https://github.com/dropbox/PyHive/blob/master/pyhive/common.py
- """
-
- def __init__(self, service, project_id, use_legacy_sql=True):
- super(BigQueryCursor, self).__init__(
- service=service,
- project_id=project_id,
- use_legacy_sql=use_legacy_sql)
- self.buffersize = None
- self.page_token = None
- self.job_id = None
- self.buffer = []
- self.all_pages_loaded = False
-
- @property
- def description(self):
- """ The schema description method is not currently implemented. """
- raise NotImplementedError
-
- def close(self):
- """ By default, do nothing """
- pass
-
- @property
- def rowcount(self):
- """ By default, return -1 to indicate that this is not supported. """
- return -1
-
- def execute(self, operation, parameters=None):
- """
- Executes a BigQuery query, and returns the job ID.
-
- :param operation: The query to execute.
- :type operation: string
- :param parameters: Parameters to substitute into the query.
- :type parameters: dict
- """
- sql = _bind_parameters(operation,
- parameters) if parameters else operation
- self.job_id = self.run_query(sql)
-
- def executemany(self, operation, seq_of_parameters):
- """
- Execute a BigQuery query multiple times with different parameters.
-
- :param operation: The query to execute.
- :type operation: string
- :param seq_of_parameters: List of dictionary parameters to substitute into the
- query.
- :type seq_of_parameters: list
- """
- for parameters in seq_of_parameters:
- self.execute(operation, parameters)
-
- def fetchone(self):
- """ Fetch the next row of a query result set. """
- return self.next()
-
- def next(self):
- """
- Helper method for fetchone, which returns the next row from a buffer.
- If the buffer is empty, attempts to paginate through the result set for
- the next page, and load it into the buffer.
- """
- if not self.job_id:
- return None
-
- if len(self.buffer) == 0:
- if self.all_pages_loaded:
- return None
-
- query_results = (self.service.jobs().getQueryResults(
- projectId=self.project_id,
- jobId=self.job_id,
- pageToken=self.page_token).execute())
-
- if 'rows' in query_results and query_results['rows']:
- self.page_token = query_results.get('pageToken')
- fields = query_results['schema']['fields']
- col_types = [field['type'] for field in fields]
- rows = query_results['rows']
-
- for dict_row in rows:
- typed_row = ([
- _bq_cast(vs['v'], col_types[idx])
- for idx, vs in enumerate(dict_row['f'])
- ])
- self.buffer.append(typed_row)
-
- if not self.page_token:
- self.all_pages_loaded = True
-
- else:
- # Reset all state since we've exhausted the results.
- self.page_token = None
- self.job_id = None
- self.page_token = None
- return None
-
- return self.buffer.pop(0)
-
- def fetchmany(self, size=None):
- """
- Fetch the next set of rows of a query result, returning a sequence of sequences
- (e.g. a list of tuples). An empty sequence is returned when no more rows are
- available. The number of rows to fetch per call is specified by the parameter.
- If it is not given, the cursor's arraysize determines the number of rows to be
- fetched. The method should try to fetch as many rows as indicated by the size
- parameter. If this is not possible due to the specified number of rows not being
- available, fewer rows may be returned. An :py:class:`~pyhive.exc.Error`
- (or subclass) exception is raised if the previous call to
- :py:meth:`execute` did not produce any result set or no call was issued yet.
- """
- if size is None:
- size = self.arraysize
- result = []
- for _ in range(size):
- one = self.fetchone()
- if one is None:
- break
- else:
- result.append(one)
- return result
-
- def fetchall(self):
- """
- Fetch all (remaining) rows of a query result, returning them as a sequence of
- sequences (e.g. a list of tuples).
- """
- result = []
- while True:
- one = self.fetchone()
- if one is None:
- break
- else:
- result.append(one)
- return result
-
- def get_arraysize(self):
- """ Specifies the number of rows to fetch at a time with .fetchmany() """
- return self._buffersize if self.buffersize else 1
-
- def set_arraysize(self, arraysize):
- """ Specifies the number of rows to fetch at a time with .fetchmany() """
- self.buffersize = arraysize
-
- arraysize = property(get_arraysize, set_arraysize)
-
- def setinputsizes(self, sizes):
- """ Does nothing by default """
- pass
-
- def setoutputsize(self, size, column=None):
- """ Does nothing by default """
- pass
-
-
-def _bind_parameters(operation, parameters):
- """ Helper method that binds parameters to a SQL query. """
- # inspired by MySQL Python Connector (conversion.py)
- string_parameters = {}
- for (name, value) in parameters.iteritems():
- if value is None:
- string_parameters[name] = 'NULL'
- elif isinstance(value, basestring):
- string_parameters[name] = "'" + _escape(value) + "'"
- else:
- string_parameters[name] = str(value)
- return operation % string_parameters
-
-
-def _escape(s):
- """ Helper method that escapes parameters to a SQL query. """
- e = s
- e = e.replace('\\', '\\\\')
- e = e.replace('\n', '\\n')
- e = e.replace('\r', '\\r')
- e = e.replace("'", "\\'")
- e = e.replace('"', '\\"')
- return e
-
-
-def _bq_cast(string_field, bq_type):
- """
- Helper method that casts a BigQuery row to the appropriate data types.
- This is useful because BigQuery returns all fields as strings.
- """
- if string_field is None:
- return None
- elif bq_type == 'INTEGER':
- return int(string_field)
- elif bq_type == 'FLOAT' or bq_type == 'TIMESTAMP':
- return float(string_field)
- elif bq_type == 'BOOLEAN':
- assert string_field in set(['true', 'false'])
- return string_field == 'true'
- else:
- return string_field
-
-
-def _split_tablename(table_input, default_project_id, var_name=None):
- assert default_project_id is not None, "INTERNAL: No default project is specified"
-
- def var_print(var_name):
- if var_name is None:
- return ""
- else:
- return "Format exception for {var}: ".format(var=var_name)
-
- if table_input.count('.') + table_input.count(':') > 3:
- raise Exception(('{var}Use either : or . to specify project '
- 'got {input}').format(
- var=var_print(var_name), input=table_input))
-
- cmpt = table_input.rsplit(':', 1)
- project_id = None
- rest = table_input
- if len(cmpt) == 1:
- project_id = None
- rest = cmpt[0]
- elif len(cmpt) == 2 and cmpt[0].count(':') <= 1:
- if cmpt[-1].count('.') != 2:
- project_id = cmpt[0]
- rest = cmpt[1]
- else:
- raise Exception(('{var}Expect format of (<project:)<dataset>.<table>, '
- 'got {input}').format(
- var=var_print(var_name), input=table_input))
-
- cmpt = rest.split('.')
- if len(cmpt) == 3:
- assert project_id is None, ("{var}Use either : or . to specify project"
- ).format(var=var_print(var_name))
- project_id = cmpt[0]
- dataset_id = cmpt[1]
- table_id = cmpt[2]
-
- elif len(cmpt) == 2:
- dataset_id = cmpt[0]
- table_id = cmpt[1]
- else:
- raise Exception(
- ('{var}Expect format of (<project.|<project:)<dataset>.<table>, '
- 'got {input}').format(var=var_print(var_name), input=table_input))
-
- if project_id is None:
- if var_name is not None:
- log = LoggingMixin().log
- log.info('Project not included in {var}: {input}; '
- 'using project "{project}"'.format(
- var=var_name,
- input=table_input,
- project=default_project_id))
- project_id = default_project_id
-
- return project_id, dataset_id, table_id
-
-
-def _cleanse_time_partitioning(destination_dataset_table, time_partitioning_in):
- # if it is a partitioned table ($ is in the table name) add partition load option
- time_partitioning_out = {}
- if destination_dataset_table and '$' in destination_dataset_table:
- assert not time_partitioning_in.get('field'), (
- "Cannot specify field partition and partition name "
- "(dataset.table$partition) at the same time"
- )
- time_partitioning_out['type'] = 'DAY'
-
- time_partitioning_out.update(time_partitioning_in)
- return time_partitioning_out
diff --git a/airflow/contrib/hooks/cassandra_hook.py b/airflow/contrib/hooks/cassandra_hook.py
deleted file mode 100644
index 0e0b47708d..0000000000
--- a/airflow/contrib/hooks/cassandra_hook.py
+++ /dev/null
@@ -1,197 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 cassandra.cluster import Cluster
-from cassandra.policies import (RoundRobinPolicy, DCAwareRoundRobinPolicy,
- TokenAwarePolicy, WhiteListRoundRobinPolicy)
-from cassandra.auth import PlainTextAuthProvider
-
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class CassandraHook(BaseHook, LoggingMixin):
- """
- Hook used to interact with Cassandra
-
- Contact points can be specified as a comma-separated string in the 'hosts'
- field of the connection.
-
- Port can be specified in the port field of the connection.
-
- If SSL is enabled in Cassandra, pass in a dict in the extra field as kwargs for
- ``ssl.wrap_socket()``. For example:
- {
- 'ssl_options' : {
- 'ca_certs' : PATH_TO_CA_CERTS
- }
- }
-
- Default load balancing policy is RoundRobinPolicy. To specify a different LB policy:
- - DCAwareRoundRobinPolicy
- {
- 'load_balancing_policy': 'DCAwareRoundRobinPolicy',
- 'load_balancing_policy_args': {
- 'local_dc': LOCAL_DC_NAME, // optional
- 'used_hosts_per_remote_dc': SOME_INT_VALUE, // optional
- }
- }
- - WhiteListRoundRobinPolicy
- {
- 'load_balancing_policy': 'WhiteListRoundRobinPolicy',
- 'load_balancing_policy_args': {
- 'hosts': ['HOST1', 'HOST2', 'HOST3']
- }
- }
- - TokenAwarePolicy
- {
- 'load_balancing_policy': 'TokenAwarePolicy',
- 'load_balancing_policy_args': {
- 'child_load_balancing_policy': CHILD_POLICY_NAME, // optional
- 'child_load_balancing_policy_args': { ... } // optional
- }
- }
-
- For details of the Cluster config, see cassandra.cluster.
- """
- def __init__(self, cassandra_conn_id='cassandra_default'):
- conn = self.get_connection(cassandra_conn_id)
-
- conn_config = {}
- if conn.host:
- conn_config['contact_points'] = conn.host.split(',')
-
- if conn.port:
- conn_config['port'] = int(conn.port)
-
- if conn.login:
- conn_config['auth_provider'] = PlainTextAuthProvider(
- username=conn.login, password=conn.password)
-
- policy_name = conn.extra_dejson.get('load_balancing_policy', None)
- policy_args = conn.extra_dejson.get('load_balancing_policy_args', {})
- lb_policy = self.get_lb_policy(policy_name, policy_args)
- if lb_policy:
- conn_config['load_balancing_policy'] = lb_policy
-
- cql_version = conn.extra_dejson.get('cql_version', None)
- if cql_version:
- conn_config['cql_version'] = cql_version
-
- ssl_options = conn.extra_dejson.get('ssl_options', None)
- if ssl_options:
- conn_config['ssl_options'] = ssl_options
-
- self.cluster = Cluster(**conn_config)
- self.keyspace = conn.schema
- self.session = None
-
- def get_conn(self):
- """
- Returns a cassandra Session object
- """
- if self.session and not self.session.is_shutdown:
- return self.session
- self.session = self.cluster.connect(self.keyspace)
- return self.session
-
- def get_cluster(self):
- return self.cluster
-
- def shutdown_cluster(self):
- """
- Closes all sessions and connections associated with this Cluster.
- """
- if not self.cluster.is_shutdown:
- self.cluster.shutdown()
-
- @staticmethod
- def get_lb_policy(policy_name, policy_args):
- policies = {
- 'RoundRobinPolicy': RoundRobinPolicy,
- 'DCAwareRoundRobinPolicy': DCAwareRoundRobinPolicy,
- 'WhiteListRoundRobinPolicy': WhiteListRoundRobinPolicy,
- 'TokenAwarePolicy': TokenAwarePolicy,
- }
-
- if not policies.get(policy_name) or policy_name == 'RoundRobinPolicy':
- return RoundRobinPolicy()
-
- if policy_name == 'DCAwareRoundRobinPolicy':
- local_dc = policy_args.get('local_dc', '')
- used_hosts_per_remote_dc = int(policy_args.get('used_hosts_per_remote_dc', 0))
- return DCAwareRoundRobinPolicy(local_dc, used_hosts_per_remote_dc)
-
- if policy_name == 'WhiteListRoundRobinPolicy':
- hosts = policy_args.get('hosts')
- if not hosts:
- raise Exception('Hosts must be specified for WhiteListRoundRobinPolicy')
- return WhiteListRoundRobinPolicy(hosts)
-
- if policy_name == 'TokenAwarePolicy':
- allowed_child_policies = ('RoundRobinPolicy',
- 'DCAwareRoundRobinPolicy',
- 'WhiteListRoundRobinPolicy',)
- child_policy_name = policy_args.get('child_load_balancing_policy',
- 'RoundRobinPolicy')
- child_policy_args = policy_args.get('child_load_balancing_policy_args', {})
- if child_policy_name not in allowed_child_policies:
- return TokenAwarePolicy(RoundRobinPolicy())
- else:
- child_policy = CassandraHook.get_lb_policy(child_policy_name,
- child_policy_args)
- return TokenAwarePolicy(child_policy)
-
- def table_exists(self, table):
- """
- Checks if a table exists in Cassandra
-
- :param table: Target Cassandra table.
- Use dot notation to target a specific keyspace.
- :type table: string
- """
- keyspace = self.keyspace
- if '.' in table:
- keyspace, table = table.split('.', 1)
- cluster_metadata = self.get_conn().cluster.metadata
- return (keyspace in cluster_metadata.keyspaces and
- table in cluster_metadata.keyspaces[keyspace].tables)
-
- def record_exists(self, table, keys):
- """
- Checks if a record exists in Cassandra
-
- :param table: Target Cassandra table.
- Use dot notation to target a specific keyspace.
- :type table: string
- :param keys: The keys and their values to check the existence.
- :type keys: dict
- """
- keyspace = self.keyspace
- if '.' in table:
- keyspace, table = table.split('.', 1)
- ks = " AND ".join("{}=%({})s".format(key, key) for key in keys.keys())
- cql = "SELECT * FROM {keyspace}.{table} WHERE {keys}".format(
- keyspace=keyspace, table=table, keys=ks)
-
- try:
- rs = self.get_conn().execute(cql, keys)
- return rs.one() is not None
- except Exception:
- return False
diff --git a/airflow/contrib/hooks/cloudant_hook.py b/airflow/contrib/hooks/cloudant_hook.py
deleted file mode 100644
index 5d39f3fa8a..0000000000
--- a/airflow/contrib/hooks/cloudant_hook.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 past.builtins import unicode
-
-import cloudant
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class CloudantHook(BaseHook):
- """Interact with Cloudant.
-
- This class is a thin wrapper around the cloudant python library. See the
- documentation `here <https://github.com/cloudant-labs/cloudant-python>`_.
- """
- def __init__(self, cloudant_conn_id='cloudant_default'):
- super(CloudantHook, self).__init__('cloudant')
- self.cloudant_conn_id = cloudant_conn_id
-
- def get_conn(self):
- def _str(s):
- # cloudant-python doesn't support unicode.
- if isinstance(s, unicode):
- log = LoggingMixin().log
- log.debug(
- 'cloudant-python does not support unicode. Encoding %s as '
- 'ascii using "ignore".', s
- )
- return s.encode('ascii', 'ignore')
-
- return s
-
- conn = self.get_connection(self.cloudant_conn_id)
-
- for conn_param in ['host', 'password', 'schema']:
- if not hasattr(conn, conn_param) or not getattr(conn, conn_param):
- raise AirflowException(
- 'missing connection parameter {0}'.format(conn_param)
- )
-
- # In the connection form:
- # - 'host' is renamed to 'Account'
- # - 'login' is renamed 'Username (or API Key)'
- # - 'schema' is renamed to 'Database'
- #
- # So, use the 'host' attribute as the account name, and, if login is
- # defined, use that as the username.
- account = cloudant.Account(_str(conn.host))
-
- username = _str(conn.login or conn.host)
-
- account.login(
- username,
- _str(conn.password)).raise_for_status()
-
- return account.database(_str(conn.schema))
-
- def db(self):
- """Returns the Database object for this hook.
-
- See the documentation for cloudant-python here
- https://github.com/cloudant-labs/cloudant-python.
- """
- return self.get_conn()
diff --git a/airflow/contrib/hooks/databricks_hook.py b/airflow/contrib/hooks/databricks_hook.py
deleted file mode 100644
index 1443ff4740..0000000000
--- a/airflow/contrib/hooks/databricks_hook.py
+++ /dev/null
@@ -1,228 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 requests
-
-from airflow import __version__
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from requests import exceptions as requests_exceptions
-from requests.auth import AuthBase
-
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-try:
- from urllib import parse as urlparse
-except ImportError:
- import urlparse
-
-
-SUBMIT_RUN_ENDPOINT = ('POST', 'api/2.0/jobs/runs/submit')
-GET_RUN_ENDPOINT = ('GET', 'api/2.0/jobs/runs/get')
-CANCEL_RUN_ENDPOINT = ('POST', 'api/2.0/jobs/runs/cancel')
-USER_AGENT_HEADER = {'user-agent': 'airflow-{v}'.format(v=__version__)}
-
-
-class DatabricksHook(BaseHook, LoggingMixin):
- """
- Interact with Databricks.
- """
- def __init__(
- self,
- databricks_conn_id='databricks_default',
- timeout_seconds=180,
- retry_limit=3):
- """
- :param databricks_conn_id: The name of the databricks connection to use.
- :type databricks_conn_id: string
- :param timeout_seconds: The amount of time in seconds the requests library
- will wait before timing-out.
- :type timeout_seconds: int
- :param retry_limit: The number of times to retry the connection in case of
- service outages.
- :type retry_limit: int
- """
- self.databricks_conn_id = databricks_conn_id
- self.databricks_conn = self.get_connection(databricks_conn_id)
- self.timeout_seconds = timeout_seconds
- assert retry_limit >= 1, 'Retry limit must be greater than equal to 1'
- self.retry_limit = retry_limit
-
- def _parse_host(self, host):
- """
- The purpose of this function is to be robust to improper connections
- settings provided by users, specifically in the host field.
-
-
- For example -- when users supply ``https://xx.cloud.databricks.com`` as the
- host, we must strip out the protocol to get the host.
- >>> h = DatabricksHook()
- >>> assert h._parse_host('https://xx.cloud.databricks.com') == \
- 'xx.cloud.databricks.com'
-
- In the case where users supply the correct ``xx.cloud.databricks.com`` as the
- host, this function is a no-op.
- >>> assert h._parse_host('xx.cloud.databricks.com') == 'xx.cloud.databricks.com'
- """
- urlparse_host = urlparse.urlparse(host).hostname
- if urlparse_host:
- # In this case, host = https://xx.cloud.databricks.com
- return urlparse_host
- else:
- # In this case, host = xx.cloud.databricks.com
- return host
-
- def _do_api_call(self, endpoint_info, json):
- """
- Utility function to perform an API call with retries
- :param endpoint_info: Tuple of method and endpoint
- :type endpoint_info: (string, string)
- :param json: Parameters for this API call.
- :type json: dict
- :return: If the api call returns a OK status code,
- this function returns the response in JSON. Otherwise,
- we throw an AirflowException.
- :rtype: dict
- """
- method, endpoint = endpoint_info
- url = 'https://{host}/{endpoint}'.format(
- host=self._parse_host(self.databricks_conn.host),
- endpoint=endpoint)
- if 'token' in self.databricks_conn.extra_dejson:
- self.log.info('Using token auth.')
- auth = _TokenAuth(self.databricks_conn.extra_dejson['token'])
- else:
- self.log.info('Using basic auth.')
- auth = (self.databricks_conn.login, self.databricks_conn.password)
- if method == 'GET':
- request_func = requests.get
- elif method == 'POST':
- request_func = requests.post
- else:
- raise AirflowException('Unexpected HTTP Method: ' + method)
-
- for attempt_num in range(1, self.retry_limit + 1):
- try:
- response = request_func(
- url,
- json=json,
- auth=auth,
- headers=USER_AGENT_HEADER,
- timeout=self.timeout_seconds)
- if response.status_code == requests.codes.ok:
- return response.json()
- else:
- # In this case, the user probably made a mistake.
- # Don't retry.
- raise AirflowException('Response: {0}, Status Code: {1}'.format(
- response.content, response.status_code))
- except (requests_exceptions.ConnectionError,
- requests_exceptions.Timeout) as e:
- self.log.error(
- 'Attempt %s API Request to Databricks failed with reason: %s',
- attempt_num, e
- )
- raise AirflowException(('API requests to Databricks failed {} times. ' +
- 'Giving up.').format(self.retry_limit))
-
- def submit_run(self, json):
- """
- Utility function to call the ``api/2.0/jobs/runs/submit`` endpoint.
-
- :param json: The data used in the body of the request to the ``submit`` endpoint.
- :type json: dict
- :return: the run_id as a string
- :rtype: string
- """
- response = self._do_api_call(SUBMIT_RUN_ENDPOINT, json)
- return response['run_id']
-
- def get_run_page_url(self, run_id):
- json = {'run_id': run_id}
- response = self._do_api_call(GET_RUN_ENDPOINT, json)
- return response['run_page_url']
-
- def get_run_state(self, run_id):
- json = {'run_id': run_id}
- response = self._do_api_call(GET_RUN_ENDPOINT, json)
- state = response['state']
- life_cycle_state = state['life_cycle_state']
- # result_state may not be in the state if not terminal
- result_state = state.get('result_state', None)
- state_message = state['state_message']
- return RunState(life_cycle_state, result_state, state_message)
-
- def cancel_run(self, run_id):
- json = {'run_id': run_id}
- self._do_api_call(CANCEL_RUN_ENDPOINT, json)
-
-
-RUN_LIFE_CYCLE_STATES = [
- 'PENDING',
- 'RUNNING',
- 'TERMINATING',
- 'TERMINATED',
- 'SKIPPED',
- 'INTERNAL_ERROR'
-]
-
-
-class RunState:
- """
- Utility class for the run state concept of Databricks runs.
- """
- def __init__(self, life_cycle_state, result_state, state_message):
- self.life_cycle_state = life_cycle_state
- self.result_state = result_state
- self.state_message = state_message
-
- @property
- def is_terminal(self):
- if self.life_cycle_state not in RUN_LIFE_CYCLE_STATES:
- raise AirflowException(
- ('Unexpected life cycle state: {}: If the state has '
- 'been introduced recently, please check the Databricks user '
- 'guide for troubleshooting information').format(
- self.life_cycle_state))
- return self.life_cycle_state in ('TERMINATED', 'SKIPPED', 'INTERNAL_ERROR')
-
- @property
- def is_successful(self):
- return self.result_state == 'SUCCESS'
-
- def __eq__(self, other):
- return self.life_cycle_state == other.life_cycle_state and \
- self.result_state == other.result_state and \
- self.state_message == other.state_message
-
- def __repr__(self):
- return str(self.__dict__)
-
-
-class _TokenAuth(AuthBase):
- """
- Helper class for requests Auth field. AuthBase requires you to implement the __call__
- magic function.
- """
- def __init__(self, token):
- self.token = token
-
- def __call__(self, r):
- r.headers['Authorization'] = 'Bearer ' + self.token
- return r
diff --git a/airflow/contrib/hooks/datadog_hook.py b/airflow/contrib/hooks/datadog_hook.py
deleted file mode 100644
index 3dfeb781ae..0000000000
--- a/airflow/contrib/hooks/datadog_hook.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 time
-from airflow.hooks.base_hook import BaseHook
-from airflow.exceptions import AirflowException
-from datadog import initialize, api
-
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class DatadogHook(BaseHook, LoggingMixin):
- """
- Uses datadog API to send metrics of practically anything measurable,
- so it's possible to track # of db records inserted/deleted, records read
- from file and many other useful metrics.
-
- Depends on the datadog API, which has to be deployed on the same server where
- Airflow runs.
-
- :param datadog_conn_id: The connection to datadog, containing metadata for api keys.
- :param datadog_conn_id: string
- """
- def __init__(self, datadog_conn_id='datadog_default'):
- conn = self.get_connection(datadog_conn_id)
- self.api_key = conn.extra_dejson.get('api_key', None)
- self.app_key = conn.extra_dejson.get('app_key', None)
- self.source_type_name = conn.extra_dejson.get('source_type_name', None)
-
- # If the host is populated, it will use that hostname instead.
- # for all metric submissions.
- self.host = conn.host
-
- if self.api_key is None:
- raise AirflowException("api_key must be specified in the "
- "Datadog connection details")
- if self.app_key is None:
- raise AirflowException("app_key must be specified in the "
- "Datadog connection details")
-
- self.log.info("Setting up api keys for Datadog")
- options = {
- 'api_key': self.api_key,
- 'app_key': self.app_key
- }
- initialize(**options)
-
- def validate_response(self, response):
- if response['status'] != 'ok':
- self.log.error("Datadog returned: %s", response)
- raise AirflowException("Error status received from Datadog")
-
- def send_metric(self, metric_name, datapoint, tags=None):
- """
- Sends a single datapoint metric to DataDog
-
- :param metric_name: The name of the metric
- :type metric_name: string
- :param datapoint: A single integer or float related to the metric
- :type datapoint: integer or float
- :param tags: A list of tags associated with the metric
- :type tags: list
- """
- response = api.Metric.send(
- metric=metric_name,
- points=datapoint,
- host=self.host,
- tags=tags)
-
- self.validate_response(response)
- return response
-
- def query_metric(self,
- query,
- from_seconds_ago,
- to_seconds_ago):
- """
- Queries datadog for a specific metric, potentially with some
- function applied to it and returns the results.
-
- :param query: The datadog query to execute (see datadog docs)
- :type query: string
- :param from_seconds_ago: How many seconds ago to start querying for.
- :type from_seconds_ago: int
- :param to_seconds_ago: Up to how many seconds ago to query for.
- :type to_seconds_ago: int
- """
- now = int(time.time())
-
- response = api.Metric.query(
- start=now - from_seconds_ago,
- end=now - to_seconds_ago,
- query=query)
-
- self.validate_response(response)
- return response
-
- def post_event(self, title, text, tags=None, alert_type=None, aggregation_key=None):
- """
- Posts an event to datadog (processing finished, potentially alerts, other issues)
- Think about this as a means to maintain persistence of alerts, rather than
- alerting itself.
-
- :param title: The title of the event
- :type title: string
- :param text: The body of the event (more information)
- :type text: string
- :param tags: List of string tags to apply to the event
- :type tags: list
- :param alert_type: The alert type for the event, one of
- ["error", "warning", "info", "success"]
- :type alert_type: string
- :param aggregation_key: Key that can be used to aggregate this event in a stream
- :type aggregation_key: string
- """
- response = api.Event.create(
- title=title,
- text=text,
- host=self.host,
- tags=tags,
- alert_type=alert_type,
- aggregation_key=aggregation_key,
- source_type_name=self.source_type_name)
-
- self.validate_response(response)
- return response
diff --git a/airflow/contrib/hooks/datastore_hook.py b/airflow/contrib/hooks/datastore_hook.py
deleted file mode 100644
index 5e54cf2a65..0000000000
--- a/airflow/contrib/hooks/datastore_hook.py
+++ /dev/null
@@ -1,206 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 time
-from apiclient.discovery import build
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-
-
-class DatastoreHook(GoogleCloudBaseHook):
- """
- Interact with Google Cloud Datastore. This hook uses the Google Cloud Platform
- connection.
-
- This object is not threads safe. If you want to make multiple requests
- simultaneously, you will need to create a hook per thread.
- """
-
- def __init__(self,
- datastore_conn_id='google_cloud_datastore_default',
- delegate_to=None):
- super(DatastoreHook, self).__init__(datastore_conn_id, delegate_to)
- self.connection = self.get_conn()
- self.admin_connection = self.get_conn('v1beta1')
-
- def get_conn(self, version='v1'):
- """
- Returns a Google Cloud Storage service object.
- """
- http_authorized = self._authorize()
- return build(
- 'datastore', version, http=http_authorized, cache_discovery=False)
-
- def allocate_ids(self, partialKeys):
- """
- Allocate IDs for incomplete keys.
- see https://cloud.google.com/datastore/docs/reference/rest/v1/projects/allocateIds
-
- :param partialKeys: a list of partial keys
- :return: a list of full keys.
- """
- resp = self.connection.projects().allocateIds(
- projectId=self.project_id, body={'keys': partialKeys}
- ).execute()
- return resp['keys']
-
- def begin_transaction(self):
- """
- Get a new transaction handle
-
- .. seealso::
- https://cloud.google.com/datastore/docs/reference/rest/v1/projects/beginTransaction
-
- :return: a transaction handle
- """
- resp = self.connection.projects().beginTransaction(
- projectId=self.project_id, body={}).execute()
- return resp['transaction']
-
- def commit(self, body):
- """
- Commit a transaction, optionally creating, deleting or modifying some entities.
-
- .. seealso::
- https://cloud.google.com/datastore/docs/reference/rest/v1/projects/commit
-
- :param body: the body of the commit request
- :return: the response body of the commit request
- """
- resp = self.connection.projects().commit(
- projectId=self.project_id, body=body).execute()
- return resp
-
- def lookup(self, keys, read_consistency=None, transaction=None):
- """
- Lookup some entities by key
-
- .. seealso::
- https://cloud.google.com/datastore/docs/reference/rest/v1/projects/lookup
-
- :param keys: the keys to lookup
- :param read_consistency: the read consistency to use. default, strong or eventual.
- Cannot be used with a transaction.
- :param transaction: the transaction to use, if any.
- :return: the response body of the lookup request.
- """
- body = {'keys': keys}
- if read_consistency:
- body['readConsistency'] = read_consistency
- if transaction:
- body['transaction'] = transaction
- return self.connection.projects().lookup(
- projectId=self.project_id, body=body).execute()
-
- def rollback(self, transaction):
- """
- Roll back a transaction
-
- .. seealso::
- https://cloud.google.com/datastore/docs/reference/rest/v1/projects/rollback
-
- :param transaction: the transaction to roll back
- """
- self.connection.projects().rollback(
- projectId=self.project_id, body={'transaction': transaction})\
- .execute()
-
- def run_query(self, body):
- """
- Run a query for entities.
-
- .. seealso::
- https://cloud.google.com/datastore/docs/reference/rest/v1/projects/runQuery
-
- :param body: the body of the query request
- :return: the batch of query results.
- """
- resp = self.connection.projects().runQuery(
- projectId=self.project_id, body=body).execute()
- return resp['batch']
-
- def get_operation(self, name):
- """
- Gets the latest state of a long-running operation
-
- :param name: the name of the operation resource
- """
- resp = self.connection.projects().operations().get(name=name).execute()
- return resp
-
- def delete_operation(self, name):
- """
- Deletes the long-running operation
-
- :param name: the name of the operation resource
- """
- resp = self.connection.projects().operations().delete(name=name).execute()
- return resp
-
- def poll_operation_until_done(self, name, polling_interval_in_seconds):
- """
- Poll backup operation state until it's completed
- """
- while True:
- result = self.get_operation(name)
- state = result['metadata']['common']['state']
- if state == 'PROCESSING':
- self.log.info('Operation is processing. Re-polling state in {} seconds'
- .format(polling_interval_in_seconds))
- time.sleep(polling_interval_in_seconds)
- else:
- return result
-
- def export_to_storage_bucket(self, bucket, namespace=None,
- entity_filter=None, labels=None):
- """
- Export entities from Cloud Datastore to Cloud Storage for backup
- """
- output_uri_prefix = 'gs://' + ('/').join(filter(None, [bucket, namespace]))
- if not entity_filter:
- entity_filter = {}
- if not labels:
- labels = {}
- body = {
- 'outputUrlPrefix': output_uri_prefix,
- 'entityFilter': entity_filter,
- 'labels': labels,
- }
- resp = self.admin_connection.projects().export(
- projectId=self.project_id, body=body).execute()
- return resp
-
- def import_from_storage_bucket(self, bucket, file,
- namespace=None, entity_filter=None, labels=None):
- """
- Import a backup from Cloud Storage to Cloud Datastore
- """
- input_url = 'gs://' + ('/').join(filter(None, [bucket, namespace, file]))
- if not entity_filter:
- entity_filter = {}
- if not labels:
- labels = {}
- body = {
- 'inputUrl': input_url,
- 'entityFilter': entity_filter,
- 'labels': labels,
- }
- resp = self.admin_connection.projects().import_(
- projectId=self.project_id, body=body).execute()
- return resp
diff --git a/airflow/contrib/hooks/discord_webhook_hook.py b/airflow/contrib/hooks/discord_webhook_hook.py
deleted file mode 100644
index 731d9d575a..0000000000
--- a/airflow/contrib/hooks/discord_webhook_hook.py
+++ /dev/null
@@ -1,140 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-import re
-
-from airflow.hooks.http_hook import HttpHook
-from airflow.exceptions import AirflowException
-
-
-class DiscordWebhookHook(HttpHook):
- """
- This hook allows you to post messages to Discord using incoming webhooks.
- Takes a Discord connection ID with a default relative webhook endpoint. The
- default endpoint can be overridden using the webhook_endpoint parameter
- (https://discordapp.com/developers/docs/resources/webhook).
-
- Each Discord webhook can be pre-configured to use a specific username and
- avatar_url. You can override these defaults in this hook.
-
- :param http_conn_id: Http connection ID with host as "https://discord.com/api/" and
- default webhook endpoint in the extra field in the form of
- {"webhook_endpoint": "webhooks/{webhook.id}/{webhook.token}"}
- :type http_conn_id: str
- :param webhook_endpoint: Discord webhook endpoint in the form of
- "webhooks/{webhook.id}/{webhook.token}"
- :type webhook_endpoint: str
- :param message: The message you want to send to your Discord channel
- (max 2000 characters)
- :type message: str
- :param username: Override the default username of the webhook
- :type username: str
- :param avatar_url: Override the default avatar of the webhook
- :type avatar_url: str
- :param tts: Is a text-to-speech message
- :type tts: bool
- :param proxy: Proxy to use to make the Discord webhook call
- :type proxy: str
- """
-
- def __init__(self,
- http_conn_id=None,
- webhook_endpoint=None,
- message="",
- username=None,
- avatar_url=None,
- tts=False,
- proxy=None,
- *args,
- **kwargs):
- super(DiscordWebhookHook, self).__init__(*args, **kwargs)
- self.http_conn_id = http_conn_id
- self.webhook_endpoint = self._get_webhook_endpoint(http_conn_id, webhook_endpoint)
- self.message = message
- self.username = username
- self.avatar_url = avatar_url
- self.tts = tts
- self.proxy = proxy
-
- def _get_webhook_endpoint(self, http_conn_id, webhook_endpoint):
- """
- Given a Discord http_conn_id, return the default webhook endpoint or override if a
- webhook_endpoint is manually supplied.
-
- :param http_conn_id: The provided connection ID
- :param webhook_endpoint: The manually provided webhook endpoint
- :return: Webhook endpoint (str) to use
- """
- if webhook_endpoint:
- endpoint = webhook_endpoint
- elif http_conn_id:
- conn = self.get_connection(http_conn_id)
- extra = conn.extra_dejson
- endpoint = extra.get('webhook_endpoint', '')
- else:
- raise AirflowException('Cannot get webhook endpoint: No valid Discord '
- 'webhook endpoint or http_conn_id supplied.')
-
- # make sure endpoint matches the expected Discord webhook format
- if not re.match('^webhooks/[0-9]+/[a-zA-Z0-9_-]+$', endpoint):
- raise AirflowException('Expected Discord webhook endpoint in the form '
- 'of "webhooks/{webhook.id}/{webhook.token}".')
-
- return endpoint
-
- def _build_discord_payload(self):
- """
- Construct the Discord JSON payload. All relevant parameters are combined here
- to a valid Discord JSON payload.
-
- :return: Discord payload (str) to send
- """
- payload = {}
-
- if self.username:
- payload['username'] = self.username
- if self.avatar_url:
- payload['avatar_url'] = self.avatar_url
-
- payload['tts'] = self.tts
-
- if len(self.message) <= 2000:
- payload['content'] = self.message
- else:
- raise AirflowException('Discord message length must be 2000 or fewer '
- 'characters.')
-
- return json.dumps(payload)
-
- def execute(self):
- """
- Execute the Discord webhook call
- """
- proxies = {}
- if self.proxy:
- # we only need https proxy for Discord
- proxies = {'https': self.proxy}
-
- discord_payload = self._build_discord_payload()
-
- self.run(endpoint=self.webhook_endpoint,
- data=discord_payload,
- headers={'Content-type': 'application/json'},
- extra_options={'proxies': proxies})
diff --git a/airflow/contrib/hooks/emr_hook.py b/airflow/contrib/hooks/emr_hook.py
deleted file mode 100644
index 6cd92c6d85..0000000000
--- a/airflow/contrib/hooks/emr_hook.py
+++ /dev/null
@@ -1,69 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.exceptions import AirflowException
-from airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class EmrHook(AwsHook):
- """
- Interact with AWS EMR. emr_conn_id is only neccessary for using the
- create_job_flow method.
- """
-
- def __init__(self, emr_conn_id=None, *args, **kwargs):
- self.emr_conn_id = emr_conn_id
- super(EmrHook, self).__init__(*args, **kwargs)
-
- def get_conn(self):
- self.conn = self.get_client_type('emr')
- return self.conn
-
- def create_job_flow(self, job_flow_overrides):
- """
- Creates a job flow using the config from the EMR connection.
- Keys of the json extra hash may have the arguments of the boto3
- run_job_flow method.
- Overrides for this config may be passed as the job_flow_overrides.
- """
-
- if not self.emr_conn_id:
- raise AirflowException('emr_conn_id must be present to use create_job_flow')
-
- emr_conn = self.get_connection(self.emr_conn_id)
-
- config = emr_conn.extra_dejson.copy()
- config.update(job_flow_overrides)
-
- response = self.get_conn().run_job_flow(
- Name=config.get('Name'),
- LogUri=config.get('LogUri'),
- ReleaseLabel=config.get('ReleaseLabel'),
- Instances=config.get('Instances'),
- Steps=config.get('Steps', []),
- BootstrapActions=config.get('BootstrapActions', []),
- Applications=config.get('Applications'),
- Configurations=config.get('Configurations', []),
- VisibleToAllUsers=config.get('VisibleToAllUsers'),
- JobFlowRole=config.get('JobFlowRole'),
- ServiceRole=config.get('ServiceRole'),
- Tags=config.get('Tags')
- )
-
- return response
diff --git a/airflow/contrib/hooks/fs_hook.py b/airflow/contrib/hooks/fs_hook.py
deleted file mode 100644
index 6832f20c22..0000000000
--- a/airflow/contrib/hooks/fs_hook.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.hooks.base_hook import BaseHook
-
-
-class FSHook(BaseHook):
- """
- Allows for interaction with an file server.
-
- Connection should have a name and a path specified under extra:
-
- example:
- Conn Id: fs_test
- Conn Type: File (path)
- Host, Shchema, Login, Password, Port: empty
- Extra: {"path": "/tmp"}
- """
-
- def __init__(self, conn_id='fs_default'):
- conn = self.get_connection(conn_id)
- self.basepath = conn.extra_dejson.get('path', '')
- self.conn = conn
-
- def get_conn(self):
- pass
-
- def get_path(self):
- return self.basepath
diff --git a/airflow/contrib/hooks/ftp_hook.py b/airflow/contrib/hooks/ftp_hook.py
deleted file mode 100644
index 8beefb3729..0000000000
--- a/airflow/contrib/hooks/ftp_hook.py
+++ /dev/null
@@ -1,259 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 datetime
-import ftplib
-import os.path
-from airflow.hooks.base_hook import BaseHook
-from past.builtins import basestring
-
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-def mlsd(conn, path="", facts=None):
- """
- BACKPORT FROM PYTHON3 FTPLIB.
-
- List a directory in a standardized format by using MLSD
- command (RFC-3659). If path is omitted the current directory
- is assumed. "facts" is a list of strings representing the type
- of information desired (e.g. ["type", "size", "perm"]).
-
- Return a generator object yielding a tuple of two elements
- for every file found in path.
- First element is the file name, the second one is a dictionary
- including a variable number of "facts" depending on the server
- and whether "facts" argument has been provided.
- """
- facts = facts or []
- if facts:
- conn.sendcmd("OPTS MLST " + ";".join(facts) + ";")
- if path:
- cmd = "MLSD %s" % path
- else:
- cmd = "MLSD"
- lines = []
- conn.retrlines(cmd, lines.append)
- for line in lines:
- facts_found, _, name = line.rstrip(ftplib.CRLF).partition(' ')
- entry = {}
- for fact in facts_found[:-1].split(";"):
- key, _, value = fact.partition("=")
- entry[key.lower()] = value
- yield (name, entry)
-
-
-class FTPHook(BaseHook, LoggingMixin):
- """
- Interact with FTP.
-
- Errors that may occur throughout but should be handled
- downstream.
- """
-
- def __init__(self, ftp_conn_id='ftp_default'):
- self.ftp_conn_id = ftp_conn_id
- self.conn = None
-
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- if self.conn is not None:
- self.close_conn()
-
- def get_conn(self):
- """
- Returns a FTP connection object
- """
- if self.conn is None:
- params = self.get_connection(self.ftp_conn_id)
- self.conn = ftplib.FTP(params.host, params.login, params.password)
-
- return self.conn
-
- def close_conn(self):
- """
- Closes the connection. An error will occur if the
- connection wasn't ever opened.
- """
- conn = self.conn
- conn.quit()
- self.conn = None
-
- def describe_directory(self, path):
- """
- Returns a dictionary of {filename: {attributes}} for all files
- on the remote system (where the MLSD command is supported).
-
- :param path: full path to the remote directory
- :type path: str
- """
- conn = self.get_conn()
- conn.cwd(path)
- try:
- # only works in Python 3
- files = dict(conn.mlsd())
- except AttributeError:
- files = dict(mlsd(conn))
- return files
-
- def list_directory(self, path, nlst=False):
- """
- Returns a list of files on the remote system.
-
- :param path: full path to the remote directory to list
- :type path: str
- """
- conn = self.get_conn()
- conn.cwd(path)
-
- files = conn.nlst()
- return files
-
- def create_directory(self, path):
- """
- Creates a directory on the remote system.
-
- :param path: full path to the remote directory to create
- :type path: str
- """
- conn = self.get_conn()
- conn.mkd(path)
-
- def delete_directory(self, path):
- """
- Deletes a directory on the remote system.
-
- :param path: full path to the remote directory to delete
- :type path: str
- """
- conn = self.get_conn()
- conn.rmd(path)
-
- def retrieve_file(self, remote_full_path, local_full_path_or_buffer):
- """
- Transfers the remote file to a local location.
-
- If local_full_path_or_buffer is a string path, the file will be put
- at that location; if it is a file-like buffer, the file will
- be written to the buffer but not closed.
-
- :param remote_full_path: full path to the remote file
- :type remote_full_path: str
- :param local_full_path_or_buffer: full path to the local file or a
- file-like buffer
- :type local_full_path_or_buffer: str or file-like buffer
- """
- conn = self.get_conn()
-
- is_path = isinstance(local_full_path_or_buffer, basestring)
-
- if is_path:
- output_handle = open(local_full_path_or_buffer, 'wb')
- else:
- output_handle = local_full_path_or_buffer
-
- remote_path, remote_file_name = os.path.split(remote_full_path)
- conn.cwd(remote_path)
- self.log.info('Retrieving file from FTP: %s', remote_full_path)
- conn.retrbinary('RETR %s' % remote_file_name, output_handle.write)
- self.log.info('Finished retrieving file from FTP: %s', remote_full_path)
-
- if is_path:
- output_handle.close()
-
- def store_file(self, remote_full_path, local_full_path_or_buffer):
- """
- Transfers a local file to the remote location.
-
- If local_full_path_or_buffer is a string path, the file will be read
- from that location; if it is a file-like buffer, the file will
- be read from the buffer but not closed.
-
- :param remote_full_path: full path to the remote file
- :type remote_full_path: str
- :param local_full_path_or_buffer: full path to the local file or a
- file-like buffer
- :type local_full_path_or_buffer: str or file-like buffer
- """
- conn = self.get_conn()
-
- is_path = isinstance(local_full_path_or_buffer, basestring)
-
- if is_path:
- input_handle = open(local_full_path_or_buffer, 'rb')
- else:
- input_handle = local_full_path_or_buffer
- remote_path, remote_file_name = os.path.split(remote_full_path)
- conn.cwd(remote_path)
- conn.storbinary('STOR %s' % remote_file_name, input_handle)
-
- if is_path:
- input_handle.close()
-
- def delete_file(self, path):
- """
- Removes a file on the FTP Server.
-
- :param path: full path to the remote file
- :type path: str
- """
- conn = self.get_conn()
- conn.delete(path)
-
- def rename(self, from_name, to_name):
- """
- Rename a file.
-
- :param from_name: rename file from name
- :param to_name: rename file to name
- """
- conn = self.get_conn()
- return conn.rename(from_name, to_name)
-
- def get_mod_time(self, path):
- conn = self.get_conn()
- ftp_mdtm = conn.sendcmd('MDTM ' + path)
- time_val = ftp_mdtm[4:]
- # time_val optionally has microseconds
- try:
- return datetime.datetime.strptime(time_val, "%Y%m%d%H%M%S.%f")
- except ValueError:
- return datetime.datetime.strptime(time_val, '%Y%m%d%H%M%S')
-
-
-class FTPSHook(FTPHook):
-
- def get_conn(self):
- """
- Returns a FTPS connection object.
- """
- if self.conn is None:
- params = self.get_connection(self.ftp_conn_id)
-
- if params.port:
- ftplib.FTP_TLS.port = params.port
-
- self.conn = ftplib.FTP_TLS(
- params.host, params.login, params.password
- )
-
- return self.conn
diff --git a/airflow/contrib/hooks/gcp_api_base_hook.py b/airflow/contrib/hooks/gcp_api_base_hook.py
deleted file mode 100644
index 053494743f..0000000000
--- a/airflow/contrib/hooks/gcp_api_base_hook.py
+++ /dev/null
@@ -1,152 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-
-import httplib2
-import google.auth
-import google_auth_httplib2
-import google.oauth2.service_account
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-_DEFAULT_SCOPES = ('https://www.googleapis.com/auth/cloud-platform',)
-
-
-class GoogleCloudBaseHook(BaseHook, LoggingMixin):
- """
- A base hook for Google cloud-related hooks. Google cloud has a shared REST
- API client that is built in the same way no matter which service you use.
- This class helps construct and authorize the credentials needed to then
- call apiclient.discovery.build() to actually discover and build a client
- for a Google cloud service.
-
- The class also contains some miscellaneous helper functions.
-
- All hook derived from this base hook use the 'Google Cloud Platform' connection
- type. Two ways of authentication are supported:
-
- Default credentials: Only the 'Project Id' is required. You'll need to
- have set up default credentials, such as by the
- ``GOOGLE_APPLICATION_DEFAULT`` environment variable or from the metadata
- server on Google Compute Engine.
-
- JSON key file: Specify 'Project Id', 'Key Path' and 'Scope'.
-
- Legacy P12 key files are not supported.
- """
-
- def __init__(self, gcp_conn_id='google_cloud_default', delegate_to=None):
- """
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- """
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.extras = self.get_connection(self.gcp_conn_id).extra_dejson
-
- def _get_credentials(self):
- """
- Returns the Credentials object for Google API
- """
- key_path = self._get_field('key_path', False)
- keyfile_dict = self._get_field('keyfile_dict', False)
- scope = self._get_field('scope', None)
- if scope is not None:
- scopes = [s.strip() for s in scope.split(',')]
- else:
- scopes = _DEFAULT_SCOPES
-
- if not key_path and not keyfile_dict:
- self.log.info('Getting connection using `google.auth.default()` '
- 'since no key file is defined for hook.')
- credentials, _ = google.auth.default(scopes=scopes)
- elif key_path:
- # Get credentials from a JSON file.
- if key_path.endswith('.json'):
- self.log.info('Getting connection using a JSON key file.')
- credentials = (
- google.oauth2.service_account.Credentials.from_service_account_file(
- key_path, scopes=scopes)
- )
- elif key_path.endswith('.p12'):
- raise AirflowException('Legacy P12 key file are not supported, '
- 'use a JSON key file.')
- else:
- raise AirflowException('Unrecognised extension for key file.')
- else:
- # Get credentials from JSON data provided in the UI.
- try:
- keyfile_dict = json.loads(keyfile_dict)
-
- # Depending on how the JSON was formatted, it may contain
- # escaped newlines. Convert those to actual newlines.
- keyfile_dict['private_key'] = keyfile_dict['private_key'].replace(
- '\\n', '\n')
-
- credentials = (
- google.oauth2.service_account.Credentials.from_service_account_info(
- keyfile_dict, scopes=scopes)
- )
- except json.decoder.JSONDecodeError:
- raise AirflowException('Invalid key JSON.')
-
- return credentials.with_subject(self.delegate_to) \
- if self.delegate_to else credentials
-
- def _get_access_token(self):
- """
- Returns a valid access token from Google API Credentials
- """
- return self._get_credentials().token
-
- def _authorize(self):
- """
- Returns an authorized HTTP object to be used to build a Google cloud
- service hook connection.
- """
- credentials = self._get_credentials()
- http = httplib2.Http()
- authed_http = google_auth_httplib2.AuthorizedHttp(
- credentials, http=http)
- return authed_http
-
- def _get_field(self, f, default=None):
- """
- Fetches a field from extras, and returns it. This is some Airflow
- magic. The google_cloud_platform hook type adds custom UI elements
- to the hook page, which allow admins to specify service_account,
- key_path, etc. They get formatted as shown below.
- """
- long_f = 'extra__google_cloud_platform__{}'.format(f)
- if long_f in self.extras:
- return self.extras[long_f]
- else:
- return default
-
- @property
- def project_id(self):
- return self._get_field('project')
diff --git a/airflow/contrib/hooks/gcp_container_hook.py b/airflow/contrib/hooks/gcp_container_hook.py
deleted file mode 100644
index d36d796d76..0000000000
--- a/airflow/contrib/hooks/gcp_container_hook.py
+++ /dev/null
@@ -1,217 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-import time
-
-from airflow import AirflowException, version
-from airflow.hooks.base_hook import BaseHook
-
-from google.api_core.exceptions import AlreadyExists
-from google.api_core.gapic_v1.method import DEFAULT
-from google.cloud import container_v1, exceptions
-from google.cloud.container_v1.gapic.enums import Operation
-from google.cloud.container_v1.types import Cluster
-from google.protobuf import json_format
-from google.api_core.gapic_v1.client_info import ClientInfo
-
-OPERATIONAL_POLL_INTERVAL = 15
-
-
-class GKEClusterHook(BaseHook):
-
- def __init__(self, project_id, location):
- self.project_id = project_id
- self.location = location
-
- # Add client library info for better error tracking
- client_info = ClientInfo(client_library_version='airflow_v' + version.version)
- self.client = container_v1.ClusterManagerClient(client_info=client_info)
-
- def _dict_to_proto(self, py_dict, proto):
- """
- Converts a python dictionary to the proto supplied
- :param py_dict: The dictionary to convert
- :type py_dict: dict
- :param proto: The proto object to merge with dictionary
- :type proto: protobuf
- :return: A parsed python dictionary in provided proto format
- :raises:
- ParseError: On JSON parsing problems.
- """
- dict_json_str = json.dumps(py_dict)
- return json_format.Parse(dict_json_str, proto)
-
- def wait_for_operation(self, operation):
- """
- Given an operation, continuously fetches the status from Google Cloud until either
- completion or an error occurring
- :param operation: The Operation to wait for
- :type operation: A google.cloud.container_V1.gapic.enums.Operator
- :return: A new, updated operation fetched from Google Cloud
- """
- self.log.info("Waiting for OPERATION_NAME %s" % operation.name)
- time.sleep(OPERATIONAL_POLL_INTERVAL)
- while operation.status != Operation.Status.DONE:
- if operation.status == Operation.Status.RUNNING or operation.status == \
- Operation.Status.PENDING:
- time.sleep(OPERATIONAL_POLL_INTERVAL)
- else:
- raise exceptions.GoogleCloudError(
- "Operation has failed with status: %s" % operation.status)
- # To update status of operation
- operation = self.get_operation(operation.name)
- return operation
-
- def get_operation(self, operation_name):
- """
- Fetches the operation from Google Cloud
- :param operation_name: Name of operation to fetch
- :type operation_name: str
- :return: The new, updated operation from Google Cloud
- """
- return self.client.get_operation(project_id=self.project_id,
- zone=self.location,
- operation_id=operation_name)
-
- def _append_label(self, cluster_proto, key, val):
- """
- Append labels to provided Cluster Protobuf
-
- Labels must fit the regex [a-z]([-a-z0-9]*[a-z0-9])? (current airflow version
- string follows semantic versioning spec: x.y.z).
- :param cluster_proto: The proto to append resource_label airflow version to
- :type cluster_proto: google.cloud.container_v1.types.Cluster
- :param key: The key label
- :type key: str
- :param val:
- :type val: str
- :return: The cluster proto updated with new label
- """
- val = val.replace('.', '-').replace('+', '-')
- cluster_proto.resource_labels.update({key: val})
- return cluster_proto
-
- def delete_cluster(self, name, retry=DEFAULT, timeout=DEFAULT):
- """
- Deletes the cluster, including the Kubernetes endpoint and all
- worker nodes. Firewalls and routes that were configured during
- cluster creation are also deleted. Other Google Compute Engine
- resources that might be in use by the cluster (e.g. load balancer
- resources) will not be deleted if they weren’t present at the
- initial create time.
-
- :param name: The name of the cluster to delete
- :type name: str
- :param retry: Retry object used to determine when/if to retry requests.
- If None is specified, requests will not be retried.
- :type retry: google.api_core.retry.Retry
- :param timeout: The amount of time, in seconds, to wait for the request to
- complete. Note that if retry is specified, the timeout applies to each
- individual attempt.
- :type timeout: float
- :return: The full url to the delete operation if successful, else None
- """
-
- self.log.info("Deleting (project_id={}, zone={}, cluster_id={})".format(
- self.project_id, self.location, name))
-
- try:
- op = self.client.delete_cluster(project_id=self.project_id,
- zone=self.location,
- cluster_id=name,
- retry=retry,
- timeout=timeout)
- op = self.wait_for_operation(op)
- # Returns server-defined url for the resource
- return op.self_link
- except exceptions.NotFound as error:
- self.log.info('Assuming Success: ' + error.message)
-
- def create_cluster(self, cluster, retry=DEFAULT, timeout=DEFAULT):
- """
- Creates a cluster, consisting of the specified number and type of Google Compute
- Engine instances.
-
- :param cluster: A Cluster protobuf or dict. If dict is provided, it must be of
- the same form as the protobuf message google.cloud.container_v1.types.Cluster
- :type cluster: dict or google.cloud.container_v1.types.Cluster
- :param retry: A retry object (google.api_core.retry.Retry) used to retry requests.
- If None is specified, requests will not be retried.
- :type retry: google.api_core.retry.Retry
- :param timeout: The amount of time, in seconds, to wait for the request to
- complete. Note that if retry is specified, the timeout applies to each
- individual attempt.
- :type timeout: float
- :return: The full url to the new, or existing, cluster
- :raises
- ParseError: On JSON parsing problems when trying to convert dict
- AirflowException: cluster is not dict type nor Cluster proto type
- """
-
- if isinstance(cluster, dict):
- cluster_proto = Cluster()
- cluster = self._dict_to_proto(py_dict=cluster, proto=cluster_proto)
- elif not isinstance(cluster, Cluster):
- raise AirflowException(
- "cluster is not instance of Cluster proto or python dict")
-
- self._append_label(cluster, 'airflow-version', 'v' + version.version)
-
- self.log.info("Creating (project_id={}, zone={}, cluster_name={})".format(
- self.project_id,
- self.location,
- cluster.name))
- try:
- op = self.client.create_cluster(project_id=self.project_id,
- zone=self.location,
- cluster=cluster,
- retry=retry,
- timeout=timeout)
- op = self.wait_for_operation(op)
-
- return op.target_link
- except AlreadyExists as error:
- self.log.info('Assuming Success: ' + error.message)
- return self.get_cluster(name=cluster.name).self_link
-
- def get_cluster(self, name, retry=DEFAULT, timeout=DEFAULT):
- """
- Gets details of specified cluster
- :param name: The name of the cluster to retrieve
- :type name: str
- :param retry: A retry object used to retry requests. If None is specified,
- requests will not be retried.
- :type retry: google.api_core.retry.Retry
- :param timeout: The amount of time, in seconds, to wait for the request to
- complete. Note that if retry is specified, the timeout applies to each
- individual attempt.
- :type timeout: float
- :return: A google.cloud.container_v1.types.Cluster instance
- """
- self.log.info("Fetching cluster (project_id={}, zone={}, cluster_name={})".format(
- self.project_id,
- self.location,
- name))
-
- return self.client.get_cluster(project_id=self.project_id,
- zone=self.location,
- cluster_id=name,
- retry=retry,
- timeout=timeout).self_link
diff --git a/airflow/contrib/hooks/gcp_dataflow_hook.py b/airflow/contrib/hooks/gcp_dataflow_hook.py
deleted file mode 100644
index 7abb413a1a..0000000000
--- a/airflow/contrib/hooks/gcp_dataflow_hook.py
+++ /dev/null
@@ -1,263 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-import select
-import subprocess
-import time
-import uuid
-
-from apiclient.discovery import build
-
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-# This is the default location
-# https://cloud.google.com/dataflow/pipelines/specifying-exec-params
-DEFAULT_DATAFLOW_LOCATION = 'us-central1'
-
-
-class _DataflowJob(LoggingMixin):
- def __init__(self, dataflow, project_number, name, location, poll_sleep=10):
- self._dataflow = dataflow
- self._project_number = project_number
- self._job_name = name
- self._job_location = location
- self._job_id = None
- self._job = self._get_job()
- self._poll_sleep = poll_sleep
-
- def _get_job_id_from_name(self):
- jobs = self._dataflow.projects().locations().jobs().list(
- projectId=self._project_number,
- location=self._job_location
- ).execute()
- for job in jobs['jobs']:
- if job['name'] == self._job_name:
- self._job_id = job['id']
- return job
- return None
-
- def _get_job(self):
- if self._job_name:
- job = self._get_job_id_from_name()
- else:
- job = self._dataflow.projects().jobs().get(
- projectId=self._project_number,
- jobId=self._job_id
- ).execute()
-
- if job and 'currentState' in job:
- self.log.info(
- 'Google Cloud DataFlow job %s is %s',
- job['name'], job['currentState']
- )
- elif job:
- self.log.info(
- 'Google Cloud DataFlow with job_id %s has name %s',
- self._job_id, job['name']
- )
- else:
- self.log.info(
- 'Google Cloud DataFlow job not available yet..'
- )
-
- return job
-
- def wait_for_done(self):
- while True:
- if self._job and 'currentState' in self._job:
- if 'JOB_STATE_DONE' == self._job['currentState']:
- return True
- elif 'JOB_STATE_RUNNING' == self._job['currentState'] and \
- 'JOB_TYPE_STREAMING' == self._job['type']:
- return True
- elif 'JOB_STATE_FAILED' == self._job['currentState']:
- raise Exception("Google Cloud Dataflow job {} has failed.".format(
- self._job['name']))
- elif 'JOB_STATE_CANCELLED' == self._job['currentState']:
- raise Exception("Google Cloud Dataflow job {} was cancelled.".format(
- self._job['name']))
- elif 'JOB_STATE_RUNNING' == self._job['currentState']:
- time.sleep(self._poll_sleep)
- elif 'JOB_STATE_PENDING' == self._job['currentState']:
- time.sleep(15)
- else:
- self.log.debug(str(self._job))
- raise Exception(
- "Google Cloud Dataflow job {} was unknown state: {}".format(
- self._job['name'], self._job['currentState']))
- else:
- time.sleep(15)
-
- self._job = self._get_job()
-
- def get(self):
- return self._job
-
-
-class _Dataflow(LoggingMixin):
- def __init__(self, cmd):
- self.log.info("Running command: %s", ' '.join(cmd))
- self._proc = subprocess.Popen(
- cmd,
- shell=False,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- close_fds=True)
-
- def _line(self, fd):
- if fd == self._proc.stderr.fileno():
- lines = self._proc.stderr.readlines()
- for line in lines:
- self.log.warning(line[:-1])
- if lines:
- return lines[-1]
- if fd == self._proc.stdout.fileno():
- line = self._proc.stdout.readline()
- return line
-
- @staticmethod
- def _extract_job(line):
- if line is not None:
- if line.startswith("Submitted job: "):
- return line[15:-1]
-
- def wait_for_done(self):
- reads = [self._proc.stderr.fileno(), self._proc.stdout.fileno()]
- self.log.info("Start waiting for DataFlow process to complete.")
- while self._proc.poll() is None:
- ret = select.select(reads, [], [], 5)
- if ret is not None:
- for fd in ret[0]:
- line = self._line(fd)
- if line:
- self.log.debug(line[:-1])
- else:
- self.log.info("Waiting for DataFlow process to complete.")
- if self._proc.returncode is not 0:
- raise Exception("DataFlow failed with return code {}".format(
- self._proc.returncode))
-
-
-class DataFlowHook(GoogleCloudBaseHook):
-
- def __init__(self,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- poll_sleep=10):
- self.poll_sleep = poll_sleep
- super(DataFlowHook, self).__init__(gcp_conn_id, delegate_to)
-
- def get_conn(self):
- """
- Returns a Google Cloud Storage service object.
- """
- http_authorized = self._authorize()
- return build(
- 'dataflow', 'v1b3', http=http_authorized, cache_discovery=False)
-
- def _start_dataflow(self, task_id, variables, name,
- command_prefix, label_formatter):
- variables = self._set_variables(variables)
- cmd = command_prefix + self._build_cmd(task_id, variables,
- label_formatter)
- _Dataflow(cmd).wait_for_done()
- _DataflowJob(self.get_conn(), variables['project'], name,
- variables['region'], self.poll_sleep).wait_for_done()
-
- @staticmethod
- def _set_variables(variables):
- if variables['project'] is None:
- raise Exception('Project not specified')
- if 'region' not in variables.keys():
- variables['region'] = DEFAULT_DATAFLOW_LOCATION
- return variables
-
- def start_java_dataflow(self, task_id, variables, dataflow, job_class=None,
- append_job_name=True):
- if append_job_name:
- name = task_id + "-" + str(uuid.uuid1())[:8]
- else:
- name = task_id
- variables['jobName'] = name
-
- def label_formatter(labels_dict):
- return ['--labels={}'.format(
- json.dumps(labels_dict).replace(' ', ''))]
- command_prefix = (["java", "-cp", dataflow, job_class] if job_class
- else ["java", "-jar", dataflow])
- self._start_dataflow(task_id, variables, name,
- command_prefix, label_formatter)
-
- def start_template_dataflow(self, task_id, variables, parameters, dataflow_template,
- append_job_name=True):
- if append_job_name:
- name = task_id + "-" + str(uuid.uuid1())[:8]
- else:
- name = task_id
- self._start_template_dataflow(
- name, variables, parameters, dataflow_template)
-
- def start_python_dataflow(self, task_id, variables, dataflow, py_options,
- append_job_name=True):
- if append_job_name:
- name = task_id + "-" + str(uuid.uuid1())[:8]
- else:
- name = task_id
- variables['job_name'] = name
-
- def label_formatter(labels_dict):
- return ['--labels={}={}'.format(key, value)
- for key, value in labels_dict.items()]
- self._start_dataflow(task_id, variables, name,
- ["python"] + py_options + [dataflow],
- label_formatter)
-
- def _build_cmd(self, task_id, variables, label_formatter):
- command = ["--runner=DataflowRunner"]
- if variables is not None:
- for attr, value in variables.items():
- if attr == 'labels':
- command += label_formatter(value)
- elif value is None or value.__len__() < 1:
- command.append("--" + attr)
- else:
- command.append("--" + attr + "=" + value)
- return command
-
- def _start_template_dataflow(self, name, variables, parameters, dataflow_template):
- # Builds RuntimeEnvironment from variables dictionary
- # https://cloud.google.com/dataflow/docs/reference/rest/v1b3/RuntimeEnvironment
- environment = {}
- for key in ['maxWorkers', 'zone', 'serviceAccountEmail', 'tempLocation',
- 'bypassTempDirValidation', 'machineType']:
- if key in variables:
- environment.update({key: variables[key]})
- body = {"jobName": name,
- "parameters": parameters,
- "environment": environment}
- service = self.get_conn()
- request = service.projects().templates().launch(projectId=variables['project'],
- gcsPath=dataflow_template,
- body=body)
- response = request.execute()
- variables = self._set_variables(variables)
- _DataflowJob(self.get_conn(), variables['project'], name, variables['region'],
- self.poll_sleep).wait_for_done()
- return response
diff --git a/airflow/contrib/hooks/gcp_dataproc_hook.py b/airflow/contrib/hooks/gcp_dataproc_hook.py
deleted file mode 100644
index fc15137cbf..0000000000
--- a/airflow/contrib/hooks/gcp_dataproc_hook.py
+++ /dev/null
@@ -1,230 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 time
-import uuid
-
-from apiclient.discovery import build
-
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class _DataProcJob(LoggingMixin):
- def __init__(self, dataproc_api, project_id, job, region='global'):
- self.dataproc_api = dataproc_api
- self.project_id = project_id
- self.region = region
- self.job = dataproc_api.projects().regions().jobs().submit(
- projectId=self.project_id,
- region=self.region,
- body=job).execute()
- self.job_id = self.job['reference']['jobId']
- self.log.info(
- 'DataProc job %s is %s',
- self.job_id, str(self.job['status']['state'])
- )
-
- def wait_for_done(self):
- while True:
- self.job = self.dataproc_api.projects().regions().jobs().get(
- projectId=self.project_id,
- region=self.region,
- jobId=self.job_id).execute(num_retries=5)
- if 'ERROR' == self.job['status']['state']:
- print(str(self.job))
- self.log.error('DataProc job %s has errors', self.job_id)
- self.log.error(self.job['status']['details'])
- self.log.debug(str(self.job))
- return False
- if 'CANCELLED' == self.job['status']['state']:
- print(str(self.job))
- self.log.warning('DataProc job %s is cancelled', self.job_id)
- if 'details' in self.job['status']:
- self.log.warning(self.job['status']['details'])
- self.log.debug(str(self.job))
- return False
- if 'DONE' == self.job['status']['state']:
- return True
- self.log.debug(
- 'DataProc job %s is %s',
- self.job_id, str(self.job['status']['state'])
- )
- time.sleep(5)
-
- def raise_error(self, message=None):
- if 'ERROR' == self.job['status']['state']:
- if message is None:
- message = "Google DataProc job has error"
- raise Exception(message + ": " + str(self.job['status']['details']))
-
- def get(self):
- return self.job
-
-
-class _DataProcJobBuilder:
- def __init__(self, project_id, task_id, cluster_name, job_type, properties):
- name = task_id + "_" + str(uuid.uuid1())[:8]
- self.job_type = job_type
- self.job = {
- "job": {
- "reference": {
- "projectId": project_id,
- "jobId": name,
- },
- "placement": {
- "clusterName": cluster_name
- },
- job_type: {
- }
- }
- }
- if properties is not None:
- self.job["job"][job_type]["properties"] = properties
-
- def add_variables(self, variables):
- if variables is not None:
- self.job["job"][self.job_type]["scriptVariables"] = variables
-
- def add_args(self, args):
- if args is not None:
- self.job["job"][self.job_type]["args"] = args
-
- def add_query(self, query):
- self.job["job"][self.job_type]["queryList"] = {'queries': [query]}
-
- def add_query_uri(self, query_uri):
- self.job["job"][self.job_type]["queryFileUri"] = query_uri
-
- def add_jar_file_uris(self, jars):
- if jars is not None:
- self.job["job"][self.job_type]["jarFileUris"] = jars
-
- def add_archive_uris(self, archives):
- if archives is not None:
- self.job["job"][self.job_type]["archiveUris"] = archives
-
- def add_file_uris(self, files):
- if files is not None:
- self.job["job"][self.job_type]["fileUris"] = files
-
- def add_python_file_uris(self, pyfiles):
- if pyfiles is not None:
- self.job["job"][self.job_type]["pythonFileUris"] = pyfiles
-
- def set_main(self, main_jar, main_class):
- if main_class is not None and main_jar is not None:
- raise Exception("Set either main_jar or main_class")
- if main_jar:
- self.job["job"][self.job_type]["mainJarFileUri"] = main_jar
- else:
- self.job["job"][self.job_type]["mainClass"] = main_class
-
- def set_python_main(self, main):
- self.job["job"][self.job_type]["mainPythonFileUri"] = main
-
- def set_job_name(self, name):
- self.job["job"]["reference"]["jobId"] = name + "_" + str(uuid.uuid1())[:8]
-
- def build(self):
- return self.job
-
-
-class _DataProcOperation(LoggingMixin):
- """Continuously polls Dataproc Operation until it completes."""
- def __init__(self, dataproc_api, operation):
- self.dataproc_api = dataproc_api
- self.operation = operation
- self.operation_name = self.operation['name']
-
- def wait_for_done(self):
- if self._check_done():
- return True
-
- self.log.info(
- 'Waiting for Dataproc Operation %s to finish', self.operation_name)
- while True:
- time.sleep(10)
- self.operation = (
- self.dataproc_api.projects()
- .regions()
- .operations()
- .get(name=self.operation_name)
- .execute(num_retries=5))
-
- if self._check_done():
- return True
-
- def get(self):
- return self.operation
-
- def _check_done(self):
- if 'done' in self.operation:
- if 'error' in self.operation:
- self.log.warning(
- 'Dataproc Operation %s failed with error: %s',
- self.operation_name, self.operation['error']['message'])
- self._raise_error()
- else:
- self.log.info(
- 'Dataproc Operation %s done', self.operation['name'])
- return True
- return False
-
- def _raise_error(self):
- raise Exception('Google Dataproc Operation %s failed: %s' %
- (self.operation_name, self.operation['error']['message']))
-
-
-class DataProcHook(GoogleCloudBaseHook):
- """Hook for Google Cloud Dataproc APIs."""
- def __init__(self,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- api_version='v1beta2'):
- super(DataProcHook, self).__init__(gcp_conn_id, delegate_to)
- self.api_version = api_version
-
- def get_conn(self):
- """Returns a Google Cloud Dataproc service object."""
- http_authorized = self._authorize()
- return build(
- 'dataproc', self.api_version, http=http_authorized,
- cache_discovery=False)
-
- def get_cluster(self, project_id, region, cluster_name):
- return self.get_conn().projects().regions().clusters().get(
- projectId=project_id,
- region=region,
- clusterName=cluster_name
- ).execute(num_retries=5)
-
- def submit(self, project_id, job, region='global'):
- submitted = _DataProcJob(self.get_conn(), project_id, job, region)
- if not submitted.wait_for_done():
- submitted.raise_error('DataProcTask has errors')
-
- def create_job_template(self, task_id, cluster_name, job_type, properties):
- return _DataProcJobBuilder(self.project_id, task_id, cluster_name,
- job_type, properties)
-
- def await(self, operation):
- """Awaits for Google Cloud Dataproc Operation to complete."""
- submitted = _DataProcOperation(self.get_conn(), operation)
- submitted.wait_for_done()
diff --git a/airflow/contrib/hooks/gcp_mlengine_hook.py b/airflow/contrib/hooks/gcp_mlengine_hook.py
deleted file mode 100644
index 66f392b156..0000000000
--- a/airflow/contrib/hooks/gcp_mlengine_hook.py
+++ /dev/null
@@ -1,266 +0,0 @@
-#
-# 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 random
-import time
-from apiclient import errors
-from apiclient.discovery import build
-
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-def _poll_with_exponential_delay(request, max_n, is_done_func, is_error_func):
- log = LoggingMixin().log
-
- for i in range(0, max_n):
- try:
- response = request.execute()
- if is_error_func(response):
- raise ValueError(
- 'The response contained an error: {}'.format(response)
- )
- elif is_done_func(response):
- log.info('Operation is done: %s', response)
- return response
- else:
- time.sleep((2**i) + (random.randint(0, 1000) / 1000))
- except errors.HttpError as e:
- if e.resp.status != 429:
- log.info('Something went wrong. Not retrying: %s', format(e))
- raise
- else:
- time.sleep((2**i) + (random.randint(0, 1000) / 1000))
-
-
-class MLEngineHook(GoogleCloudBaseHook):
- def __init__(self, gcp_conn_id='google_cloud_default', delegate_to=None):
- super(MLEngineHook, self).__init__(gcp_conn_id, delegate_to)
- self._mlengine = self.get_conn()
-
- def get_conn(self):
- """
- Returns a Google MLEngine service object.
- """
- authed_http = self._authorize()
- return build('ml', 'v1', http=authed_http, cache_discovery=False)
-
- def create_job(self, project_id, job, use_existing_job_fn=None):
- """
- Launches a MLEngine job and wait for it to reach a terminal state.
-
- :param project_id: The Google Cloud project id within which MLEngine
- job will be launched.
- :type project_id: string
-
- :param job: MLEngine Job object that should be provided to the MLEngine
- API, such as: ::
- {
- 'jobId': 'my_job_id',
- 'trainingInput': {
- 'scaleTier': 'STANDARD_1',
- ...
- }
- }
- :type job: dict
-
- :param use_existing_job_fn: In case that a MLEngine job with the same
- job_id already exist, this method (if provided) will decide whether
- we should use this existing job, continue waiting for it to finish
- and returning the job object. It should accepts a MLEngine job
- object, and returns a boolean value indicating whether it is OK to
- reuse the existing job. If 'use_existing_job_fn' is not provided,
- we by default reuse the existing MLEngine job.
- :type use_existing_job_fn: function
-
- :return: The MLEngine job object if the job successfully reach a
- terminal state (which might be FAILED or CANCELLED state).
- :rtype: dict
- """
- request = self._mlengine.projects().jobs().create(
- parent='projects/{}'.format(project_id),
- body=job)
- job_id = job['jobId']
-
- try:
- request.execute()
- except errors.HttpError as e:
- # 409 means there is an existing job with the same job ID.
- if e.resp.status == 409:
- if use_existing_job_fn is not None:
- existing_job = self._get_job(project_id, job_id)
- if not use_existing_job_fn(existing_job):
- self.log.error(
- 'Job with job_id %s already exist, but it does '
- 'not match our expectation: %s',
- job_id, existing_job
- )
- raise
- self.log.info(
- 'Job with job_id %s already exist. Will waiting for it to finish',
- job_id
- )
- else:
- self.log.error('Failed to create MLEngine job: {}'.format(e))
- raise
-
- return self._wait_for_job_done(project_id, job_id)
-
- def _get_job(self, project_id, job_id):
- """
- Gets a MLEngine job based on the job name.
-
- :return: MLEngine job object if succeed.
- :rtype: dict
-
- Raises:
- apiclient.errors.HttpError: if HTTP error is returned from server
- """
- job_name = 'projects/{}/jobs/{}'.format(project_id, job_id)
- request = self._mlengine.projects().jobs().get(name=job_name)
- while True:
- try:
- return request.execute()
- except errors.HttpError as e:
- if e.resp.status == 429:
- # polling after 30 seconds when quota failure occurs
- time.sleep(30)
- else:
- self.log.error('Failed to get MLEngine job: {}'.format(e))
- raise
-
- def _wait_for_job_done(self, project_id, job_id, interval=30):
- """
- Waits for the Job to reach a terminal state.
-
- This method will periodically check the job state until the job reach
- a terminal state.
-
- Raises:
- apiclient.errors.HttpError: if HTTP error is returned when getting
- the job
- """
- assert interval > 0
- while True:
- job = self._get_job(project_id, job_id)
- if job['state'] in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
- return job
- time.sleep(interval)
-
- def create_version(self, project_id, model_name, version_spec):
- """
- Creates the Version on Google Cloud ML Engine.
-
- Returns the operation if the version was created successfully and
- raises an error otherwise.
- """
- parent_name = 'projects/{}/models/{}'.format(project_id, model_name)
- create_request = self._mlengine.projects().models().versions().create(
- parent=parent_name, body=version_spec)
- response = create_request.execute()
- get_request = self._mlengine.projects().operations().get(
- name=response['name'])
-
- return _poll_with_exponential_delay(
- request=get_request,
- max_n=9,
- is_done_func=lambda resp: resp.get('done', False),
- is_error_func=lambda resp: resp.get('error', None) is not None)
-
- def set_default_version(self, project_id, model_name, version_name):
- """
- Sets a version to be the default. Blocks until finished.
- """
- full_version_name = 'projects/{}/models/{}/versions/{}'.format(
- project_id, model_name, version_name)
- request = self._mlengine.projects().models().versions().setDefault(
- name=full_version_name, body={})
-
- try:
- response = request.execute()
- self.log.info('Successfully set version: %s to default', response)
- return response
- except errors.HttpError as e:
- self.log.error('Something went wrong: %s', e)
- raise
-
- def list_versions(self, project_id, model_name):
- """
- Lists all available versions of a model. Blocks until finished.
- """
- result = []
- full_parent_name = 'projects/{}/models/{}'.format(
- project_id, model_name)
- request = self._mlengine.projects().models().versions().list(
- parent=full_parent_name, pageSize=100)
-
- response = request.execute()
- next_page_token = response.get('nextPageToken', None)
- result.extend(response.get('versions', []))
- while next_page_token is not None:
- next_request = self._mlengine.projects().models().versions().list(
- parent=full_parent_name,
- pageToken=next_page_token,
- pageSize=100)
- response = next_request.execute()
- next_page_token = response.get('nextPageToken', None)
- result.extend(response.get('versions', []))
- time.sleep(5)
- return result
-
- def delete_version(self, project_id, model_name, version_name):
- """
- Deletes the given version of a model. Blocks until finished.
- """
- full_name = 'projects/{}/models/{}/versions/{}'.format(
- project_id, model_name, version_name)
- delete_request = self._mlengine.projects().models().versions().delete(
- name=full_name)
- response = delete_request.execute()
- get_request = self._mlengine.projects().operations().get(
- name=response['name'])
-
- return _poll_with_exponential_delay(
- request=get_request,
- max_n=9,
- is_done_func=lambda resp: resp.get('done', False),
- is_error_func=lambda resp: resp.get('error', None) is not None)
-
- def create_model(self, project_id, model):
- """
- Create a Model. Blocks until finished.
- """
- assert model['name'] is not None and model['name'] is not ''
- project = 'projects/{}'.format(project_id)
-
- request = self._mlengine.projects().models().create(
- parent=project, body=model)
- return request.execute()
-
- def get_model(self, project_id, model_name):
- """
- Gets a Model. Blocks until finished.
- """
- assert model_name is not None and model_name is not ''
- full_model_name = 'projects/{}/models/{}'.format(
- project_id, model_name)
- request = self._mlengine.projects().models().get(name=full_model_name)
- try:
- return request.execute()
- except errors.HttpError as e:
- if e.resp.status == 404:
- self.log.error('Model was not found: %s', e)
- return None
- raise
diff --git a/airflow/contrib/hooks/gcp_pubsub_hook.py b/airflow/contrib/hooks/gcp_pubsub_hook.py
deleted file mode 100644
index 1d55d5d487..0000000000
--- a/airflow/contrib/hooks/gcp_pubsub_hook.py
+++ /dev/null
@@ -1,285 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 uuid import uuid4
-
-from apiclient.discovery import build
-from apiclient import errors
-
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-
-
-def _format_subscription(project, subscription):
- return 'projects/{}/subscriptions/{}'.format(project, subscription)
-
-
-def _format_topic(project, topic):
- return 'projects/{}/topics/{}'.format(project, topic)
-
-
-class PubSubException(Exception):
- pass
-
-
-class PubSubHook(GoogleCloudBaseHook):
- """Hook for accessing Google Pub/Sub.
-
- The GCP project against which actions are applied is determined by
- the project embedded in the Connection referenced by gcp_conn_id.
- """
-
- def __init__(self, gcp_conn_id='google_cloud_default', delegate_to=None):
- super(PubSubHook, self).__init__(gcp_conn_id, delegate_to=delegate_to)
-
- def get_conn(self):
- """Returns a Pub/Sub service object.
-
- :rtype: apiclient.discovery.Resource
- """
- http_authorized = self._authorize()
- return build(
- 'pubsub', 'v1', http=http_authorized, cache_discovery=False)
-
- def publish(self, project, topic, messages):
- """Publishes messages to a Pub/Sub topic.
-
- :param project: the GCP project ID in which to publish
- :type project: string
- :param topic: the Pub/Sub topic to which to publish; do not
- include the ``projects/{project}/topics/`` prefix.
- :type topic: string
- :param messages: messages to publish; if the data field in a
- message is set, it should already be base64 encoded.
- :type messages: list of PubSub messages; see
- http://cloud.google.com/pubsub/docs/reference/rest/v1/PubsubMessage
- """
- body = {'messages': messages}
- full_topic = _format_topic(project, topic)
- request = self.get_conn().projects().topics().publish(
- topic=full_topic, body=body)
- try:
- request.execute()
- except errors.HttpError as e:
- raise PubSubException(
- 'Error publishing to topic {}'.format(full_topic), e)
-
- def create_topic(self, project, topic, fail_if_exists=False):
- """Creates a Pub/Sub topic, if it does not already exist.
-
- :param project: the GCP project ID in which to create
- the topic
- :type project: string
- :param topic: the Pub/Sub topic name to create; do not
- include the ``projects/{project}/topics/`` prefix.
- :type topic: string
- :param fail_if_exists: if set, raise an exception if the topic
- already exists
- :type fail_if_exists: bool
- """
- service = self.get_conn()
- full_topic = _format_topic(project, topic)
- try:
- service.projects().topics().create(
- name=full_topic, body={}).execute()
- except errors.HttpError as e:
- # Status code 409 indicates that the topic already exists.
- if str(e.resp['status']) == '409':
- message = 'Topic already exists: {}'.format(full_topic)
- self.log.warning(message)
- if fail_if_exists:
- raise PubSubException(message)
- else:
- raise PubSubException(
- 'Error creating topic {}'.format(full_topic), e)
-
- def delete_topic(self, project, topic, fail_if_not_exists=False):
- """Deletes a Pub/Sub topic if it exists.
-
- :param project: the GCP project ID in which to delete the topic
- :type project: string
- :param topic: the Pub/Sub topic name to delete; do not
- include the ``projects/{project}/topics/`` prefix.
- :type topic: string
- :param fail_if_not_exists: if set, raise an exception if the topic
- does not exist
- :type fail_if_not_exists: bool
- """
- service = self.get_conn()
- full_topic = _format_topic(project, topic)
- try:
- service.projects().topics().delete(topic=full_topic).execute()
- except errors.HttpError as e:
- # Status code 409 indicates that the topic was not found
- if str(e.resp['status']) == '404':
- message = 'Topic does not exist: {}'.format(full_topic)
- self.log.warning(message)
- if fail_if_not_exists:
- raise PubSubException(message)
- else:
- raise PubSubException(
- 'Error deleting topic {}'.format(full_topic), e)
-
- def create_subscription(self, topic_project, topic, subscription=None,
- subscription_project=None, ack_deadline_secs=10,
- fail_if_exists=False):
- """Creates a Pub/Sub subscription, if it does not already exist.
-
- :param topic_project: the GCP project ID of the topic that the
- subscription will be bound to.
- :type topic_project: string
- :param topic: the Pub/Sub topic name that the subscription will be bound
- to create; do not include the ``projects/{project}/subscriptions/``
- prefix.
- :type topic: string
- :param subscription: the Pub/Sub subscription name. If empty, a random
- name will be generated using the uuid module
- :type subscription: string
- :param subscription_project: the GCP project ID where the subscription
- will be created. If unspecified, ``topic_project`` will be used.
- :type subscription_project: string
- :param ack_deadline_secs: Number of seconds that a subscriber has to
- acknowledge each message pulled from the subscription
- :type ack_deadline_secs: int
- :param fail_if_exists: if set, raise an exception if the topic
- already exists
- :type fail_if_exists: bool
- :return: subscription name which will be the system-generated value if
- the ``subscription`` parameter is not supplied
- :rtype: string
- """
- service = self.get_conn()
- full_topic = _format_topic(topic_project, topic)
- if not subscription:
- subscription = 'sub-{}'.format(uuid4())
- if not subscription_project:
- subscription_project = topic_project
- full_subscription = _format_subscription(subscription_project,
- subscription)
- body = {
- 'topic': full_topic,
- 'ackDeadlineSeconds': ack_deadline_secs
- }
- try:
- service.projects().subscriptions().create(
- name=full_subscription, body=body).execute()
- except errors.HttpError as e:
- # Status code 409 indicates that the subscription already exists.
- if str(e.resp['status']) == '409':
- message = 'Subscription already exists: {}'.format(
- full_subscription)
- self.log.warning(message)
- if fail_if_exists:
- raise PubSubException(message)
- else:
- raise PubSubException(
- 'Error creating subscription {}'.format(full_subscription),
- e)
- return subscription
-
- def delete_subscription(self, project, subscription,
- fail_if_not_exists=False):
- """Deletes a Pub/Sub subscription, if it exists.
-
- :param project: the GCP project ID where the subscription exists
- :type project: string
- :param subscription: the Pub/Sub subscription name to delete; do not
- include the ``projects/{project}/subscriptions/`` prefix.
- :type subscription: string
- :param fail_if_not_exists: if set, raise an exception if the topic
- does not exist
- :type fail_if_not_exists: bool
- """
- service = self.get_conn()
- full_subscription = _format_subscription(project, subscription)
- try:
- service.projects().subscriptions().delete(
- subscription=full_subscription).execute()
- except errors.HttpError as e:
- # Status code 404 indicates that the subscription was not found
- if str(e.resp['status']) == '404':
- message = 'Subscription does not exist: {}'.format(
- full_subscription)
- self.log.warning(message)
- if fail_if_not_exists:
- raise PubSubException(message)
- else:
- raise PubSubException(
- 'Error deleting subscription {}'.format(full_subscription),
- e)
-
- def pull(self, project, subscription, max_messages,
- return_immediately=False):
- """Pulls up to ``max_messages`` messages from Pub/Sub subscription.
-
- :param project: the GCP project ID where the subscription exists
- :type project: string
- :param subscription: the Pub/Sub subscription name to pull from; do not
- include the 'projects/{project}/topics/' prefix.
- :type subscription: string
- :param max_messages: The maximum number of messages to return from
- the Pub/Sub API.
- :type max_messages: int
- :param return_immediately: If set, the Pub/Sub API will immediately
- return if no messages are available. Otherwise, the request will
- block for an undisclosed, but bounded period of time
- :type return_immediately: bool
- :return A list of Pub/Sub ReceivedMessage objects each containing
- an ``ackId`` property and a ``message`` property, which includes
- the base64-encoded message content. See
- https://cloud.google.com/pubsub/docs/reference/rest/v1/\
- projects.subscriptions/pull#ReceivedMessage
- """
- service = self.get_conn()
- full_subscription = _format_subscription(project, subscription)
- body = {
- 'maxMessages': max_messages,
- 'returnImmediately': return_immediately
- }
- try:
- response = service.projects().subscriptions().pull(
- subscription=full_subscription, body=body).execute()
- return response.get('receivedMessages', [])
- except errors.HttpError as e:
- raise PubSubException(
- 'Error pulling messages from subscription {}'.format(
- full_subscription), e)
-
- def acknowledge(self, project, subscription, ack_ids):
- """Pulls up to ``max_messages`` messages from Pub/Sub subscription.
-
- :param project: the GCP project name or ID in which to create
- the topic
- :type project: string
- :param subscription: the Pub/Sub subscription name to delete; do not
- include the 'projects/{project}/topics/' prefix.
- :type subscription: string
- :param ack_ids: List of ReceivedMessage ackIds from a previous pull
- response
- :type ack_ids: list
- """
- service = self.get_conn()
- full_subscription = _format_subscription(project, subscription)
- try:
- service.projects().subscriptions().acknowledge(
- subscription=full_subscription, body={'ackIds': ack_ids}
- ).execute()
- except errors.HttpError as e:
- raise PubSubException(
- 'Error acknowledging {} messages pulled from subscription {}'
- .format(len(ack_ids), full_subscription), e)
diff --git a/airflow/contrib/hooks/gcs_hook.py b/airflow/contrib/hooks/gcs_hook.py
deleted file mode 100644
index c5e356f41c..0000000000
--- a/airflow/contrib/hooks/gcs_hook.py
+++ /dev/null
@@ -1,537 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 apiclient.discovery import build
-from apiclient.http import MediaFileUpload
-from googleapiclient import errors
-
-from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
-from airflow.exceptions import AirflowException
-
-import re
-
-
-class GoogleCloudStorageHook(GoogleCloudBaseHook):
- """
- Interact with Google Cloud Storage. This hook uses the Google Cloud Platform
- connection.
- """
-
- def __init__(self,
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None):
- super(GoogleCloudStorageHook, self).__init__(google_cloud_storage_conn_id,
- delegate_to)
-
- def get_conn(self):
- """
- Returns a Google Cloud Storage service object.
- """
- http_authorized = self._authorize()
- return build(
- 'storage', 'v1', http=http_authorized, cache_discovery=False)
-
- # pylint:disable=redefined-builtin
- def copy(self, source_bucket, source_object, destination_bucket=None,
- destination_object=None):
- """
- Copies an object from a bucket to another, with renaming if requested.
-
- destination_bucket or destination_object can be omitted, in which case
- source bucket/object is used, but not both.
-
- :param source_bucket: The bucket of the object to copy from.
- :type source_bucket: string
- :param source_object: The object to copy.
- :type source_object: string
- :param destination_bucket: The destination of the object to copied to.
- Can be omitted; then the same bucket is used.
- :type destination_bucket: string
- :param destination_object: The (renamed) path of the object if given.
- Can be omitted; then the same name is used.
- """
- destination_bucket = destination_bucket or source_bucket
- destination_object = destination_object or source_object
- if source_bucket == destination_bucket and \
- source_object == destination_object:
-
- raise ValueError(
- 'Either source/destination bucket or source/destination object '
- 'must be different, not both the same: bucket=%s, object=%s' %
- (source_bucket, source_object))
- if not source_bucket or not source_object:
- raise ValueError('source_bucket and source_object cannot be empty.')
-
- service = self.get_conn()
- try:
- service \
- .objects() \
- .copy(sourceBucket=source_bucket, sourceObject=source_object,
- destinationBucket=destination_bucket,
- destinationObject=destination_object, body='') \
- .execute()
- return True
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- return False
- raise
-
- def rewrite(self, source_bucket, source_object, destination_bucket,
- destination_object=None):
- """
- Has the same functionality as copy, except that will work on files
- over 5 TB, as well as when copying between locations and/or storage
- classes.
-
- destination_object can be omitted, in which case source_object is used.
-
- :param source_bucket: The bucket of the object to copy from.
- :type source_bucket: string
- :param source_object: The object to copy.
- :type source_object: string
- :param destination_bucket: The destination of the object to copied to.
- :type destination_bucket: string
- :param destination_object: The (renamed) path of the object if given.
- Can be omitted; then the same name is used.
- """
- destination_object = destination_object or source_object
- if (source_bucket == destination_bucket and
- source_object == destination_object):
- raise ValueError(
- 'Either source/destination bucket or source/destination object '
- 'must be different, not both the same: bucket=%s, object=%s' %
- (source_bucket, source_object))
- if not source_bucket or not source_object:
- raise ValueError('source_bucket and source_object cannot be empty.')
-
- service = self.get_conn()
- request_count = 1
- try:
- result = service.objects() \
- .rewrite(sourceBucket=source_bucket, sourceObject=source_object,
- destinationBucket=destination_bucket,
- destinationObject=destination_object, body='') \
- .execute()
- self.log.info('Rewrite request #%s: %s', request_count, result)
- while not result['done']:
- request_count += 1
- result = service.objects() \
- .rewrite(sourceBucket=source_bucket, sourceObject=source_object,
- destinationBucket=destination_bucket,
- destinationObject=destination_object,
- rewriteToken=result['rewriteToken'], body='') \
- .execute()
- self.log.info('Rewrite request #%s: %s', request_count, result)
- return True
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- return False
- raise
-
- # pylint:disable=redefined-builtin
- def download(self, bucket, object, filename=None):
- """
- Get a file from Google Cloud Storage.
-
- :param bucket: The bucket to fetch from.
- :type bucket: string
- :param object: The object to fetch.
- :type object: string
- :param filename: If set, a local file path where the file should be written to.
- :type filename: string
- """
- service = self.get_conn()
- downloaded_file_bytes = service \
- .objects() \
- .get_media(bucket=bucket, object=object) \
- .execute()
-
- # Write the file to local file path, if requested.
- if filename:
- write_argument = 'wb' if isinstance(downloaded_file_bytes, bytes) else 'w'
- with open(filename, write_argument) as file_fd:
- file_fd.write(downloaded_file_bytes)
-
- return downloaded_file_bytes
-
- # pylint:disable=redefined-builtin
- def upload(self, bucket, object, filename, mime_type='application/octet-stream'):
- """
- Uploads a local file to Google Cloud Storage.
-
- :param bucket: The bucket to upload to.
- :type bucket: string
- :param object: The object name to set when uploading the local file.
- :type object: string
- :param filename: The local file path to the file to be uploaded.
- :type filename: string
- :param mime_type: The MIME type to set when uploading the file.
- :type mime_type: string
- """
- service = self.get_conn()
- media = MediaFileUpload(filename, mime_type)
- try:
- service \
- .objects() \
- .insert(bucket=bucket, name=object, media_body=media) \
- .execute()
- return True
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- return False
- raise
-
- # pylint:disable=redefined-builtin
- def exists(self, bucket, object):
- """
- Checks for the existence of a file in Google Cloud Storage.
-
- :param bucket: The Google cloud storage bucket where the object is.
- :type bucket: string
- :param object: The name of the object to check in the Google cloud
- storage bucket.
- :type object: string
- """
- service = self.get_conn()
- try:
- service \
- .objects() \
- .get(bucket=bucket, object=object) \
- .execute()
- return True
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- return False
- raise
-
- # pylint:disable=redefined-builtin
- def is_updated_after(self, bucket, object, ts):
- """
- Checks if an object is updated in Google Cloud Storage.
-
- :param bucket: The Google cloud storage bucket where the object is.
- :type bucket: string
- :param object: The name of the object to check in the Google cloud
- storage bucket.
- :type object: string
- :param ts: The timestamp to check against.
- :type ts: datetime
- """
- service = self.get_conn()
- try:
- response = (service
- .objects()
- .get(bucket=bucket, object=object)
- .execute())
-
- if 'updated' in response:
- import dateutil.parser
- import dateutil.tz
-
- if not ts.tzinfo:
- ts = ts.replace(tzinfo=dateutil.tz.tzutc())
-
- updated = dateutil.parser.parse(response['updated'])
- self.log.info("Verify object date: %s > %s", updated, ts)
-
- if updated > ts:
- return True
-
- except errors.HttpError as ex:
- if ex.resp['status'] != '404':
- raise
-
- return False
-
- def delete(self, bucket, object, generation=None):
- """
- Delete an object if versioning is not enabled for the bucket, or if generation
- parameter is used.
-
- :param bucket: name of the bucket, where the object resides
- :type bucket: string
- :param object: name of the object to delete
- :type object: string
- :param generation: if present, permanently delete the object of this generation
- :type generation: string
- :return: True if succeeded
- """
- service = self.get_conn()
-
- try:
- service \
- .objects() \
- .delete(bucket=bucket, object=object, generation=generation) \
- .execute()
- return True
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- return False
- raise
-
- def list(self, bucket, versions=None, maxResults=None, prefix=None, delimiter=None):
- """
- List all objects from the bucket with the give string prefix in name
-
- :param bucket: bucket name
- :type bucket: string
- :param versions: if true, list all versions of the objects
- :type versions: boolean
- :param maxResults: max count of items to return in a single page of responses
- :type maxResults: integer
- :param prefix: prefix string which filters objects whose name begin with
- this prefix
- :type prefix: string
- :param delimiter: filters objects based on the delimiter (for e.g '.csv')
- :type delimiter: string
- :return: a stream of object names matching the filtering criteria
- """
- service = self.get_conn()
-
- ids = list()
- pageToken = None
- while(True):
- response = service.objects().list(
- bucket=bucket,
- versions=versions,
- maxResults=maxResults,
- pageToken=pageToken,
- prefix=prefix,
- delimiter=delimiter
- ).execute()
-
- if 'prefixes' not in response:
- if 'items' not in response:
- self.log.info("No items found for prefix: %s", prefix)
- break
-
- for item in response['items']:
- if item and 'name' in item:
- ids.append(item['name'])
- else:
- for item in response['prefixes']:
- ids.append(item)
-
- if 'nextPageToken' not in response:
- # no further pages of results, so stop the loop
- break
-
- pageToken = response['nextPageToken']
- if not pageToken:
- # empty next page token
- break
- return ids
-
- def get_size(self, bucket, object):
- """
- Gets the size of a file in Google Cloud Storage.
-
- :param bucket: The Google cloud storage bucket where the object is.
- :type bucket: string
- :param object: The name of the object to check in the Google cloud storage bucket.
- :type object: string
-
- """
- self.log.info('Checking the file size of object: %s in bucket: %s',
- object,
- bucket)
- service = self.get_conn()
- try:
- response = service.objects().get(
- bucket=bucket,
- object=object
- ).execute()
-
- if 'name' in response and response['name'][-1] != '/':
- # Remove Directories & Just check size of files
- size = response['size']
- self.log.info('The file size of %s is %s bytes.', object, size)
- return size
- else:
- raise ValueError('Object is not a file')
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- raise ValueError('Object Not Found')
-
- def get_crc32c(self, bucket, object):
- """
- Gets the CRC32c checksum of an object in Google Cloud Storage.
-
- :param bucket: The Google cloud storage bucket where the object is.
- :type bucket: string
- :param object: The name of the object to check in the Google cloud
- storage bucket.
- :type object: string
- """
- self.log.info('Retrieving the crc32c checksum of '
- 'object: %s in bucket: %s', object, bucket)
- service = self.get_conn()
- try:
- response = service.objects().get(
- bucket=bucket,
- object=object
- ).execute()
-
- crc32c = response['crc32c']
- self.log.info('The crc32c checksum of %s is %s', object, crc32c)
- return crc32c
-
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- raise ValueError('Object Not Found')
-
- def get_md5hash(self, bucket, object):
- """
- Gets the MD5 hash of an object in Google Cloud Storage.
-
- :param bucket: The Google cloud storage bucket where the object is.
- :type bucket: string
- :param object: The name of the object to check in the Google cloud
- storage bucket.
- :type object: string
- """
- self.log.info('Retrieving the MD5 hash of '
- 'object: %s in bucket: %s', object, bucket)
- service = self.get_conn()
- try:
- response = service.objects().get(
- bucket=bucket,
- object=object
- ).execute()
-
- md5hash = response['md5Hash']
- self.log.info('The md5Hash of %s is %s', object, md5hash)
- return md5hash
-
- except errors.HttpError as ex:
- if ex.resp['status'] == '404':
- raise ValueError('Object Not Found')
-
- def create_bucket(self,
- bucket_name,
- storage_class='MULTI_REGIONAL',
- location='US',
- project_id=None,
- labels=None
- ):
- """
- Creates a new bucket. Google Cloud Storage uses a flat namespace, so
- you can't create a bucket with a name that is already in use.
-
- .. seealso::
- For more information, see Bucket Naming Guidelines:
- https://cloud.google.com/storage/docs/bucketnaming.html#requirements
-
- :param bucket_name: The name of the bucket.
- :type bucket_name: string
- :param storage_class: This defines how objects in the bucket are stored
- and determines the SLA and the cost of storage. Values include
-
- - ``MULTI_REGIONAL``
- - ``REGIONAL``
- - ``STANDARD``
- - ``NEARLINE``
- - ``COLDLINE``.
- If this value is not specified when the bucket is
- created, it will default to STANDARD.
- :type storage_class: string
- :param location: The location of the bucket.
- Object data for objects in the bucket resides in physical storage
- within this region. Defaults to US.
-
- .. seealso::
- https://developers.google.com/storage/docs/bucket-locations
-
- :type location: string
- :param project_id: The ID of the GCP Project.
- :type project_id: string
- :param labels: User-provided labels, in key/value pairs.
- :type labels: dict
- :return: If successful, it returns the ``id`` of the bucket.
- """
-
- project_id = project_id if project_id is not None else self.project_id
- storage_classes = [
- 'MULTI_REGIONAL',
- 'REGIONAL',
- 'NEARLINE',
- 'COLDLINE',
- 'STANDARD', # alias for MULTI_REGIONAL/REGIONAL, based on location
- ]
-
- self.log.info('Creating Bucket: %s; Location: %s; Storage Class: %s',
- bucket_name, location, storage_class)
- assert storage_class in storage_classes, \
- 'Invalid value ({}) passed to storage_class. Value should be ' \
- 'one of {}'.format(storage_class, storage_classes)
-
- assert re.match('[a-zA-Z0-9]+', bucket_name[0]), \
- 'Bucket names must start with a number or letter.'
-
- assert re.match('[a-zA-Z0-9]+', bucket_name[-1]), \
- 'Bucket names must end with a number or letter.'
-
- service = self.get_conn()
- bucket_resource = {
- 'name': bucket_name,
- 'location': location,
- 'storageClass': storage_class
- }
-
- self.log.info('The Default Project ID is %s', self.project_id)
-
- if labels is not None:
- bucket_resource['labels'] = labels
-
- try:
- response = service.buckets().insert(
- project=project_id,
- body=bucket_resource
- ).execute()
-
- self.log.info('Bucket: %s created successfully.', bucket_name)
-
- return response['id']
-
- except errors.HttpError as ex:
- raise AirflowException(
- 'Bucket creation failed. Error was: {}'.format(ex.content)
- )
-
-
-def _parse_gcs_url(gsurl):
- """
- Given a Google Cloud Storage URL (gs://<bucket>/<blob>), returns a
- tuple containing the corresponding bucket and blob.
- """
- # Python 3
- try:
- from urllib.parse import urlparse
- # Python 2
- except ImportError:
- from urlparse import urlparse
-
- parsed_url = urlparse(gsurl)
- if not parsed_url.netloc:
- raise AirflowException('Please provide a bucket name')
- else:
- bucket = parsed_url.netloc
- # Remove leading '/' but NOT trailing one
- blob = parsed_url.path.lstrip('/')
- return bucket, blob
diff --git a/airflow/contrib/hooks/jenkins_hook.py b/airflow/contrib/hooks/jenkins_hook.py
deleted file mode 100644
index b1bea30b6e..0000000000
--- a/airflow/contrib/hooks/jenkins_hook.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.hooks.base_hook import BaseHook
-
-import jenkins
-import distutils
-
-
-class JenkinsHook(BaseHook):
- """
- Hook to manage connection to jenkins server
- """
-
- def __init__(self, conn_id='jenkins_default'):
- connection = self.get_connection(conn_id)
- self.connection = connection
- connectionPrefix = 'http'
- # connection.extra contains info about using https (true) or http (false)
- if connection.extra is None or connection.extra == '':
- connection.extra = 'false'
- # set a default value to connection.extra
- # to avoid rising ValueError in strtobool
- if distutils.util.strtobool(connection.extra):
- connectionPrefix = 'https'
- url = '%s://%s:%d' % (connectionPrefix, connection.host, connection.port)
- self.log.info('Trying to connect to %s', url)
- self.jenkins_server = jenkins.Jenkins(url, connection.login, connection.password)
-
- def get_jenkins_server(self):
- return self.jenkins_server
diff --git a/airflow/contrib/hooks/jira_hook.py b/airflow/contrib/hooks/jira_hook.py
deleted file mode 100644
index c6806d935a..0000000000
--- a/airflow/contrib/hooks/jira_hook.py
+++ /dev/null
@@ -1,87 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 jira import JIRA
-from jira.exceptions import JIRAError
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class JiraHook(BaseHook, LoggingMixin):
- """
- Jira interaction hook, a Wrapper around JIRA Python SDK.
-
- :param jira_conn_id: reference to a pre-defined Jira Connection
- :type jira_conn_id: string
- """
- def __init__(self,
- jira_conn_id='jira_default',
- proxies=None):
- super(JiraHook, self).__init__(jira_conn_id)
- self.jira_conn_id = jira_conn_id
- self.proxies = proxies
- self.client = None
- self.get_conn()
-
- def get_conn(self):
- if not self.client:
- self.log.debug('Creating Jira client for conn_id: %s', self.jira_conn_id)
-
- get_server_info = True
- validate = True
- extra_options = {}
- conn = None
-
- if self.jira_conn_id is not None:
- conn = self.get_connection(self.jira_conn_id)
- if conn.extra is not None:
- extra_options = conn.extra_dejson
- # only required attributes are taken for now,
- # more can be added ex: async, logging, max_retries
-
- # verify
- if 'verify' in extra_options \
- and extra_options['verify'].lower() == 'false':
- extra_options['verify'] = False
-
- # validate
- if 'validate' in extra_options \
- and extra_options['validate'].lower() == 'false':
- validate = False
-
- if 'get_server_info' in extra_options \
- and extra_options['get_server_info'].lower() == 'false':
- get_server_info = False
-
- try:
- self.client = JIRA(conn.host,
- options=extra_options,
- basic_auth=(conn.login, conn.password),
- get_server_info=get_server_info,
- validate=validate,
- proxies=self.proxies)
- except JIRAError as jira_error:
- raise AirflowException('Failed to create jira client, jira error: %s'
- % str(jira_error))
- except Exception as e:
- raise AirflowException('Failed to create jira client, error: %s'
- % str(e))
-
- return self.client
diff --git a/airflow/contrib/hooks/mongo_hook.py b/airflow/contrib/hooks/mongo_hook.py
deleted file mode 100644
index 6ae71a8c8d..0000000000
--- a/airflow/contrib/hooks/mongo_hook.py
+++ /dev/null
@@ -1,119 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Licensed 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 ssl import CERT_NONE
-
-from airflow.hooks.base_hook import BaseHook
-from pymongo import MongoClient
-
-
-class MongoHook(BaseHook):
- """
- PyMongo Wrapper to Interact With Mongo Database
- Mongo Connection Documentation
- https://docs.mongodb.com/manual/reference/connection-string/index.html
- You can specify connection string options in extra field of your connection
- https://docs.mongodb.com/manual/reference/connection-string/index.html#connection-string-options
- ex.
- {replicaSet: test, ssl: True, connectTimeoutMS: 30000}
- """
- conn_type = 'MongoDb'
-
- def __init__(self, conn_id='mongo_default', *args, **kwargs):
- super(MongoHook, self).__init__(source='mongo')
-
- self.mongo_conn_id = conn_id
- self.connection = self.get_connection(conn_id)
- self.extras = self.connection.extra_dejson
- self.client = None
-
- def get_conn(self):
- """
- Fetches PyMongo Client
- """
- if self.client is not None:
- return self.client
-
- conn = self.connection
-
- uri = 'mongodb://{creds}{host}{port}/{database}'.format(
- creds='{}:{}@'.format(
- conn.login, conn.password
- ) if conn.login is not None else '',
-
- host=conn.host,
- port='' if conn.port is None else ':{}'.format(conn.port),
- database='' if conn.schema is None else conn.schema
- )
-
- # Mongo Connection Options dict that is unpacked when passed to MongoClient
- options = self.extras
-
- # If we are using SSL disable requiring certs from specific hostname
- if options.get('ssl', False):
- options.update({'ssl_cert_reqs': CERT_NONE})
-
- self.client = MongoClient(uri, **options)
-
- return self.client
-
- def get_collection(self, mongo_collection, mongo_db=None):
- """
- Fetches a mongo collection object for querying.
-
- Uses connection schema as DB unless specified.
- """
- mongo_db = mongo_db if mongo_db is not None else self.connection.schema
- mongo_conn = self.get_conn()
-
- return mongo_conn.get_database(mongo_db).get_collection(mongo_collection)
-
- def aggregate(self, mongo_collection, aggregate_query, mongo_db=None, **kwargs):
- """
- Runs an aggregation pipeline and returns the results
- https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.aggregate
- https://api.mongodb.com/python/current/examples/aggregation.html
- """
- collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
-
- return collection.aggregate(aggregate_query, **kwargs)
-
- def find(self, mongo_collection, query, find_one=False, mongo_db=None, **kwargs):
- """
- Runs a mongo find query and returns the results
- https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.find
- """
- collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
-
- if find_one:
- return collection.find_one(query, **kwargs)
- else:
- return collection.find(query, **kwargs)
-
- def insert_one(self, mongo_collection, doc, mongo_db=None, **kwargs):
- """
- Inserts a single document into a mongo collection
- https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.insert_one
- """
- collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
-
- return collection.insert_one(doc, **kwargs)
-
- def insert_many(self, mongo_collection, docs, mongo_db=None, **kwargs):
- """
- Inserts many docs into a mongo collection.
- https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.insert_many
- """
- collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
-
- return collection.insert_many(docs, **kwargs)
diff --git a/airflow/contrib/hooks/pinot_hook.py b/airflow/contrib/hooks/pinot_hook.py
deleted file mode 100644
index d731211efc..0000000000
--- a/airflow/contrib/hooks/pinot_hook.py
+++ /dev/null
@@ -1,105 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 six
-
-from pinotdb import connect
-
-from airflow.hooks.dbapi_hook import DbApiHook
-
-
-class PinotDbApiHook(DbApiHook):
- """
- Connect to pinot db(https://github.com/linkedin/pinot) to issue pql
- """
- conn_name_attr = 'pinot_broker_conn_id'
- default_conn_name = 'pinot_broker_default'
- supports_autocommit = False
-
- def __init__(self, *args, **kwargs):
- super(PinotDbApiHook, self).__init__(*args, **kwargs)
-
- def get_conn(self):
- """
- Establish a connection to pinot broker through pinot dbqpi.
- """
- conn = self.get_connection(self.pinot_broker_conn_id)
- pinot_broker_conn = connect(
- host=conn.host,
- port=conn.port,
- path=conn.extra_dejson.get('endpoint', '/pql'),
- scheme=conn.extra_dejson.get('schema', 'http')
- )
- self.log.info('Get the connection to pinot '
- 'broker on {host}'.format(host=conn.host))
- return pinot_broker_conn
-
- def get_uri(self):
- """
- Get the connection uri for pinot broker.
-
- e.g: http://localhost:9000/pql
- """
- conn = self.get_connection(getattr(self, self.conn_name_attr))
- host = conn.host
- if conn.port is not None:
- host += ':{port}'.format(port=conn.port)
- conn_type = 'http' if not conn.conn_type else conn.conn_type
- endpoint = conn.extra_dejson.get('endpoint', 'pql')
- return '{conn_type}://{host}/{endpoint}'.format(
- conn_type=conn_type, host=host, endpoint=endpoint)
-
- def get_records(self, sql):
- """
- Executes the sql and returns a set of records.
-
- :param sql: the sql statement to be executed (str) or a list of
- sql statements to execute
- :type sql: str
- """
- if six.PY2:
- sql = sql.encode('utf-8')
-
- with self.get_conn() as cur:
- cur.execute(sql)
- return cur.fetchall()
-
- def get_first(self, sql):
- """
- Executes the sql and returns the first resulting row.
-
- :param sql: the sql statement to be executed (str) or a list of
- sql statements to execute
- :type sql: str or list
- """
- if six.PY2:
- sql = sql.encode('utf-8')
-
- with self.get_conn() as cur:
- cur.execute(sql)
- return cur.fetchone()
-
- def set_autocommit(self, conn, autocommit):
- raise NotImplementedError()
-
- def get_pandas_df(self, sql, parameters=None):
- raise NotImplementedError()
-
- def insert_rows(self, table, rows, target_fields=None, commit_every=1000):
- raise NotImplementedError()
diff --git a/airflow/contrib/hooks/qubole_check_hook.py b/airflow/contrib/hooks/qubole_check_hook.py
deleted file mode 100644
index 303c19b648..0000000000
--- a/airflow/contrib/hooks/qubole_check_hook.py
+++ /dev/null
@@ -1,117 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.contrib.hooks.qubole_hook import QuboleHook
-from airflow.exceptions import AirflowException
-from qds_sdk.commands import Command
-
-try:
- from cStringIO import StringIO
-except ImportError:
- from io import StringIO
-
-
-COL_DELIM = '\t'
-ROW_DELIM = '\r\n'
-
-
-def isint(value):
- try:
- int(value)
- return True
- except ValueError:
- return False
-
-
-def isfloat(value):
- try:
- float(value)
- return True
- except ValueError:
- return False
-
-
-def isbool(value):
- try:
- if value.lower() in ["true", "false"]:
- return True
- except ValueError:
- return False
-
-
-def parse_first_row(row_list):
- record_list = []
- first_row = row_list[0] if row_list else ""
-
- for col_value in first_row.split(COL_DELIM):
- if isint(col_value):
- col_value = int(col_value)
- elif isfloat(col_value):
- col_value = float(col_value)
- elif isbool(col_value):
- col_value = (col_value.lower() == "true")
- record_list.append(col_value)
-
- return record_list
-
-
-class QuboleCheckHook(QuboleHook):
- def __init__(self, context, *args, **kwargs):
- super(QuboleCheckHook, self).__init__(*args, **kwargs)
- self.results_parser_callable = parse_first_row
- if 'results_parser_callable' in kwargs and \
- kwargs['results_parser_callable'] is not None:
- if not callable(kwargs['results_parser_callable']):
- raise AirflowException('`results_parser_callable` param must be callable')
- self.results_parser_callable = kwargs['results_parser_callable']
- self.context = context
-
- @staticmethod
- def handle_failure_retry(context):
- ti = context['ti']
- cmd_id = ti.xcom_pull(key='qbol_cmd_id', task_ids=ti.task_id)
-
- if cmd_id is not None:
- cmd = Command.find(cmd_id)
- if cmd is not None:
- if cmd.status == 'running':
- log = LoggingMixin().log
- log.info('Cancelling the Qubole Command Id: %s', cmd_id)
- cmd.cancel()
-
- def get_first(self, sql):
- self.execute(context=self.context)
- query_result = self.get_query_results()
- row_list = list(filter(None, query_result.split(ROW_DELIM)))
- record_list = self.results_parser_callable(row_list)
- return record_list
-
- def get_query_results(self):
- log = LoggingMixin().log
- if self.cmd is not None:
- cmd_id = self.cmd.id
- log.info("command id: " + str(cmd_id))
- query_result_buffer = StringIO()
- self.cmd.get_results(fp=query_result_buffer, inline=True, delim=COL_DELIM)
- query_result = query_result_buffer.getvalue()
- query_result_buffer.close()
- return query_result
- else:
- log.info("Qubole command not found")
diff --git a/airflow/contrib/hooks/qubole_hook.py b/airflow/contrib/hooks/qubole_hook.py
deleted file mode 100755
index fb45862e88..0000000000
--- a/airflow/contrib/hooks/qubole_hook.py
+++ /dev/null
@@ -1,214 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 os
-import time
-import datetime
-import six
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow import configuration
-from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.state import State
-
-from qds_sdk.qubole import Qubole
-from qds_sdk.commands import Command, HiveCommand, PrestoCommand, HadoopCommand, \
- PigCommand, ShellCommand, SparkCommand, DbTapQueryCommand, DbExportCommand, \
- DbImportCommand
-
-
-COMMAND_CLASSES = {
- "hivecmd": HiveCommand,
- "prestocmd": PrestoCommand,
- "hadoopcmd": HadoopCommand,
- "shellcmd": ShellCommand,
- "pigcmd": PigCommand,
- "sparkcmd": SparkCommand,
- "dbtapquerycmd": DbTapQueryCommand,
- "dbexportcmd": DbExportCommand,
- "dbimportcmd": DbImportCommand
-}
-
-HYPHEN_ARGS = ['cluster_label', 'app_id', 'note_id']
-
-POSITIONAL_ARGS = ['sub_command', 'parameters']
-
-COMMAND_ARGS = {
- "hivecmd": ['query', 'script_location', 'macros', 'tags', 'sample_size',
- 'cluster_label', 'name'],
- 'prestocmd': ['query', 'script_location', 'macros', 'tags', 'cluster_label', 'name'],
- 'hadoopcmd': ['sub_command', 'tags', 'cluster_label', 'name'],
- 'shellcmd': ['script', 'script_location', 'files', 'archives', 'parameters', 'tags',
- 'cluster_label', 'name'],
- 'pigcmd': ['script', 'script_location', 'parameters', 'tags', 'cluster_label',
- 'name'],
- 'dbtapquerycmd': ['db_tap_id', 'query', 'macros', 'tags', 'name'],
- 'sparkcmd': ['program', 'cmdline', 'sql', 'script_location', 'macros', 'tags',
- 'cluster_label', 'language', 'app_id', 'name', 'arguments', 'note_id',
- 'user_program_arguments'],
- 'dbexportcmd': ['mode', 'hive_table', 'partition_spec', 'dbtap_id', 'db_table',
- 'db_update_mode', 'db_update_keys', 'export_dir',
- 'fields_terminated_by', 'tags', 'name'],
- 'dbimportcmd': ['mode', 'hive_table', 'dbtap_id', 'db_table', 'where_clause',
- 'parallelism', 'extract_query', 'boundary_query', 'split_column',
- 'tags', 'name']
-}
-
-
-class QuboleHook(BaseHook, LoggingMixin):
- def __init__(self, *args, **kwargs):
- conn = self.get_connection(kwargs['qubole_conn_id'])
- Qubole.configure(api_token=conn.password, api_url=conn.host)
- self.task_id = kwargs['task_id']
- self.dag_id = kwargs['dag'].dag_id
- self.kwargs = kwargs
- self.cls = COMMAND_CLASSES[self.kwargs['command_type']]
- self.cmd = None
-
- @staticmethod
- def handle_failure_retry(context):
- ti = context['ti']
- cmd_id = ti.xcom_pull(key='qbol_cmd_id', task_ids=ti.task_id)
-
- if cmd_id is not None:
- cmd = Command.find(cmd_id)
- if cmd is not None:
- log = LoggingMixin().log
- if cmd.status == 'done':
- log.info('Command ID: %s has been succeeded, hence marking this '
- 'TI as Success.', cmd_id)
- ti.state = State.SUCCESS
- elif cmd.status == 'running':
- log.info('Cancelling the Qubole Command Id: %s', cmd_id)
- cmd.cancel()
-
- def execute(self, context):
- args = self.cls.parse(self.create_cmd_args(context))
- self.cmd = self.cls.create(**args)
- context['task_instance'].xcom_push(key='qbol_cmd_id', value=self.cmd.id)
- self.log.info(
- "Qubole command created with Id: %s and Status: %s",
- self.cmd.id, self.cmd.status
- )
-
- while not Command.is_done(self.cmd.status):
- time.sleep(Qubole.poll_interval)
- self.cmd = self.cls.find(self.cmd.id)
- self.log.info("Command Id: %s and Status: %s", self.cmd.id, self.cmd.status)
-
- if 'fetch_logs' in self.kwargs and self.kwargs['fetch_logs'] is True:
- self.log.info("Logs for Command Id: %s \n%s", self.cmd.id, self.cmd.get_log())
-
- if self.cmd.status != 'done':
- raise AirflowException('Command Id: {0} failed with Status: {1}'.format(
- self.cmd.id, self.cmd.status))
-
- def kill(self, ti):
- """
- Kill (cancel) a Qubole commmand
- :param ti: Task Instance of the dag, used to determine the Quboles command id
- :return: response from Qubole
- """
- if self.cmd is None:
- cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=ti.task_id)
- self.cmd = self.cls.find(cmd_id)
- if self.cls and self.cmd:
- self.log.info('Sending KILL signal to Qubole Command Id: %s', self.cmd.id)
- self.cmd.cancel()
-
- def get_results(self, ti=None, fp=None, inline=True, delim=None, fetch=True):
- """
- Get results (or just s3 locations) of a command from Qubole and save into a file
- :param ti: Task Instance of the dag, used to determine the Quboles command id
- :param fp: Optional file pointer, will create one and return if None passed
- :param inline: True to download actual results, False to get s3 locations only
- :param delim: Replaces the CTL-A chars with the given delim, defaults to ','
- :param fetch: when inline is True, get results directly from s3 (if large)
- :return: file location containing actual results or s3 locations of results
- """
- if fp is None:
- iso = datetime.datetime.utcnow().isoformat()
- logpath = os.path.expanduser(
- configuration.conf.get('core', 'BASE_LOG_FOLDER')
- )
- resultpath = logpath + '/' + self.dag_id + '/' + self.task_id + '/results'
- configuration.mkdir_p(resultpath)
- fp = open(resultpath + '/' + iso, 'wb')
-
- if self.cmd is None:
- cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
- self.cmd = self.cls.find(cmd_id)
-
- self.cmd.get_results(fp, inline, delim, fetch)
- fp.flush()
- fp.close()
- return fp.name
-
- def get_log(self, ti):
- """
- Get Logs of a command from Qubole
- :param ti: Task Instance of the dag, used to determine the Quboles command id
- :return: command log as text
- """
- if self.cmd is None:
- cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
- Command.get_log_id(self.cls, cmd_id)
-
- def get_jobs_id(self, ti):
- """
- Get jobs associated with a Qubole commands
- :param ti: Task Instance of the dag, used to determine the Quboles command id
- :return: Job informations assoiciated with command
- """
- if self.cmd is None:
- cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
- Command.get_jobs_id(self.cls, cmd_id)
-
- def create_cmd_args(self, context):
- args = []
- cmd_type = self.kwargs['command_type']
- inplace_args = None
- tags = set([self.dag_id, self.task_id, context['run_id']])
-
- for k, v in self.kwargs.items():
- if k in COMMAND_ARGS[cmd_type]:
- if k in HYPHEN_ARGS:
- args.append("--{0}={1}".format(k.replace('_', '-'), v))
- elif k in POSITIONAL_ARGS:
- inplace_args = v
- elif k == 'tags':
- if isinstance(v, six.string_types):
- tags.add(v)
- elif isinstance(v, (list, tuple)):
- for val in v:
- tags.add(val)
- else:
- args.append("--{0}={1}".format(k, v))
-
- if k == 'notify' and v is True:
- args.append("--notify")
-
- args.append("--tags={0}".format(','.join(filter(None, tags))))
-
- if inplace_args is not None:
- args += inplace_args.split(' ')
-
- return args
diff --git a/airflow/contrib/hooks/redis_hook.py b/airflow/contrib/hooks/redis_hook.py
deleted file mode 100644
index 1de75dbca9..0000000000
--- a/airflow/contrib/hooks/redis_hook.py
+++ /dev/null
@@ -1,93 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-
-"""
-RedisHook module
-"""
-from redis import StrictRedis
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class RedisHook(BaseHook, LoggingMixin):
- """
- Hook to interact with Redis database
- """
- def __init__(self, redis_conn_id='redis_default'):
- """
- Prepares hook to connect to a Redis database.
-
- :param conn_id: the name of the connection that has the parameters
- we need to connect to Redis.
- """
- self.redis_conn_id = redis_conn_id
- self.client = None
- conn = self.get_connection(self.redis_conn_id)
- self.host = conn.host
- self.port = int(conn.port)
- self.password = conn.password
- self.db = int(conn.extra_dejson.get('db', 0))
-
- self.log.debug(
- '''Connection "{conn}":
- \thost: {host}
- \tport: {port}
- \textra: {extra}
- '''.format(
- conn=self.redis_conn_id,
- host=self.host,
- port=self.port,
- extra=conn.extra_dejson
- )
- )
-
- def get_conn(self):
- """
- Returns a Redis connection.
- """
- if not self.client:
- self.log.debug(
- 'generating Redis client for conn_id "%s" on %s:%s:%s',
- self.redis_conn_id, self.host, self.port, self.db
- )
- try:
- self.client = StrictRedis(
- host=self.host,
- port=self.port,
- password=self.password,
- db=self.db)
- except Exception as general_error:
- raise AirflowException(
- 'Failed to create Redis client, error: {error}'.format(
- error=str(general_error)
- )
- )
-
- return self.client
-
- def key_exists(self, key):
- """
- Checks if a key exists in Redis database
-
- :param key: The key to check the existence.
- :type key: string
- """
- return self.get_conn().exists(key)
diff --git a/airflow/contrib/hooks/redshift_hook.py b/airflow/contrib/hooks/redshift_hook.py
deleted file mode 100644
index fcce683f15..0000000000
--- a/airflow/contrib/hooks/redshift_hook.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class RedshiftHook(AwsHook):
- """
- Interact with AWS Redshift, using the boto3 library
- """
- def get_conn(self):
- return self.get_client_type('redshift')
-
- # TODO: Wrap create_cluster_snapshot
- def cluster_status(self, cluster_identifier):
- """
- Return status of a cluster
-
- :param cluster_identifier: unique identifier of a cluster
- :type cluster_identifier: str
- """
- conn = self.get_conn()
- try:
- response = conn.describe_clusters(
- ClusterIdentifier=cluster_identifier)['Clusters']
- return response[0]['ClusterStatus'] if response else None
- except conn.exceptions.ClusterNotFoundFault:
- return 'cluster_not_found'
-
- def delete_cluster(
- self,
- cluster_identifier,
- skip_final_cluster_snapshot=True,
- final_cluster_snapshot_identifier=''):
- """
- Delete a cluster and optionally create a snapshot
-
- :param cluster_identifier: unique identifier of a cluster
- :type cluster_identifier: str
- :param skip_final_cluster_snapshot: determines cluster snapshot creation
- :type skip_final_cluster_snapshot: bool
- :param final_cluster_snapshot_identifier: name of final cluster snapshot
- :type final_cluster_snapshot_identifier: str
- """
- response = self.get_conn().delete_cluster(
- ClusterIdentifier=cluster_identifier,
- SkipFinalClusterSnapshot=skip_final_cluster_snapshot,
- FinalClusterSnapshotIdentifier=final_cluster_snapshot_identifier
- )
- return response['Cluster'] if response['Cluster'] else None
-
- def describe_cluster_snapshots(self, cluster_identifier):
- """
- Gets a list of snapshots for a cluster
-
- :param cluster_identifier: unique identifier of a cluster
- :type cluster_identifier: str
- """
- response = self.get_conn().describe_cluster_snapshots(
- ClusterIdentifier=cluster_identifier
- )
- if 'Snapshots' not in response:
- return None
- snapshots = response['Snapshots']
- snapshots = filter(lambda x: x['Status'], snapshots)
- snapshots.sort(key=lambda x: x['SnapshotCreateTime'], reverse=True)
- return snapshots
-
- def restore_from_cluster_snapshot(self, cluster_identifier, snapshot_identifier):
- """
- Restores a cluster from its snapshot
-
- :param cluster_identifier: unique identifier of a cluster
- :type cluster_identifier: str
- :param snapshot_identifier: unique identifier for a snapshot of a cluster
- :type snapshot_identifier: str
- """
- response = self.get_conn().restore_from_cluster_snapshot(
- ClusterIdentifier=cluster_identifier,
- SnapshotIdentifier=snapshot_identifier
- )
- return response['Cluster'] if response['Cluster'] else None
-
- def create_cluster_snapshot(self, snapshot_identifier, cluster_identifier):
- """
- Creates a snapshot of a cluster
-
- :param snapshot_identifier: unique identifier for a snapshot of a cluster
- :type snapshot_identifier: str
- :param cluster_identifier: unique identifier of a cluster
- :type cluster_identifier: str
- """
- response = self.get_conn().create_cluster_snapshot(
- SnapshotIdentifier=snapshot_identifier,
- ClusterIdentifier=cluster_identifier,
- )
- return response['Snapshot'] if response['Snapshot'] else None
diff --git a/airflow/contrib/hooks/salesforce_hook.py b/airflow/contrib/hooks/salesforce_hook.py
deleted file mode 100644
index ee18b353d2..0000000000
--- a/airflow/contrib/hooks/salesforce_hook.py
+++ /dev/null
@@ -1,324 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-#
-"""
-This module contains a Salesforce Hook
-which allows you to connect to your Salesforce instance,
-retrieve data from it, and write that data to a file
-for other uses.
-
-NOTE: this hook also relies on the simple_salesforce package:
- https://github.com/simple-salesforce/simple-salesforce
-"""
-from simple_salesforce import Salesforce
-from airflow.hooks.base_hook import BaseHook
-
-import json
-
-import pandas as pd
-import time
-
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class SalesforceHook(BaseHook, LoggingMixin):
- def __init__(
- self,
- conn_id,
- *args,
- **kwargs
- ):
- """
- Create new connection to Salesforce
- and allows you to pull data out of SFDC and save it to a file.
-
- You can then use that file with other
- Airflow operators to move the data into another data source
-
- :param conn_id: the name of the connection that has the parameters
- we need to connect to Salesforce.
- The conenction shoud be type `http` and include a
- user's security token in the `Extras` field.
- .. note::
- For the HTTP connection type, you can include a
- JSON structure in the `Extras` field.
- We need a user's security token to connect to Salesforce.
- So we define it in the `Extras` field as:
- `{"security_token":"YOUR_SECRUITY_TOKEN"}`
- """
- self.conn_id = conn_id
- self._args = args
- self._kwargs = kwargs
-
- # get the connection parameters
- self.connection = self.get_connection(conn_id)
- self.extras = self.connection.extra_dejson
-
- def sign_in(self):
- """
- Sign into Salesforce.
-
- If we have already signed it, this will just return the original object
- """
- if hasattr(self, 'sf'):
- return self.sf
-
- # connect to Salesforce
- sf = Salesforce(
- username=self.connection.login,
- password=self.connection.password,
- security_token=self.extras['security_token'],
- instance_url=self.connection.host,
- sandbox=self.extras.get('sandbox', False)
- )
- self.sf = sf
- return sf
-
- def make_query(self, query):
- """
- Make a query to Salesforce. Returns result in dictionary
-
- :param query: The query to make to Salesforce
- """
- self.sign_in()
-
- self.log.info("Querying for all objects")
- query = self.sf.query_all(query)
-
- self.log.info(
- "Received results: Total size: %s; Done: %s",
- query['totalSize'], query['done']
- )
-
- query = json.loads(json.dumps(query))
- return query
-
- def describe_object(self, obj):
- """
- Get the description of an object from Salesforce.
-
- This description is the object's schema
- and some extra metadata that Salesforce stores for each object
-
- :param obj: Name of the Salesforce object
- that we are getting a description of.
- """
- self.sign_in()
-
- return json.loads(json.dumps(self.sf.__getattr__(obj).describe()))
-
- def get_available_fields(self, obj):
- """
- Get a list of all available fields for an object.
-
- This only returns the names of the fields.
- """
- self.sign_in()
-
- desc = self.describe_object(obj)
-
- return [f['name'] for f in desc['fields']]
-
- def _build_field_list(self, fields):
- # join all of the fields in a comma separated list
- return ",".join(fields)
-
- def get_object_from_salesforce(self, obj, fields):
- """
- Get all instances of the `object` from Salesforce.
- For each model, only get the fields specified in fields.
-
- All we really do underneath the hood is run:
- SELECT <fields> FROM <obj>;
- """
- field_string = self._build_field_list(fields)
-
- query = "SELECT {0} FROM {1}".format(field_string, obj)
- self.log.info(
- "Making query to Salesforce: %s",
- query if len(query) < 30 else " ... ".join([query[:15], query[-15:]])
- )
- return self.make_query(query)
-
- @classmethod
- def _to_timestamp(cls, col):
- """
- Convert a column of a dataframe to UNIX timestamps if applicable
-
- :param col: A Series object representing a column of a dataframe.
- """
- # try and convert the column to datetimes
- # the column MUST have a four digit year somewhere in the string
- # there should be a better way to do this,
- # but just letting pandas try and convert every column without a format
- # caused it to convert floats as well
- # For example, a column of integers
- # between 0 and 10 are turned into timestamps
- # if the column cannot be converted,
- # just return the original column untouched
- try:
- col = pd.to_datetime(col)
- except ValueError:
- log = LoggingMixin().log
- log.warning(
- "Could not convert field to timestamps: %s", col.name
- )
- return col
-
- # now convert the newly created datetimes into timestamps
- # we have to be careful here
- # because NaT cannot be converted to a timestamp
- # so we have to return NaN
- converted = []
- for i in col:
- try:
- converted.append(i.timestamp())
- except ValueError:
- converted.append(pd.np.NaN)
- except AttributeError:
- converted.append(pd.np.NaN)
-
- # return a new series that maintains the same index as the original
- return pd.Series(converted, index=col.index)
-
- def write_object_to_file(
- self,
- query_results,
- filename,
- fmt="csv",
- coerce_to_timestamp=False,
- record_time_added=False
- ):
- """
- Write query results to file.
-
- Acceptable formats are:
- - csv:
- comma-separated-values file. This is the default format.
- - json:
- JSON array. Each element in the array is a different row.
- - ndjson:
- JSON array but each element is new-line delimited
- instead of comma delimited like in `json`
-
- This requires a significant amount of cleanup.
- Pandas doesn't handle output to CSV and json in a uniform way.
- This is especially painful for datetime types.
- Pandas wants to write them as strings in CSV,
- but as millisecond Unix timestamps.
-
- By default, this function will try and leave all values as
- they are represented in Salesforce.
- You use the `coerce_to_timestamp` flag to force all datetimes
- to become Unix timestamps (UTC).
- This is can be greatly beneficial as it will make all of your
- datetime fields look the same,
- and makes it easier to work with in other database environments
-
- :param query_results: the results from a SQL query
- :param filename: the name of the file where the data
- should be dumped to
- :param fmt: the format you want the output in.
- *Default:* csv.
- :param coerce_to_timestamp: True if you want all datetime fields to be
- converted into Unix timestamps.
- False if you want them to be left in the
- same format as they were in Salesforce.
- Leaving the value as False will result
- in datetimes being strings.
- *Defaults to False*
- :param record_time_added: *(optional)* True if you want to add a
- Unix timestamp field to the resulting data
- that marks when the data
- was fetched from Salesforce.
- *Default: False*.
- """
- fmt = fmt.lower()
- if fmt not in ['csv', 'json', 'ndjson']:
- raise ValueError("Format value is not recognized: {0}".format(fmt))
-
- # this line right here will convert all integers to floats if there are
- # any None/np.nan values in the column
- # that's because None/np.nan cannot exist in an integer column
- # we should write all of our timestamps as FLOATS in our final schema
- df = pd.DataFrame.from_records(query_results, exclude=["attributes"])
-
- df.columns = [c.lower() for c in df.columns]
-
- # convert columns with datetime strings to datetimes
- # not all strings will be datetimes, so we ignore any errors that occur
- # we get the object's definition at this point and only consider
- # features that are DATE or DATETIME
- if coerce_to_timestamp and df.shape[0] > 0:
- # get the object name out of the query results
- # it's stored in the "attributes" dictionary
- # for each returned record
- object_name = query_results[0]['attributes']['type']
-
- self.log.info("Coercing timestamps for: %s", object_name)
-
- schema = self.describe_object(object_name)
-
- # possible columns that can be convereted to timestamps
- # are the ones that are either date or datetime types
- # strings are too general and we risk unintentional conversion
- possible_timestamp_cols = [
- i['name'].lower()
- for i in schema['fields']
- if i['type'] in ["date", "datetime"] and
- i['name'].lower() in df.columns
- ]
- df[possible_timestamp_cols] = df[possible_timestamp_cols].apply(
- lambda x: self._to_timestamp(x)
- )
-
- if record_time_added:
- fetched_time = time.time()
- df["time_fetched_from_salesforce"] = fetched_time
-
- # write the CSV or JSON file depending on the option
- # NOTE:
- # datetimes here are an issue.
- # There is no good way to manage the difference
- # for to_json, the options are an epoch or a ISO string
- # but for to_csv, it will be a string output by datetime
- # For JSON we decided to output the epoch timestamp in seconds
- # (as is fairly standard for JavaScript)
- # And for csv, we do a string
- if fmt == "csv":
- # there are also a ton of newline objects
- # that mess up our ability to write to csv
- # we remove these newlines so that the output is a valid CSV format
- self.log.info("Cleaning data and writing to CSV")
- possible_strings = df.columns[df.dtypes == "object"]
- df[possible_strings] = df[possible_strings].apply(
- lambda x: x.str.replace("\r\n", "")
- )
- df[possible_strings] = df[possible_strings].apply(
- lambda x: x.str.replace("\n", "")
- )
-
- # write the dataframe
- df.to_csv(filename, index=False)
- elif fmt == "json":
- df.to_json(filename, "records", date_unit="s")
- elif fmt == "ndjson":
- df.to_json(filename, "records", lines=True, date_unit="s")
-
- return df
diff --git a/airflow/contrib/hooks/segment_hook.py b/airflow/contrib/hooks/segment_hook.py
deleted file mode 100644
index 874d35d074..0000000000
--- a/airflow/contrib/hooks/segment_hook.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-#
-"""
-This module contains a Segment Hook
-which allows you to connect to your Segment account,
-retrieve data from it or write to that file.
-
-NOTE: this hook also relies on the Segment analytics package:
- https://github.com/segmentio/analytics-python
-"""
-import analytics
-from airflow.hooks.base_hook import BaseHook
-from airflow.exceptions import AirflowException
-
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class SegmentHook(BaseHook, LoggingMixin):
- def __init__(
- self,
- segment_conn_id='segment_default',
- segment_debug_mode=False,
- *args,
- **kwargs
- ):
- """
- Create new connection to Segment
- and allows you to pull data out of Segment or write to it.
-
- You can then use that file with other
- Airflow operators to move the data around or interact with segment.
-
- :param segment_conn_id: the name of the connection that has the parameters
- we need to connect to Segment.
- The connection should be type `json` and include a
- write_key security token in the `Extras` field.
- :type segment_conn_id: str
- :param segment_debug_mode: Determines whether Segment should run in debug mode.
- Defaults to False
- :type segment_debug_mode: boolean
- .. note::
- You must include a JSON structure in the `Extras` field.
- We need a user's security token to connect to Segment.
- So we define it in the `Extras` field as:
- `{"write_key":"YOUR_SECURITY_TOKEN"}`
- """
- self.segment_conn_id = segment_conn_id
- self.segment_debug_mode = segment_debug_mode
- self._args = args
- self._kwargs = kwargs
-
- # get the connection parameters
- self.connection = self.get_connection(self.segment_conn_id)
- self.extras = self.connection.extra_dejson
- self.write_key = self.extras.get('write_key')
- if self.write_key is None:
- raise AirflowException('No Segment write key provided')
-
- def get_conn(self):
- self.log.info('Setting write key for Segment analytics connection')
- analytics.debug = self.segment_debug_mode
- if self.segment_debug_mode:
- self.log.info('Setting Segment analytics connection to debug mode')
- analytics.on_error = self.on_error
- analytics.write_key = self.write_key
- return analytics
-
- def on_error(self, error, items):
- """
- Handles error callbacks when using Segment with segment_debug_mode set to True
- """
- self.log.error('Encountered Segment error: {segment_error} with '
- 'items: {with_items}'.format(segment_error=error,
- with_items=items))
- raise AirflowException('Segment error: {}'.format(error))
diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
deleted file mode 100644
index 33c8b19e0a..0000000000
--- a/airflow/contrib/hooks/sftp_hook.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 stat
-import pysftp
-import logging
-import datetime
-from airflow.hooks.base_hook import BaseHook
-
-
-class SFTPHook(BaseHook):
- """
- Interact with SFTP. Aims to be interchangeable with FTPHook.
-
- Pitfalls: - In contrast with FTPHook describe_directory only returns size, type and
- modify. It doesn't return unix.owner, unix.mode, perm, unix.group and
- unique.
- - retrieve_file and store_file only take a local full path and not a
- buffer.
- - If no mode is passed to create_directory it will be created with 777
- permissions.
-
- Errors that may occur throughout but should be handled downstream.
- """
-
- def __init__(self, ftp_conn_id='sftp_default'):
- self.ftp_conn_id = ftp_conn_id
- self.conn = None
-
- def get_conn(self):
- """
- Returns an SFTP connection object
- """
- if self.conn is None:
- params = self.get_connection(self.ftp_conn_id)
- cnopts = pysftp.CnOpts()
- if ('ignore_hostkey_verification' in params.extra_dejson and
- params.extra_dejson['ignore_hostkey_verification']):
- cnopts.hostkeys = None
- conn_params = {
- 'host': params.host,
- 'port': params.port,
- 'username': params.login,
- 'cnopts': cnopts
- }
- if params.password is not None:
- conn_params['password'] = params.password
- if 'private_key' in params.extra_dejson:
- conn_params['private_key'] = params.extra_dejson['private_key']
- if 'private_key_pass' in params.extra_dejson:
- conn_params['private_key_pass'] = params.extra_dejson['private_key_pass']
- self.conn = pysftp.Connection(**conn_params)
- return self.conn
-
- def close_conn(self):
- """
- Closes the connection. An error will occur if the
- connection wasnt ever opened.
- """
- conn = self.conn
- conn.close()
- self.conn = None
-
- def describe_directory(self, path):
- """
- Returns a dictionary of {filename: {attributes}} for all files
- on the remote system (where the MLSD command is supported).
- :param path: full path to the remote directory
- :type path: str
- """
- conn = self.get_conn()
- flist = conn.listdir_attr(path)
- files = {}
- for f in flist:
- modify = datetime.datetime.fromtimestamp(
- f.st_mtime).strftime('%Y%m%d%H%M%S')
- files[f.filename] = {
- 'size': f.st_size,
- 'type': 'dir' if stat.S_ISDIR(f.st_mode) else 'file',
- 'modify': modify}
- return files
-
- def list_directory(self, path):
- """
- Returns a list of files on the remote system.
- :param path: full path to the remote directory to list
- :type path: str
- """
- conn = self.get_conn()
- files = conn.listdir(path)
- return files
-
- def create_directory(self, path, mode=777):
- """
- Creates a directory on the remote system.
- :param path: full path to the remote directory to create
- :type path: str
- :param mode: int representation of octal mode for directory
- """
- conn = self.get_conn()
- conn.mkdir(path, mode)
-
- def delete_directory(self, path):
- """
- Deletes a directory on the remote system.
- :param path: full path to the remote directory to delete
- :type path: str
- """
- conn = self.get_conn()
- conn.rmdir(path)
-
- def retrieve_file(self, remote_full_path, local_full_path):
- """
- Transfers the remote file to a local location.
- If local_full_path is a string path, the file will be put
- at that location
- :param remote_full_path: full path to the remote file
- :type remote_full_path: str
- :param local_full_path: full path to the local file
- :type local_full_path: str
- """
- conn = self.get_conn()
- logging.info('Retrieving file from FTP: {}'.format(remote_full_path))
- conn.get(remote_full_path, local_full_path)
- logging.info('Finished retrieving file from FTP: {}'.format(
- remote_full_path))
-
- def store_file(self, remote_full_path, local_full_path):
- """
- Transfers a local file to the remote location.
- If local_full_path_or_buffer is a string path, the file will be read
- from that location
- :param remote_full_path: full path to the remote file
- :type remote_full_path: str
- :param local_full_path: full path to the local file
- :type local_full_path: str
- """
- conn = self.get_conn()
- conn.put(local_full_path, remote_full_path)
-
- def delete_file(self, path):
- """
- Removes a file on the FTP Server
- :param path: full path to the remote file
- :type path: str
- """
- conn = self.get_conn()
- conn.remove(path)
-
- def get_mod_time(self, path):
- conn = self.get_conn()
- ftp_mdtm = conn.stat(path).st_mtime
- return datetime.datetime.fromtimestamp(ftp_mdtm).strftime('%Y%m%d%H%M%S')
diff --git a/airflow/contrib/hooks/slack_webhook_hook.py b/airflow/contrib/hooks/slack_webhook_hook.py
deleted file mode 100644
index 670d40105e..0000000000
--- a/airflow/contrib/hooks/slack_webhook_hook.py
+++ /dev/null
@@ -1,129 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-
-from airflow.hooks.http_hook import HttpHook
-from airflow.exceptions import AirflowException
-
-
-class SlackWebhookHook(HttpHook):
- """
- This hook allows you to post messages to Slack using incoming webhooks.
- Takes both Slack webhook token directly and connection that has Slack webhook token.
- If both supplied, Slack webhook token will be used.
-
- Each Slack webhook token can be pre-configured to use a specific channel, username and
- icon. You can override these defaults in this hook.
-
- :param http_conn_id: connection that has Slack webhook token in the extra field
- :type http_conn_id: str
- :param webhook_token: Slack webhook token
- :type webhook_token: str
- :param message: The message you want to send on Slack
- :type message: str
- :param channel: The channel the message should be posted to
- :type channel: str
- :param username: The username to post to slack with
- :type username: str
- :param icon_emoji: The emoji to use as icon for the user posting to Slack
- :type icon_emoji: str
- :param link_names: Whether or not to find and link channel and usernames in your
- message
- :type link_names: bool
- :param proxy: Proxy to use to make the Slack webhook call
- :type proxy: str
- """
- def __init__(self,
- http_conn_id=None,
- webhook_token=None,
- message="",
- channel=None,
- username=None,
- icon_emoji=None,
- link_names=False,
- proxy=None,
- *args,
- **kwargs
- ):
- super(SlackWebhookHook, self).__init__(*args, **kwargs)
- self.http_conn_id = http_conn_id
- self.webhook_token = self._get_token(webhook_token, http_conn_id)
- self.message = message
- self.channel = channel
- self.username = username
- self.icon_emoji = icon_emoji
- self.link_names = link_names
- self.proxy = proxy
-
- def _get_token(self, token, http_conn_id):
- """
- Given either a manually set token or a conn_id, return the webhook_token to use
- :param token: The manually provided token
- :param conn_id: The conn_id provided
- :return: webhook_token (str) to use
- """
- if token:
- return token
- elif http_conn_id:
- conn = self.get_connection(http_conn_id)
- extra = conn.extra_dejson
- return extra.get('webhook_token', '')
- else:
- raise AirflowException('Cannot get token: No valid Slack '
- 'webhook token nor conn_id supplied')
-
- def _build_slack_message(self):
- """
- Construct the Slack message. All relevant parameters are combined here to a valid
- Slack json message
- :return: Slack message (str) to send
- """
- cmd = {}
-
- if self.channel:
- cmd['channel'] = self.channel
- if self.username:
- cmd['username'] = self.username
- if self.icon_emoji:
- cmd['icon_emoji'] = self.icon_emoji
- if self.link_names:
- cmd['link_names'] = 1
-
- # there should always be a message to post ;-)
- cmd['text'] = self.message
- return json.dumps(cmd)
-
- def execute(self):
- """
- Remote Popen (actually execute the slack webhook call)
-
- :param cmd: command to remotely execute
- :param kwargs: extra arguments to Popen (see subprocess.Popen)
- """
- proxies = {}
- if self.proxy:
- # we only need https proxy for Slack, as the endpoint is https
- proxies = {'https': self.proxy}
-
- slack_message = self._build_slack_message()
- self.run(endpoint=self.webhook_token,
- data=slack_message,
- headers={'Content-type': 'application/json'},
- extra_options={'proxies': proxies})
diff --git a/airflow/contrib/hooks/snowflake_hook.py b/airflow/contrib/hooks/snowflake_hook.py
deleted file mode 100644
index 4027a9a78b..0000000000
--- a/airflow/contrib/hooks/snowflake_hook.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 snowflake.connector
-
-from airflow.hooks.dbapi_hook import DbApiHook
-
-
-class SnowflakeHook(DbApiHook):
- """
- Interact with Snowflake.
-
- get_sqlalchemy_engine() depends on snowflake-sqlalchemy
-
- """
-
- conn_name_attr = 'snowflake_conn_id'
- default_conn_name = 'snowflake_default'
- supports_autocommit = True
-
- def __init__(self, *args, **kwargs):
- super(SnowflakeHook, self).__init__(*args, **kwargs)
- self.account = kwargs.pop("account", None)
- self.warehouse = kwargs.pop("warehouse", None)
- self.database = kwargs.pop("database", None)
-
- def _get_conn_params(self):
- """
- one method to fetch connection params as a dict
- used in get_uri() and get_connection()
- """
- conn = self.get_connection(self.snowflake_conn_id)
- account = conn.extra_dejson.get('account', None)
- warehouse = conn.extra_dejson.get('warehouse', None)
- database = conn.extra_dejson.get('database', None)
-
- conn_config = {
- "user": conn.login,
- "password": conn.password or '',
- "schema": conn.schema or '',
- "database": self.database or database or '',
- "account": self.account or account or '',
- "warehouse": self.warehouse or warehouse or ''
- }
- return conn_config
-
- def get_uri(self):
- """
- override DbApiHook get_uri method for get_sqlalchemy_engine()
- """
- conn_config = self._get_conn_params()
- uri = 'snowflake://{user}:{password}@{account}/{database}/'
- uri += '{schema}?warehouse={warehouse}'
- return uri.format(
- **conn_config)
-
- def get_conn(self):
- """
- Returns a snowflake.connection object
- """
- conn_config = self._get_conn_params()
- conn = snowflake.connector.connect(**conn_config)
- return conn
-
- def _get_aws_credentials(self):
- """
- returns aws_access_key_id, aws_secret_access_key
- from extra
-
- intended to be used by external import and export statements
- """
- if self.snowflake_conn_id:
- connection_object = self.get_connection(self.snowflake_conn_id)
- if 'aws_secret_access_key' in connection_object.extra_dejson:
- aws_access_key_id = connection_object.extra_dejson.get(
- 'aws_access_key_id')
- aws_secret_access_key = connection_object.extra_dejson.get(
- 'aws_secret_access_key')
- return aws_access_key_id, aws_secret_access_key
-
- def set_autocommit(self, conn, autocommit):
- conn.autocommit(autocommit)
diff --git a/airflow/contrib/hooks/spark_jdbc_hook.py b/airflow/contrib/hooks/spark_jdbc_hook.py
deleted file mode 100644
index b55e4ef060..0000000000
--- a/airflow/contrib/hooks/spark_jdbc_hook.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 os
-from airflow.contrib.hooks.spark_submit_hook import SparkSubmitHook
-from airflow.exceptions import AirflowException
-
-
-class SparkJDBCHook(SparkSubmitHook):
- """
- This hook extends the SparkSubmitHook specifically for performing data
- transfers to/from JDBC-based databases with Apache Spark.
-
- :param spark_app_name: Name of the job (default airflow-spark-jdbc)
- :type spark_app_name: str
- :param spark_conn_id: Connection id as configured in Airflow administration
- :type spark_conn_id: str
- :param spark_conf: Any additional Spark configuration properties
- :type spark_conf: dict
- :param spark_py_files: Additional python files used (.zip, .egg, or .py)
- :type spark_py_files: str
- :param spark_files: Additional files to upload to the container running the job
- :type spark_files: str
- :param spark_jars: Additional jars to upload and add to the driver and
- executor classpath
- :type spark_jars: str
- :param num_executors: number of executor to run. This should be set so as to manage
- the number of connections made with the JDBC database
- :type num_executors: int
- :param executor_cores: Number of cores per executor
- :type executor_cores: int
- :param executor_memory: Memory per executor (e.g. 1000M, 2G)
- :type executor_memory: str
- :param driver_memory: Memory allocated to the driver (e.g. 1000M, 2G)
- :type driver_memory: str
- :param verbose: Whether to pass the verbose flag to spark-submit for debugging
- :type verbose: bool
- :param keytab: Full path to the file that contains the keytab
- :type keytab: str
- :param principal: The name of the kerberos principal used for keytab
- :type principal: str
- :param cmd_type: Which way the data should flow. 2 possible values:
- spark_to_jdbc: data written by spark from metastore to jdbc
- jdbc_to_spark: data written by spark from jdbc to metastore
- :type cmd_type: str
- :param jdbc_table: The name of the JDBC table
- :type jdbc_table: str
- :param jdbc_conn_id: Connection id used for connection to JDBC database
- :type: jdbc_conn_id: str
- :param jdbc_driver: Name of the JDBC driver to use for the JDBC connection. This
- driver (usually a jar) should be passed in the 'jars' parameter
- :type jdbc_driver: str
- :param metastore_table: The name of the metastore table,
- :type metastore_table: str
- :param jdbc_truncate: (spark_to_jdbc only) Whether or not Spark should truncate or
- drop and recreate the JDBC table. This only takes effect if
- 'save_mode' is set to Overwrite. Also, if the schema is
- different, Spark cannot truncate, and will drop and recreate
- :type jdbc_truncate: bool
- :param save_mode: The Spark save-mode to use (e.g. overwrite, append, etc.)
- :type save_mode: str
- :param save_format: (jdbc_to_spark-only) The Spark save-format to use (e.g. parquet)
- :type save_format: str
- :param batch_size: (spark_to_jdbc only) The size of the batch to insert per round
- trip to the JDBC database. Defaults to 1000
- :type batch_size: int
- :param fetch_size: (jdbc_to_spark only) The size of the batch to fetch per round trip
- from the JDBC database. Default depends on the JDBC driver
- :type fetch_size: int
- :param num_partitions: The maximum number of partitions that can be used by Spark
- simultaneously, both for spark_to_jdbc and jdbc_to_spark
- operations. This will also cap the number of JDBC connections
- that can be opened
- :type num_partitions: int
- :param partition_column: (jdbc_to_spark-only) A numeric column to be used to
- partition the metastore table by. If specified, you must
- also specify:
- num_partitions, lower_bound, upper_bound
- :type partition_column: str
- :param lower_bound: (jdbc_to_spark-only) Lower bound of the range of the numeric
- partition column to fetch. If specified, you must also specify:
- num_partitions, partition_column, upper_bound
- :type lower_bound: int
- :param upper_bound: (jdbc_to_spark-only) Upper bound of the range of the numeric
- partition column to fetch. If specified, you must also specify:
- num_partitions, partition_column, lower_bound
- :type upper_bound: int
- :param create_table_column_types: (spark_to_jdbc-only) The database column data types
- to use instead of the defaults, when creating the
- table. Data type information should be specified in
- the same format as CREATE TABLE columns syntax
- (e.g: "name CHAR(64), comments VARCHAR(1024)").
- The specified types should be valid spark sql data
- types.
- """
- def __init__(self,
- spark_app_name='airflow-spark-jdbc',
- spark_conn_id='spark-default',
- spark_conf=None,
- spark_py_files=None,
- spark_files=None,
- spark_jars=None,
- num_executors=None,
- executor_cores=None,
- executor_memory=None,
- driver_memory=None,
- verbose=False,
- principal=None,
- keytab=None,
- cmd_type='spark_to_jdbc',
- jdbc_table=None,
- jdbc_conn_id='jdbc-default',
- jdbc_driver=None,
- metastore_table=None,
- jdbc_truncate=False,
- save_mode=None,
- save_format=None,
- batch_size=None,
- fetch_size=None,
- num_partitions=None,
- partition_column=None,
- lower_bound=None,
- upper_bound=None,
- create_table_column_types=None,
- *args,
- **kwargs
- ):
- super(SparkJDBCHook, self).__init__(*args, **kwargs)
- self._name = spark_app_name
- self._conn_id = spark_conn_id
- self._conf = spark_conf
- self._py_files = spark_py_files
- self._files = spark_files
- self._jars = spark_jars
- self._num_executors = num_executors
- self._executor_cores = executor_cores
- self._executor_memory = executor_memory
- self._driver_memory = driver_memory
- self._verbose = verbose
- self._keytab = keytab
- self._principal = principal
- self._cmd_type = cmd_type
- self._jdbc_table = jdbc_table
- self._jdbc_conn_id = jdbc_conn_id
- self._jdbc_driver = jdbc_driver
- self._metastore_table = metastore_table
- self._jdbc_truncate = jdbc_truncate
- self._save_mode = save_mode
- self._save_format = save_format
- self._batch_size = batch_size
- self._fetch_size = fetch_size
- self._num_partitions = num_partitions
- self._partition_column = partition_column
- self._lower_bound = lower_bound
- self._upper_bound = upper_bound
- self._create_table_column_types = create_table_column_types
- self._jdbc_connection = self._resolve_jdbc_connection()
-
- def _resolve_jdbc_connection(self):
- conn_data = {'url': '',
- 'schema': '',
- 'conn_prefix': '',
- 'user': '',
- 'password': ''
- }
- try:
- conn = self.get_connection(self._jdbc_conn_id)
- if conn.port:
- conn_data['url'] = "{}:{}".format(conn.host, conn.port)
- else:
- conn_data['url'] = conn.host
- conn_data['schema'] = conn.schema
- conn_data['user'] = conn.login
- conn_data['password'] = conn.password
- extra = conn.extra_dejson
- conn_data['conn_prefix'] = extra.get('conn_prefix', '')
- except AirflowException:
- self.log.debug(
- "Could not load jdbc connection string %s, defaulting to %s",
- self._jdbc_conn_id, ""
- )
- return conn_data
-
- def _build_jdbc_application_arguments(self, jdbc_conn):
- arguments = []
- arguments += ["-cmdType", self._cmd_type]
- if self._jdbc_connection['url']:
- arguments += ['-url', "{0}{1}/{2}".format(
- jdbc_conn['conn_prefix'], jdbc_conn['url'], jdbc_conn['schema']
- )]
- if self._jdbc_connection['user']:
- arguments += ['-user', self._jdbc_connection['user']]
- if self._jdbc_connection['password']:
- arguments += ['-password', self._jdbc_connection['password']]
- if self._metastore_table:
- arguments += ['-metastoreTable', self._metastore_table]
- if self._jdbc_table:
- arguments += ['-jdbcTable', self._jdbc_table]
- if self._jdbc_truncate:
- arguments += ['-jdbcTruncate', str(self._jdbc_truncate)]
- if self._jdbc_driver:
- arguments += ['-jdbcDriver', self._jdbc_driver]
- if self._batch_size:
- arguments += ['-batchsize', str(self._batch_size)]
- if self._fetch_size:
- arguments += ['-fetchsize', str(self._fetch_size)]
- if self._num_partitions:
- arguments += ['-numPartitions', str(self._num_partitions)]
- if (self._partition_column and self._lower_bound and
- self._upper_bound and self._num_partitions):
- # these 3 parameters need to be used all together to take effect.
- arguments += ['-partitionColumn', self._partition_column,
- '-lowerBound', self._lower_bound,
- '-upperBound', self._upper_bound]
- if self._save_mode:
- arguments += ['-saveMode', self._save_mode]
- if self._save_format:
- arguments += ['-saveFormat', self._save_format]
- if self._create_table_column_types:
- arguments += ['-createTableColumnTypes', self._create_table_column_types]
- return arguments
-
- def submit_jdbc_job(self):
- self._application_args = \
- self._build_jdbc_application_arguments(self._jdbc_connection)
- self.submit(application=os.path.dirname(os.path.abspath(__file__)) +
- "/spark_jdbc_script.py")
-
- def get_conn(self):
- pass
diff --git a/airflow/contrib/hooks/spark_jdbc_script.py b/airflow/contrib/hooks/spark_jdbc_script.py
deleted file mode 100644
index 8e27f363f8..0000000000
--- a/airflow/contrib/hooks/spark_jdbc_script.py
+++ /dev/null
@@ -1,146 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 argparse
-from pyspark.sql import SparkSession
-
-
-def set_common_options(spark_source,
- url='localhost:5432',
- jdbc_table='default.default',
- user='root',
- password='root',
- driver='driver'):
-
- spark_source = spark_source \
- .format('jdbc') \
- .option('url', url) \
- .option('dbtable', jdbc_table) \
- .option('user', user) \
- .option('password', password) \
- .option('driver', driver)
- return spark_source
-
-
-def spark_write_to_jdbc(spark, url, user, password, metastore_table, jdbc_table, driver,
- truncate, save_mode, batch_size, num_partitions,
- create_table_column_types):
- writer = spark \
- .table(metastore_table) \
- .write \
-
- # first set common options
- writer = set_common_options(writer, url, jdbc_table, user, password, driver)
-
- # now set write-specific options
- if truncate:
- writer = writer.option('truncate', truncate)
- if batch_size:
- writer = writer.option('batchsize', batch_size)
- if num_partitions:
- writer = writer.option('numPartitions', num_partitions)
- if create_table_column_types:
- writer = writer.option("createTableColumnTypes", create_table_column_types)
-
- writer \
- .save(mode=save_mode)
-
-
-def spark_read_from_jdbc(spark, url, user, password, metastore_table, jdbc_table, driver,
- save_mode, save_format, fetch_size, num_partitions,
- partition_column, lower_bound, upper_bound):
-
- # first set common options
- reader = set_common_options(spark.read, url, jdbc_table, user, password, driver)
-
- # now set specific read options
- if fetch_size:
- reader = reader.option('fetchsize', fetch_size)
- if num_partitions:
- reader = reader.option('numPartitions', num_partitions)
- if partition_column and lower_bound and upper_bound:
- reader = reader \
- .option('partitionColumn', partition_column) \
- .option('lowerBound', lower_bound) \
- .option('upperBound', upper_bound)
-
- reader \
- .load() \
- .write \
- .saveAsTable(metastore_table, format=save_format, mode=save_mode)
-
-
-if __name__ == "__main__": # pragma: no cover
- # parse the parameters
- parser = argparse.ArgumentParser(description='Spark-JDBC')
- parser.add_argument('-cmdType', dest='cmd_type', action='store')
- parser.add_argument('-url', dest='url', action='store')
- parser.add_argument('-user', dest='user', action='store')
- parser.add_argument('-password', dest='password', action='store')
- parser.add_argument('-metastoreTable', dest='metastore_table', action='store')
- parser.add_argument('-jdbcTable', dest='jdbc_table', action='store')
- parser.add_argument('-jdbcDriver', dest='jdbc_driver', action='store')
- parser.add_argument('-jdbcTruncate', dest='truncate', action='store')
- parser.add_argument('-saveMode', dest='save_mode', action='store')
- parser.add_argument('-saveFormat', dest='save_format', action='store')
- parser.add_argument('-batchsize', dest='batch_size', action='store')
- parser.add_argument('-fetchsize', dest='fetch_size', action='store')
- parser.add_argument('-name', dest='name', action='store')
- parser.add_argument('-numPartitions', dest='num_partitions', action='store')
- parser.add_argument('-partitionColumn', dest='partition_column', action='store')
- parser.add_argument('-lowerBound', dest='lower_bound', action='store')
- parser.add_argument('-upperBound', dest='upper_bound', action='store')
- parser.add_argument('-createTableColumnTypes',
- dest='create_table_column_types', action='store')
- arguments = parser.parse_args()
-
- # Disable dynamic allocation by default to allow num_executors to take effect.
- spark = SparkSession.builder \
- .appName(arguments.name) \
- .enableHiveSupport() \
- .getOrCreate()
-
- if arguments.cmd_type == "spark_to_jdbc":
- spark_write_to_jdbc(spark,
- arguments.url,
- arguments.user,
- arguments.password,
- arguments.metastore_table,
- arguments.jdbc_table,
- arguments.jdbc_driver,
- arguments.truncate,
- arguments.save_mode,
- arguments.batch_size,
- arguments.num_partitions,
- arguments.create_table_column_types)
- elif arguments.cmd_type == "jdbc_to_spark":
- spark_read_from_jdbc(spark,
- arguments.url,
- arguments.user,
- arguments.password,
- arguments.metastore_table,
- arguments.jdbc_table,
- arguments.jdbc_driver,
- arguments.save_mode,
- arguments.save_format,
- arguments.fetch_size,
- arguments.num_partitions,
- arguments.partition_column,
- arguments.lower_bound,
- arguments.upper_bound)
diff --git a/airflow/contrib/hooks/spark_sql_hook.py b/airflow/contrib/hooks/spark_sql_hook.py
deleted file mode 100644
index c1fd2ce21c..0000000000
--- a/airflow/contrib/hooks/spark_sql_hook.py
+++ /dev/null
@@ -1,162 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 subprocess
-
-from airflow.hooks.base_hook import BaseHook
-from airflow.exceptions import AirflowException
-
-
-class SparkSqlHook(BaseHook):
- """
- This hook is a wrapper around the spark-sql binary. It requires that the
- "spark-sql" binary is in the PATH.
- :param sql: The SQL query to execute
- :type sql: str
- :param conf: arbitrary Spark configuration property
- :type conf: str (format: PROP=VALUE)
- :param conn_id: connection_id string
- :type conn_id: str
- :param total_executor_cores: (Standalone & Mesos only) Total cores for all executors
- (Default: all the available cores on the worker)
- :type total_executor_cores: int
- :param executor_cores: (Standalone & YARN only) Number of cores per
- executor (Default: 2)
- :type executor_cores: int
- :param executor_memory: Memory per executor (e.g. 1000M, 2G) (Default: 1G)
- :type executor_memory: str
- :param keytab: Full path to the file that contains the keytab
- :type keytab: str
- :param master: spark://host:port, mesos://host:port, yarn, or local
- :type master: str
- :param name: Name of the job.
- :type name: str
- :param num_executors: Number of executors to launch
- :type num_executors: int
- :param verbose: Whether to pass the verbose flag to spark-sql
- :type verbose: bool
- :param yarn_queue: The YARN queue to submit to (Default: "default")
- :type yarn_queue: str
- """
- def __init__(self,
- sql,
- conf=None,
- conn_id='spark_sql_default',
- total_executor_cores=None,
- executor_cores=None,
- executor_memory=None,
- keytab=None,
- principal=None,
- master='yarn',
- name='default-name',
- num_executors=None,
- verbose=True,
- yarn_queue='default'
- ):
- self._sql = sql
- self._conf = conf
- self._conn = self.get_connection(conn_id)
- self._total_executor_cores = total_executor_cores
- self._executor_cores = executor_cores
- self._executor_memory = executor_memory
- self._keytab = keytab
- self._principal = principal
- self._master = master
- self._name = name
- self._num_executors = num_executors
- self._verbose = verbose
- self._yarn_queue = yarn_queue
- self._sp = None
-
- def get_conn(self):
- pass
-
- def _prepare_command(self, cmd):
- """
- Construct the spark-sql command to execute. Verbose output is enabled
- as default.
- :param cmd: command to append to the spark-sql command
- :type cmd: str
- :return: full command to be executed
- """
- connection_cmd = ["spark-sql"]
- if self._conf:
- for conf_el in self._conf.split(","):
- connection_cmd += ["--conf", conf_el]
- if self._total_executor_cores:
- connection_cmd += ["--total-executor-cores", str(self._total_executor_cores)]
- if self._executor_cores:
- connection_cmd += ["--executor-cores", str(self._executor_cores)]
- if self._executor_memory:
- connection_cmd += ["--executor-memory", self._executor_memory]
- if self._keytab:
- connection_cmd += ["--keytab", self._keytab]
- if self._principal:
- connection_cmd += ["--principal", self._principal]
- if self._num_executors:
- connection_cmd += ["--num-executors", str(self._num_executors)]
- if self._sql:
- sql = self._sql.strip()
- if sql.endswith(".sql") or sql.endswith(".hql"):
- connection_cmd += ["-f", sql]
- else:
- connection_cmd += ["-e", sql]
- if self._master:
- connection_cmd += ["--master", self._master]
- if self._name:
- connection_cmd += ["--name", self._name]
- if self._verbose:
- connection_cmd += ["--verbose"]
- if self._yarn_queue:
- connection_cmd += ["--queue", self._yarn_queue]
-
- connection_cmd += cmd
- self.log.debug("Spark-Sql cmd: %s", connection_cmd)
-
- return connection_cmd
-
- def run_query(self, cmd="", **kwargs):
- """
- Remote Popen (actually execute the Spark-sql query)
-
- :param cmd: command to remotely execute
- :param kwargs: extra arguments to Popen (see subprocess.Popen)
- """
- spark_sql_cmd = self._prepare_command(cmd)
- self._sp = subprocess.Popen(spark_sql_cmd,
- stdout=subprocess.PIPE,
- stderr=subprocess.STDOUT,
- **kwargs)
-
- for line in iter(self._sp.stdout.readline, ''):
- self.log.info(line)
-
- returncode = self._sp.wait()
-
- if returncode:
- raise AirflowException(
- "Cannot execute {} on {}. Process exit code: {}.".format(
- cmd, self._conn.host, returncode
- )
- )
-
- def kill(self):
- if self._sp and self._sp.poll() is None:
- self.log.info("Killing the Spark-Sql job")
- self._sp.kill()
diff --git a/airflow/contrib/hooks/spark_submit_hook.py b/airflow/contrib/hooks/spark_submit_hook.py
deleted file mode 100644
index 0185cab283..0000000000
--- a/airflow/contrib/hooks/spark_submit_hook.py
+++ /dev/null
@@ -1,564 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 os
-import subprocess
-import re
-import time
-
-from airflow.hooks.base_hook import BaseHook
-from airflow.exceptions import AirflowException
-from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.contrib.kubernetes import kube_client
-from kubernetes.client.rest import ApiException
-
-
-class SparkSubmitHook(BaseHook, LoggingMixin):
- """
- This hook is a wrapper around the spark-submit binary to kick off a spark-submit job.
- It requires that the "spark-submit" binary is in the PATH or the spark_home to be
- supplied.
- :param conf: Arbitrary Spark configuration properties
- :type conf: dict
- :param conn_id: The connection id as configured in Airflow administration. When an
- invalid connection_id is supplied, it will default to yarn.
- :type conn_id: str
- :param files: Upload additional files to the executor running the job, separated by a
- comma. Files will be placed in the working directory of each executor.
- For example, serialized objects.
- :type files: str
- :param py_files: Additional python files used by the job, can be .zip, .egg or .py.
- :type py_files: str
- :param driver_classpath: Additional, driver-specific, classpath settings.
- :type driver_classpath: str
- :param jars: Submit additional jars to upload and place them in executor classpath.
- :type jars: str
- :param java_class: the main class of the Java application
- :type java_class: str
- :param packages: Comma-separated list of maven coordinates of jars to include on the
- driver and executor classpaths
- :type packages: str
- :param exclude_packages: Comma-separated list of maven coordinates of jars to exclude
- while resolving the dependencies provided in 'packages'
- :type exclude_packages: str
- :param repositories: Comma-separated list of additional remote repositories to search
- for the maven coordinates given with 'packages'
- :type repositories: str
- :param total_executor_cores: (Standalone & Mesos only) Total cores for all executors
- (Default: all the available cores on the worker)
- :type total_executor_cores: int
- :param executor_cores: (Standalone, YARN and Kubernetes only) Number of cores per
- executor (Default: 2)
- :type executor_cores: int
- :param executor_memory: Memory per executor (e.g. 1000M, 2G) (Default: 1G)
- :type executor_memory: str
- :param driver_memory: Memory allocated to the driver (e.g. 1000M, 2G) (Default: 1G)
- :type driver_memory: str
- :param keytab: Full path to the file that contains the keytab
- :type keytab: str
- :param principal: The name of the kerberos principal used for keytab
- :type principal: str
- :param name: Name of the job (default airflow-spark)
- :type name: str
- :param num_executors: Number of executors to launch
- :type num_executors: int
- :param application_args: Arguments for the application being submitted
- :type application_args: list
- :param env_vars: Environment variables for spark-submit. It
- supports yarn and k8s mode too.
- :type env_vars: dict
- :param verbose: Whether to pass the verbose flag to spark-submit process for debugging
- :type verbose: bool
- """
- def __init__(self,
- conf=None,
- conn_id='spark_default',
- files=None,
- py_files=None,
- driver_classpath=None,
- jars=None,
- java_class=None,
- packages=None,
- exclude_packages=None,
- repositories=None,
- total_executor_cores=None,
- executor_cores=None,
- executor_memory=None,
- driver_memory=None,
- keytab=None,
- principal=None,
- name='default-name',
- num_executors=None,
- application_args=None,
- env_vars=None,
- verbose=False):
- self._conf = conf
- self._conn_id = conn_id
- self._files = files
- self._py_files = py_files
- self._driver_classpath = driver_classpath
- self._jars = jars
- self._java_class = java_class
- self._packages = packages
- self._exclude_packages = exclude_packages
- self._repositories = repositories
- self._total_executor_cores = total_executor_cores
- self._executor_cores = executor_cores
- self._executor_memory = executor_memory
- self._driver_memory = driver_memory
- self._keytab = keytab
- self._principal = principal
- self._name = name
- self._num_executors = num_executors
- self._application_args = application_args
- self._env_vars = env_vars
- self._verbose = verbose
- self._submit_sp = None
- self._yarn_application_id = None
- self._kubernetes_driver_pod = None
-
- self._connection = self._resolve_connection()
- self._is_yarn = 'yarn' in self._connection['master']
- self._is_kubernetes = 'k8s' in self._connection['master']
-
- self._should_track_driver_status = self._resolve_should_track_driver_status()
- self._driver_id = None
- self._driver_status = None
- self._spark_exit_code = None
-
- def _resolve_should_track_driver_status(self):
- """
- Determines whether or not this hook should poll the spark driver status through
- subsequent spark-submit status requests after the initial spark-submit request
- :return: if the driver status should be tracked
- """
- return ('spark://' in self._connection['master'] and
- self._connection['deploy_mode'] == 'cluster')
-
- def _resolve_connection(self):
- # Build from connection master or default to yarn if not available
- conn_data = {'master': 'yarn',
- 'queue': None,
- 'deploy_mode': None,
- 'spark_home': None,
- 'spark_binary': 'spark-submit',
- 'namespace': 'default'}
-
- try:
- # Master can be local, yarn, spark://HOST:PORT, mesos://HOST:PORT and
- # k8s://https://<HOST>:<PORT>
- conn = self.get_connection(self._conn_id)
- if conn.port:
- conn_data['master'] = "{}:{}".format(conn.host, conn.port)
- else:
- conn_data['master'] = conn.host
-
- # Determine optional yarn queue from the extra field
- extra = conn.extra_dejson
- conn_data['queue'] = extra.get('queue', None)
- conn_data['deploy_mode'] = extra.get('deploy-mode', None)
- conn_data['spark_home'] = extra.get('spark-home', None)
- conn_data['spark_binary'] = extra.get('spark-binary', 'spark-submit')
- conn_data['namespace'] = extra.get('namespace', 'default')
- except AirflowException:
- self.log.debug(
- "Could not load connection string %s, defaulting to %s",
- self._conn_id, conn_data['master']
- )
-
- return conn_data
-
- def get_conn(self):
- pass
-
- def _get_spark_binary_path(self):
- # If the spark_home is passed then build the spark-submit executable path using
- # the spark_home; otherwise assume that spark-submit is present in the path to
- # the executing user
- if self._connection['spark_home']:
- connection_cmd = [os.path.join(self._connection['spark_home'], 'bin',
- self._connection['spark_binary'])]
- else:
- connection_cmd = [self._connection['spark_binary']]
-
- return connection_cmd
-
- def _build_spark_submit_command(self, application):
- """
- Construct the spark-submit command to execute.
- :param application: command to append to the spark-submit command
- :type application: str
- :return: full command to be executed
- """
- connection_cmd = self._get_spark_binary_path()
-
- # The url ot the spark master
- connection_cmd += ["--master", self._connection['master']]
-
- if self._conf:
- for key in self._conf:
- connection_cmd += ["--conf", "{}={}".format(key, str(self._conf[key]))]
- if self._env_vars and (self._is_kubernetes or self._is_yarn):
- if self._is_yarn:
- tmpl = "spark.yarn.appMasterEnv.{}={}"
- else:
- tmpl = "spark.kubernetes.driverEnv.{}={}"
- for key in self._env_vars:
- connection_cmd += [
- "--conf",
- tmpl.format(key, str(self._env_vars[key]))]
- elif self._env_vars and self._connection['deploy_mode'] != "cluster":
- self._env = self._env_vars # Do it on Popen of the process
- elif self._env_vars and self._connection['deploy_mode'] == "cluster":
- raise AirflowException(
- "SparkSubmitHook env_vars is not supported in standalone-cluster mode.")
- if self._is_kubernetes:
- connection_cmd += ["--conf", "spark.kubernetes.namespace={}".format(
- self._connection['namespace'])]
- if self._files:
- connection_cmd += ["--files", self._files]
- if self._py_files:
- connection_cmd += ["--py-files", self._py_files]
- if self._driver_classpath:
- connection_cmd += ["--driver-classpath", self._driver_classpath]
- if self._jars:
- connection_cmd += ["--jars", self._jars]
- if self._packages:
- connection_cmd += ["--packages", self._packages]
- if self._exclude_packages:
- connection_cmd += ["--exclude-packages", self._exclude_packages]
- if self._repositories:
- connection_cmd += ["--repositories", self._repositories]
- if self._num_executors:
- connection_cmd += ["--num-executors", str(self._num_executors)]
- if self._total_executor_cores:
- connection_cmd += ["--total-executor-cores", str(self._total_executor_cores)]
- if self._executor_cores:
- connection_cmd += ["--executor-cores", str(self._executor_cores)]
- if self._executor_memory:
- connection_cmd += ["--executor-memory", self._executor_memory]
- if self._driver_memory:
- connection_cmd += ["--driver-memory", self._driver_memory]
- if self._keytab:
- connection_cmd += ["--keytab", self._keytab]
- if self._principal:
- connection_cmd += ["--principal", self._principal]
- if self._name:
- connection_cmd += ["--name", self._name]
- if self._java_class:
- connection_cmd += ["--class", self._java_class]
- if self._verbose:
- connection_cmd += ["--verbose"]
- if self._connection['queue']:
- connection_cmd += ["--queue", self._connection['queue']]
- if self._connection['deploy_mode']:
- connection_cmd += ["--deploy-mode", self._connection['deploy_mode']]
-
- # The actual script to execute
- connection_cmd += [application]
-
- # Append any application arguments
- if self._application_args:
- connection_cmd += self._application_args
-
- self.log.info("Spark-Submit cmd: %s", connection_cmd)
-
- return connection_cmd
-
- def _build_track_driver_status_command(self):
- """
- Construct the command to poll the driver status.
-
- :return: full command to be executed
- """
- connection_cmd = self._get_spark_binary_path()
-
- # The url ot the spark master
- connection_cmd += ["--master", self._connection['master']]
-
- # The driver id so we can poll for its status
- if self._driver_id:
- connection_cmd += ["--status", self._driver_id]
- else:
- raise AirflowException(
- "Invalid status: attempted to poll driver " +
- "status but no driver id is known. Giving up.")
-
- self.log.debug("Poll driver status cmd: %s", connection_cmd)
-
- return connection_cmd
-
- def submit(self, application="", **kwargs):
- """
- Remote Popen to execute the spark-submit job
-
- :param application: Submitted application, jar or py file
- :type application: str
- :param kwargs: extra arguments to Popen (see subprocess.Popen)
- """
- spark_submit_cmd = self._build_spark_submit_command(application)
-
- if hasattr(self, '_env'):
- env = os.environ.copy()
- env.update(self._env)
- kwargs["env"] = env
-
- self._submit_sp = subprocess.Popen(spark_submit_cmd,
- stdout=subprocess.PIPE,
- stderr=subprocess.STDOUT,
- bufsize=-1,
- universal_newlines=True,
- **kwargs)
-
- self._process_spark_submit_log(iter(self._submit_sp.stdout.readline, ''))
- returncode = self._submit_sp.wait()
-
- # Check spark-submit return code. In Kubernetes mode, also check the value
- # of exit code in the log, as it may differ.
- if returncode or (self._is_kubernetes and self._spark_exit_code != 0):
- raise AirflowException(
- "Cannot execute: {}. Error code is: {}.".format(
- spark_submit_cmd, returncode
- )
- )
-
- self.log.debug("Should track driver: {}".format(self._should_track_driver_status))
-
- # We want the Airflow job to wait until the Spark driver is finished
- if self._should_track_driver_status:
- if self._driver_id is None:
- raise AirflowException(
- "No driver id is known: something went wrong when executing " +
- "the spark submit command"
- )
-
- # We start with the SUBMITTED status as initial status
- self._driver_status = "SUBMITTED"
-
- # Start tracking the driver status (blocking function)
- self._start_driver_status_tracking()
-
- if self._driver_status != "FINISHED":
- raise AirflowException(
- "ERROR : Driver {} badly exited with status {}"
- .format(self._driver_id, self._driver_status)
- )
-
- def _process_spark_submit_log(self, itr):
- """
- Processes the log files and extracts useful information out of it.
-
- If the deploy-mode is 'client', log the output of the submit command as those
- are the output logs of the Spark worker directly.
-
- Remark: If the driver needs to be tracked for its status, the log-level of the
- spark deploy needs to be at least INFO (log4j.logger.org.apache.spark.deploy=INFO)
-
- :param itr: An iterator which iterates over the input of the subprocess
- """
- # Consume the iterator
- for line in itr:
- line = line.strip()
- # If we run yarn cluster mode, we want to extract the application id from
- # the logs so we can kill the application when we stop it unexpectedly
- if self._is_yarn and self._connection['deploy_mode'] == 'cluster':
- match = re.search('(application[0-9_]+)', line)
- if match:
- self._yarn_application_id = match.groups()[0]
- self.log.info("Identified spark driver id: %s",
- self._yarn_application_id)
-
- # If we run Kubernetes cluster mode, we want to extract the driver pod id
- # from the logs so we can kill the application when we stop it unexpectedly
- elif self._is_kubernetes:
- match = re.search('\s*pod name: ((.+?)-([a-z0-9]+)-driver)', line)
- if match:
- self._kubernetes_driver_pod = match.groups()[0]
- self.log.info("Identified spark driver pod: %s",
- self._kubernetes_driver_pod)
-
- # Store the Spark Exit code
- match_exit_code = re.search('\s*Exit code: (\d+)', line)
- if match_exit_code:
- self._spark_exit_code = int(match_exit_code.groups()[0])
-
- # if we run in standalone cluster mode and we want to track the driver status
- # we need to extract the driver id from the logs. This allows us to poll for
- # the status using the driver id. Also, we can kill the driver when needed.
- elif self._should_track_driver_status and not self._driver_id:
- match_driver_id = re.search('(driver-[0-9\-]+)', line)
- if match_driver_id:
- self._driver_id = match_driver_id.groups()[0]
- self.log.info("identified spark driver id: {}"
- .format(self._driver_id))
-
- else:
- self.log.info(line)
-
- self.log.debug("spark submit log: {}".format(line))
-
- def _process_spark_status_log(self, itr):
- """
- parses the logs of the spark driver status query process
-
- :param itr: An iterator which iterates over the input of the subprocess
- """
- # Consume the iterator
- for line in itr:
- line = line.strip()
-
- # Check if the log line is about the driver status and extract the status.
- if "driverState" in line:
- self._driver_status = line.split(' : ')[1] \
- .replace(',', '').replace('\"', '').strip()
-
- self.log.debug("spark driver status log: {}".format(line))
-
- def _start_driver_status_tracking(self):
- """
- Polls the driver based on self._driver_id to get the status.
- Finish successfully when the status is FINISHED.
- Finish failed when the status is ERROR/UNKNOWN/KILLED/FAILED.
-
- Possible status:
- SUBMITTED: Submitted but not yet scheduled on a worker
- RUNNING: Has been allocated to a worker to run
- FINISHED: Previously ran and exited cleanly
- RELAUNCHING: Exited non-zero or due to worker failure, but has not yet
- started running again
- UNKNOWN: The status of the driver is temporarily not known due to
- master failure recovery
- KILLED: A user manually killed this driver
- FAILED: The driver exited non-zero and was not supervised
- ERROR: Unable to run or restart due to an unrecoverable error
- (e.g. missing jar file)
- """
-
- # When your Spark Standalone cluster is not performing well
- # due to misconfiguration or heavy loads.
- # it is possible that the polling request will timeout.
- # Therefore we use a simple retry mechanism.
- missed_job_status_reports = 0
- max_missed_job_status_reports = 10
-
- # Keep polling as long as the driver is processing
- while self._driver_status not in ["FINISHED", "UNKNOWN",
- "KILLED", "FAILED", "ERROR"]:
-
- # Sleep for 1 second as we do not want to spam the cluster
- time.sleep(1)
-
- self.log.debug("polling status of spark driver with id {}"
- .format(self._driver_id))
-
- poll_drive_status_cmd = self._build_track_driver_status_command()
- status_process = subprocess.Popen(poll_drive_status_cmd,
- stdout=subprocess.PIPE,
- stderr=subprocess.STDOUT,
- bufsize=-1,
- universal_newlines=True)
-
- self._process_spark_status_log(iter(status_process.stdout.readline, ''))
- returncode = status_process.wait()
-
- if returncode:
- if missed_job_status_reports < max_missed_job_status_reports:
- missed_job_status_reports = missed_job_status_reports + 1
- else:
- raise AirflowException(
- "Failed to poll for the driver status {} times: returncode = {}"
- .format(max_missed_job_status_reports, returncode)
- )
-
- def _build_spark_driver_kill_command(self):
- """
- Construct the spark-submit command to kill a driver.
- :return: full command to kill a driver
- """
-
- # If the spark_home is passed then build the spark-submit executable path using
- # the spark_home; otherwise assume that spark-submit is present in the path to
- # the executing user
- if self._connection['spark_home']:
- connection_cmd = [os.path.join(self._connection['spark_home'],
- 'bin',
- self._connection['spark_binary'])]
- else:
- connection_cmd = [self._connection['spark_binary']]
-
- # The url ot the spark master
- connection_cmd += ["--master", self._connection['master']]
-
- # The actual kill command
- connection_cmd += ["--kill", self._driver_id]
-
- self.log.debug("Spark-Kill cmd: %s", connection_cmd)
-
- return connection_cmd
-
- def on_kill(self):
-
- self.log.debug("Kill Command is being called")
-
- if self._should_track_driver_status:
- if self._driver_id:
- self.log.info('Killing driver {} on cluster'
- .format(self._driver_id))
-
- kill_cmd = self._build_spark_driver_kill_command()
- driver_kill = subprocess.Popen(kill_cmd,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE)
-
- self.log.info("Spark driver {} killed with return code: {}"
- .format(self._driver_id, driver_kill.wait()))
-
- if self._submit_sp and self._submit_sp.poll() is None:
- self.log.info('Sending kill signal to %s', self._connection['spark_binary'])
- self._submit_sp.kill()
-
- if self._yarn_application_id:
- self.log.info('Killing application {} on YARN'
- .format(self._yarn_application_id))
-
- kill_cmd = "yarn application -kill {}" \
- .format(self._yarn_application_id).split()
- yarn_kill = subprocess.Popen(kill_cmd,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE)
-
- self.log.info("YARN killed with return code: %s", yarn_kill.wait())
-
- if self._kubernetes_driver_pod:
- self.log.info('Killing pod %s on Kubernetes', self._kubernetes_driver_pod)
-
- # Currently only instantiate Kubernetes client for killing a spark pod.
- try:
- client = kube_client.get_kube_client()
- api_response = client.delete_namespaced_pod(
- self._kubernetes_driver_pod,
- self._connection['namespace'],
- body=client.V1DeleteOptions(),
- pretty=True)
-
- self.log.info("Spark on K8s killed with response: %s", api_response)
-
- except ApiException as e:
- self.log.info("Exception when attempting to kill Spark on K8s:")
- self.log.exception(e)
diff --git a/airflow/contrib/hooks/sqoop_hook.py b/airflow/contrib/hooks/sqoop_hook.py
deleted file mode 100644
index 74cddc2b21..0000000000
--- a/airflow/contrib/hooks/sqoop_hook.py
+++ /dev/null
@@ -1,352 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-#
-
-"""
-This module contains a sqoop 1.x hook
-"""
-import subprocess
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-from copy import deepcopy
-
-
-class SqoopHook(BaseHook, LoggingMixin):
- """
- This hook is a wrapper around the sqoop 1 binary. To be able to use the hook
- it is required that "sqoop" is in the PATH.
-
- Additional arguments that can be passed via the 'extra' JSON field of the
- sqoop connection:
- * job_tracker: Job tracker local|jobtracker:port.
- * namenode: Namenode.
- * lib_jars: Comma separated jar files to include in the classpath.
- * files: Comma separated files to be copied to the map reduce cluster.
- * archives: Comma separated archives to be unarchived on the compute
- machines.
- * password_file: Path to file containing the password.
-
- :param conn_id: Reference to the sqoop connection.
- :type conn_id: str
- :param verbose: Set sqoop to verbose.
- :type verbose: bool
- :param num_mappers: Number of map tasks to import in parallel.
- :type num_mappers: int
- :param properties: Properties to set via the -D argument
- :type properties: dict
- """
-
- def __init__(self, conn_id='sqoop_default', verbose=False,
- num_mappers=None, hcatalog_database=None,
- hcatalog_table=None, properties=None):
- # No mutable types in the default parameters
- self.conn = self.get_connection(conn_id)
- connection_parameters = self.conn.extra_dejson
- self.job_tracker = connection_parameters.get('job_tracker', None)
- self.namenode = connection_parameters.get('namenode', None)
- self.libjars = connection_parameters.get('libjars', None)
- self.files = connection_parameters.get('files', None)
- self.archives = connection_parameters.get('archives', None)
- self.password_file = connection_parameters.get('password_file', None)
- self.hcatalog_database = hcatalog_database
- self.hcatalog_table = hcatalog_table
- self.verbose = verbose
- self.num_mappers = num_mappers
- self.properties = properties or {}
- self.log.info(
- "Using connection to: {}:{}/{}".format(
- self.conn.host, self.conn.port, self.conn.schema
- )
- )
-
- def get_conn(self):
- return self.conn
-
- def cmd_mask_password(self, cmd_orig):
- cmd = deepcopy(cmd_orig)
- try:
- password_index = cmd.index('--password')
- cmd[password_index + 1] = 'MASKED'
- except ValueError:
- self.log.debug("No password in sqoop cmd")
- return cmd
-
- def Popen(self, cmd, **kwargs):
- """
- Remote Popen
-
- :param cmd: command to remotely execute
- :param kwargs: extra arguments to Popen (see subprocess.Popen)
- :return: handle to subprocess
- """
- masked_cmd = ' '.join(self.cmd_mask_password(cmd))
- self.log.info("Executing command: {}".format(masked_cmd))
- self.sp = subprocess.Popen(
- cmd,
- stdout=subprocess.PIPE,
- stderr=subprocess.STDOUT,
- **kwargs)
-
- for line in iter(self.sp.stdout):
- self.log.info(line.strip())
-
- self.sp.wait()
-
- self.log.info("Command exited with return code %s", self.sp.returncode)
-
- if self.sp.returncode:
- raise AirflowException("Sqoop command failed: {}".format(masked_cmd))
-
- def _prepare_command(self, export=False):
- sqoop_cmd_type = "export" if export else "import"
- connection_cmd = ["sqoop", sqoop_cmd_type]
-
- for key, value in self.properties.items():
- connection_cmd += ["-D", "{}={}".format(key, value)]
-
- if self.namenode:
- connection_cmd += ["-fs", self.namenode]
- if self.job_tracker:
- connection_cmd += ["-jt", self.job_tracker]
- if self.libjars:
- connection_cmd += ["-libjars", self.libjars]
- if self.files:
- connection_cmd += ["-files", self.files]
- if self.archives:
- connection_cmd += ["-archives", self.archives]
- if self.conn.login:
- connection_cmd += ["--username", self.conn.login]
- if self.conn.password:
- connection_cmd += ["--password", self.conn.password]
- if self.password_file:
- connection_cmd += ["--password-file", self.password_file]
- if self.verbose:
- connection_cmd += ["--verbose"]
- if self.num_mappers:
- connection_cmd += ["--num-mappers", str(self.num_mappers)]
- if self.hcatalog_database:
- connection_cmd += ["--hcatalog-database", self.hcatalog_database]
- if self.hcatalog_table:
- connection_cmd += ["--hcatalog-table", self.hcatalog_table]
-
- connection_cmd += ["--connect", "{}:{}/{}".format(
- self.conn.host,
- self.conn.port,
- self.conn.schema
- )]
-
- return connection_cmd
-
- @staticmethod
- def _get_export_format_argument(file_type='text'):
- if file_type == "avro":
- return ["--as-avrodatafile"]
- elif file_type == "sequence":
- return ["--as-sequencefile"]
- elif file_type == "parquet":
- return ["--as-parquetfile"]
- elif file_type == "text":
- return ["--as-textfile"]
- else:
- raise AirflowException("Argument file_type should be 'avro', "
- "'sequence', 'parquet' or 'text'.")
-
- def _import_cmd(self, target_dir, append, file_type, split_by, direct,
- driver, extra_import_options):
-
- cmd = self._prepare_command(export=False)
-
- if target_dir:
- cmd += ["--target-dir", target_dir]
-
- if append:
- cmd += ["--append"]
-
- cmd += self._get_export_format_argument(file_type)
-
- if split_by:
- cmd += ["--split-by", split_by]
-
- if direct:
- cmd += ["--direct"]
-
- if driver:
- cmd += ["--driver", driver]
-
- if extra_import_options:
- for key, value in extra_import_options.items():
- cmd += ['--{}'.format(key)]
- if value:
- cmd += [value]
-
- return cmd
-
- def import_table(self, table, target_dir=None, append=False, file_type="text",
- columns=None, split_by=None, where=None, direct=False,
- driver=None, extra_import_options=None):
- """
- Imports table from remote location to target dir. Arguments are
- copies of direct sqoop command line arguments
- :param table: Table to read
- :param target_dir: HDFS destination dir
- :param append: Append data to an existing dataset in HDFS
- :param file_type: "avro", "sequence", "text" or "parquet".
- Imports data to into the specified format. Defaults to text.
- :param columns: <col,col,col…> Columns to import from table
- :param split_by: Column of the table used to split work units
- :param where: WHERE clause to use during import
- :param direct: Use direct connector if exists for the database
- :param driver: Manually specify JDBC driver class to use
- :param extra_import_options: Extra import options to pass as dict.
- If a key doesn't have a value, just pass an empty string to it.
- Don't include prefix of -- for sqoop options.
- """
- cmd = self._import_cmd(target_dir, append, file_type, split_by, direct,
- driver, extra_import_options)
-
- cmd += ["--table", table]
-
- if columns:
- cmd += ["--columns", columns]
- if where:
- cmd += ["--where", where]
-
- self.Popen(cmd)
-
- def import_query(self, query, target_dir, append=False, file_type="text",
- split_by=None, direct=None, driver=None, extra_import_options=None):
- """
- Imports a specific query from the rdbms to hdfs
- :param query: Free format query to run
- :param target_dir: HDFS destination dir
- :param append: Append data to an existing dataset in HDFS
- :param file_type: "avro", "sequence", "text" or "parquet"
- Imports data to hdfs into the specified format. Defaults to text.
- :param split_by: Column of the table used to split work units
- :param direct: Use direct import fast path
- :param driver: Manually specify JDBC driver class to use
- :param extra_import_options: Extra import options to pass as dict.
- If a key doesn't have a value, just pass an empty string to it.
- Don't include prefix of -- for sqoop options.
- """
- cmd = self._import_cmd(target_dir, append, file_type, split_by, direct,
- driver, extra_import_options)
- cmd += ["--query", query]
-
- self.Popen(cmd)
-
- def _export_cmd(self, table, export_dir, input_null_string,
- input_null_non_string, staging_table, clear_staging_table,
- enclosed_by, escaped_by, input_fields_terminated_by,
- input_lines_terminated_by, input_optionally_enclosed_by,
- batch, relaxed_isolation, extra_export_options):
-
- cmd = self._prepare_command(export=True)
-
- if input_null_string:
- cmd += ["--input-null-string", input_null_string]
-
- if input_null_non_string:
- cmd += ["--input-null-non-string", input_null_non_string]
-
- if staging_table:
- cmd += ["--staging-table", staging_table]
-
- if clear_staging_table:
- cmd += ["--clear-staging-table"]
-
- if enclosed_by:
- cmd += ["--enclosed-by", enclosed_by]
-
- if escaped_by:
- cmd += ["--escaped-by", escaped_by]
-
- if input_fields_terminated_by:
- cmd += ["--input-fields-terminated-by", input_fields_terminated_by]
-
- if input_lines_terminated_by:
- cmd += ["--input-lines-terminated-by", input_lines_terminated_by]
-
- if input_optionally_enclosed_by:
- cmd += ["--input-optionally-enclosed-by",
- input_optionally_enclosed_by]
-
- if batch:
- cmd += ["--batch"]
-
- if relaxed_isolation:
- cmd += ["--relaxed-isolation"]
-
- if export_dir:
- cmd += ["--export-dir", export_dir]
-
- if extra_export_options:
- for key, value in extra_export_options.items():
- cmd += ['--{}'.format(key)]
- if value:
- cmd += [value]
-
- # The required option
- cmd += ["--table", table]
-
- return cmd
-
- def export_table(self, table, export_dir, input_null_string,
- input_null_non_string, staging_table,
- clear_staging_table, enclosed_by,
- escaped_by, input_fields_terminated_by,
- input_lines_terminated_by,
- input_optionally_enclosed_by, batch,
- relaxed_isolation, extra_export_options=None):
- """
- Exports Hive table to remote location. Arguments are copies of direct
- sqoop command line Arguments
- :param table: Table remote destination
- :param export_dir: Hive table to export
- :param input_null_string: The string to be interpreted as null for
- string columns
- :param input_null_non_string: The string to be interpreted as null
- for non-string columns
- :param staging_table: The table in which data will be staged before
- being inserted into the destination table
- :param clear_staging_table: Indicate that any data present in the
- staging table can be deleted
- :param enclosed_by: Sets a required field enclosing character
- :param escaped_by: Sets the escape character
- :param input_fields_terminated_by: Sets the field separator character
- :param input_lines_terminated_by: Sets the end-of-line character
- :param input_optionally_enclosed_by: Sets a field enclosing character
- :param batch: Use batch mode for underlying statement execution
- :param relaxed_isolation: Transaction isolation to read uncommitted
- for the mappers
- :param extra_export_options: Extra export options to pass as dict.
- If a key doesn't have a value, just pass an empty string to it.
- Don't include prefix of -- for sqoop options.
- """
- cmd = self._export_cmd(table, export_dir, input_null_string,
- input_null_non_string, staging_table,
- clear_staging_table, enclosed_by, escaped_by,
- input_fields_terminated_by,
- input_lines_terminated_by,
- input_optionally_enclosed_by, batch,
- relaxed_isolation, extra_export_options)
-
- self.Popen(cmd)
diff --git a/airflow/contrib/hooks/ssh_hook.py b/airflow/contrib/hooks/ssh_hook.py
deleted file mode 100755
index f51f0fbd11..0000000000
--- a/airflow/contrib/hooks/ssh_hook.py
+++ /dev/null
@@ -1,236 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright 2012-2015 Spotify AB
-# Ported to Airflow by Bolke de Bruin
-#
-# 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 getpass
-import os
-
-import paramiko
-from paramiko.config import SSH_PORT
-
-from contextlib import contextmanager
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class SSHHook(BaseHook, LoggingMixin):
- """
- Hook for ssh remote execution using Paramiko.
- ref: https://github.com/paramiko/paramiko
- This hook also lets you create ssh tunnel and serve as basis for SFTP file transfer
-
- :param ssh_conn_id: connection id from airflow Connections from where all the required
- parameters can be fetched like username, password or key_file.
- Thought the priority is given to the param passed during init
- :type ssh_conn_id: str
- :param remote_host: remote host to connect
- :type remote_host: str
- :param username: username to connect to the remote_host
- :type username: str
- :param password: password of the username to connect to the remote_host
- :type password: str
- :param key_file: key file to use to connect to the remote_host.
- :type key_file: str
- :param port: port of remote host to connect (Default is paramiko SSH_PORT)
- :type port: int
- :param timeout: timeout for the attempt to connect to the remote_host.
- :type timeout: int
- :param keepalive_interval: send a keepalive packet to remote host every
- keepalive_interval seconds
- :type keepalive_interval: int
- """
-
- def __init__(self,
- ssh_conn_id=None,
- remote_host=None,
- username=None,
- password=None,
- key_file=None,
- port=SSH_PORT,
- timeout=10,
- keepalive_interval=30
- ):
- super(SSHHook, self).__init__(ssh_conn_id)
- self.ssh_conn_id = ssh_conn_id
- self.remote_host = remote_host
- self.username = username
- self.password = password
- self.key_file = key_file
- self.timeout = timeout
- self.keepalive_interval = keepalive_interval
- # Default values, overridable from Connection
- self.compress = True
- self.no_host_key_check = True
- self.client = None
- self.port = port
-
- def get_conn(self):
- if not self.client:
- self.log.debug('Creating SSH client for conn_id: %s', self.ssh_conn_id)
- if self.ssh_conn_id is not None:
- conn = self.get_connection(self.ssh_conn_id)
- if self.username is None:
- self.username = conn.login
- if self.password is None:
- self.password = conn.password
- if self.remote_host is None:
- self.remote_host = conn.host
- if conn.port is not None:
- self.port = conn.port
- if conn.extra is not None:
- extra_options = conn.extra_dejson
- self.key_file = extra_options.get("key_file")
-
- if "timeout" in extra_options:
- self.timeout = int(extra_options["timeout"], 10)
-
- if "compress" in extra_options \
- and str(extra_options["compress"]).lower() == 'false':
- self.compress = False
- if "no_host_key_check" in extra_options \
- and \
- str(extra_options["no_host_key_check"]).lower() == 'false':
- self.no_host_key_check = False
-
- if not self.remote_host:
- raise AirflowException("Missing required param: remote_host")
-
- # Auto detecting username values from system
- if not self.username:
- self.log.debug(
- "username to ssh to host: %s is not specified for connection id"
- " %s. Using system's default provided by getpass.getuser()",
- self.remote_host, self.ssh_conn_id
- )
- self.username = getpass.getuser()
-
- host_proxy = None
- user_ssh_config_filename = os.path.expanduser('~/.ssh/config')
- if os.path.isfile(user_ssh_config_filename):
- ssh_conf = paramiko.SSHConfig()
- ssh_conf.parse(open(user_ssh_config_filename))
- host_info = ssh_conf.lookup(self.remote_host)
- if host_info and host_info.get('proxycommand'):
- host_proxy = paramiko.ProxyCommand(host_info.get('proxycommand'))
-
- if not (self.password or self.key_file):
- if host_info and host_info.get('identityfile'):
- self.key_file = host_info.get('identityfile')[0]
-
- try:
- client = paramiko.SSHClient()
- client.load_system_host_keys()
- if self.no_host_key_check:
- # Default is RejectPolicy
- client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-
- if self.password and self.password.strip():
- client.connect(hostname=self.remote_host,
- username=self.username,
- password=self.password,
- timeout=self.timeout,
- compress=self.compress,
- port=self.port,
- sock=host_proxy)
- else:
- client.connect(hostname=self.remote_host,
- username=self.username,
- key_filename=self.key_file,
- timeout=self.timeout,
- compress=self.compress,
- port=self.port,
- sock=host_proxy)
-
- if self.keepalive_interval:
- client.get_transport().set_keepalive(self.keepalive_interval)
-
- self.client = client
- except paramiko.AuthenticationException as auth_error:
- self.log.error(
- "Auth failed while connecting to host: %s, error: %s",
- self.remote_host, auth_error
- )
- except paramiko.SSHException as ssh_error:
- self.log.error(
- "Failed connecting to host: %s, error: %s",
- self.remote_host, ssh_error
- )
- except Exception as error:
- self.log.error(
- "Error connecting to host: %s, error: %s",
- self.remote_host, error
- )
- return self.client
-
- @contextmanager
- def create_tunnel(self, local_port, remote_port=None, remote_host="localhost"):
- """
- Creates a tunnel between two hosts. Like ssh -L <LOCAL_PORT>:host:<REMOTE_PORT>.
- Remember to close() the returned "tunnel" object in order to clean up
- after yourself when you are done with the tunnel.
-
- :param local_port:
- :type local_port: int
- :param remote_port:
- :type remote_port: int
- :param remote_host:
- :type remote_host: str
- :return:
- """
-
- import subprocess
- # this will ensure the connection to the ssh.remote_host from where the tunnel
- # is getting created
- self.get_conn()
-
- tunnel_host = "{0}:{1}:{2}".format(local_port, remote_host, remote_port)
-
- ssh_cmd = ["ssh", "{0}@{1}".format(self.username, self.remote_host),
- "-o", "ControlMaster=no",
- "-o", "UserKnownHostsFile=/dev/null",
- "-o", "StrictHostKeyChecking=no"]
-
- ssh_tunnel_cmd = ["-L", tunnel_host,
- "echo -n ready && cat"
- ]
-
- ssh_cmd += ssh_tunnel_cmd
- self.log.debug("Creating tunnel with cmd: %s", ssh_cmd)
-
- proc = subprocess.Popen(ssh_cmd,
- stdin=subprocess.PIPE,
- stdout=subprocess.PIPE,
- close_fds=True)
- ready = proc.stdout.read(5)
- assert ready == b"ready", \
- "Did not get 'ready' from remote, got '{0}' instead".format(ready)
- yield
- proc.communicate()
- assert proc.returncode == 0, \
- "Tunnel process did unclean exit (returncode {}".format(proc.returncode)
-
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- if self.client is not None:
- self.client.close()
diff --git a/airflow/contrib/hooks/vertica_hook.py b/airflow/contrib/hooks/vertica_hook.py
deleted file mode 100644
index f3411de994..0000000000
--- a/airflow/contrib/hooks/vertica_hook.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 vertica_python import connect
-
-from airflow.hooks.dbapi_hook import DbApiHook
-
-
-class VerticaHook(DbApiHook):
- """
- Interact with Vertica.
- """
-
- conn_name_attr = 'vertica_conn_id'
- default_conn_name = 'vertica_default'
- supports_autocommit = True
-
- def get_conn(self):
- """
- Returns verticaql connection object
- """
- conn = self.get_connection(self.vertica_conn_id)
- conn_config = {
- "user": conn.login,
- "password": conn.password or '',
- "database": conn.schema,
- }
-
- conn_config["host"] = conn.host or 'localhost'
- if not conn.port:
- conn_config["port"] = 5433
- else:
- conn_config["port"] = int(conn.port)
-
- conn = connect(**conn_config)
- return conn
diff --git a/airflow/contrib/hooks/wasb_hook.py b/airflow/contrib/hooks/wasb_hook.py
deleted file mode 100644
index 1d73abd78b..0000000000
--- a/airflow/contrib/hooks/wasb_hook.py
+++ /dev/null
@@ -1,150 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.hooks.base_hook import BaseHook
-
-from azure.storage.blob import BlockBlobService
-
-
-class WasbHook(BaseHook):
- """
- Interacts with Azure Blob Storage through the wasb:// protocol.
-
- Additional options passed in the 'extra' field of the connection will be
- passed to the `BlockBlockService()` constructor. For example, authenticate
- using a SAS token by adding {"sas_token": "YOUR_TOKEN"}.
-
- :param wasb_conn_id: Reference to the wasb connection.
- :type wasb_conn_id: str
- """
-
- def __init__(self, wasb_conn_id='wasb_default'):
- self.conn_id = wasb_conn_id
- self.connection = self.get_conn()
-
- def get_conn(self):
- """Return the BlockBlobService object."""
- conn = self.get_connection(self.conn_id)
- service_options = conn.extra_dejson
- return BlockBlobService(account_name=conn.login,
- account_key=conn.password, **service_options)
-
- def check_for_blob(self, container_name, blob_name, **kwargs):
- """
- Check if a blob exists on Azure Blob Storage.
-
- :param container_name: Name of the container.
- :type container_name: str
- :param blob_name: Name of the blob.
- :type blob_name: str
- :param kwargs: Optional keyword arguments that
- `BlockBlobService.exists()` takes.
- :type kwargs: object
- :return: True if the blob exists, False otherwise.
- :rtype bool
- """
- return self.connection.exists(container_name, blob_name, **kwargs)
-
- def check_for_prefix(self, container_name, prefix, **kwargs):
- """
- Check if a prefix exists on Azure Blob storage.
-
- :param container_name: Name of the container.
- :type container_name: str
- :param prefix: Prefix of the blob.
- :type prefix: str
- :param kwargs: Optional keyword arguments that
- `BlockBlobService.list_blobs()` takes.
- :type kwargs: object
- :return: True if blobs matching the prefix exist, False otherwise.
- :rtype bool
- """
- matches = self.connection.list_blobs(container_name, prefix,
- num_results=1, **kwargs)
- return len(list(matches)) > 0
-
- def load_file(self, file_path, container_name, blob_name, **kwargs):
- """
- Upload a file to Azure Blob Storage.
-
- :param file_path: Path to the file to load.
- :type file_path: str
- :param container_name: Name of the container.
- :type container_name: str
- :param blob_name: Name of the blob.
- :type blob_name: str
- :param kwargs: Optional keyword arguments that
- `BlockBlobService.create_blob_from_path()` takes.
- :type kwargs: object
- """
- # Reorder the argument order from airflow.hooks.S3_hook.load_file.
- self.connection.create_blob_from_path(container_name, blob_name,
- file_path, **kwargs)
-
- def load_string(self, string_data, container_name, blob_name, **kwargs):
- """
- Upload a string to Azure Blob Storage.
-
- :param string_data: String to load.
- :type string_data: str
- :param container_name: Name of the container.
- :type container_name: str
- :param blob_name: Name of the blob.
- :type blob_name: str
- :param kwargs: Optional keyword arguments that
- `BlockBlobService.create_blob_from_text()` takes.
- :type kwargs: object
- """
- # Reorder the argument order from airflow.hooks.S3_hook.load_string.
- self.connection.create_blob_from_text(container_name, blob_name,
- string_data, **kwargs)
-
- def get_file(self, file_path, container_name, blob_name, **kwargs):
- """
- Download a file from Azure Blob Storage.
-
- :param file_path: Path to the file to download.
- :type file_path: str
- :param container_name: Name of the container.
- :type container_name: str
- :param blob_name: Name of the blob.
- :type blob_name: str
- :param kwargs: Optional keyword arguments that
- `BlockBlobService.create_blob_from_path()` takes.
- :type kwargs: object
- """
- return self.connection.get_blob_to_path(container_name, blob_name,
- file_path, **kwargs)
-
- def read_file(self, container_name, blob_name, **kwargs):
- """
- Read a file from Azure Blob Storage and return as a string.
-
- :param container_name: Name of the container.
- :type container_name: str
- :param blob_name: Name of the blob.
- :type blob_name: str
- :param kwargs: Optional keyword arguments that
- `BlockBlobService.create_blob_from_path()` takes.
- :type kwargs: object
- """
- return self.connection.get_blob_to_text(container_name,
- blob_name,
- **kwargs).content
diff --git a/airflow/contrib/hooks/winrm_hook.py b/airflow/contrib/hooks/winrm_hook.py
deleted file mode 100644
index 1dd02f9951..0000000000
--- a/airflow/contrib/hooks/winrm_hook.py
+++ /dev/null
@@ -1,251 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 getpass
-
-from winrm.protocol import Protocol
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base_hook import BaseHook
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-
-class WinRMHook(BaseHook, LoggingMixin):
- """
- Hook for winrm remote execution using pywinrm.
-
- :seealso: https://github.com/diyan/pywinrm/blob/master/winrm/protocol.py
-
- :param ssh_conn_id: connection id from airflow Connections from where all
- the required parameters can be fetched like username and password.
- Thought the priority is given to the param passed during init
- :type ssh_conn_id: string
- :param endpoint: When set to `None`, endpoint will be constructed like this:
- 'http://{remote_host}:{remote_port}/wsman'
- :type endpoint: string
- :param remote_host: Remote host to connect to.
- Ignored if `endpoint` is not `None`.
- :type remote_host: string
- :param remote_port: Remote port to connect to.
- Ignored if `endpoint` is not `None`.
- :type remote_port: int
- :param transport: transport type, one of 'plaintext' (default), 'kerberos', 'ssl',
- 'ntlm', 'credssp'
- :type transport: string
- :param username: username to connect to the remote_host
- :type username: string
- :param password: password of the username to connect to the remote_host
- :type password: string
- :param service: the service name, default is HTTP
- :type service: string
- :param keytab: the path to a keytab file if you are using one
- :type keytab: string
- :param ca_trust_path: Certification Authority trust path
- :type ca_trust_path: string
- :param cert_pem: client authentication certificate file path in PEM format
- :type cert_pem: string
- :param cert_key_pem: client authentication certificate key file path in PEM format
- :type cert_key_pem: string
- :param server_cert_validation: whether server certificate should be validated on
- Python versions that suppport it; one of 'validate' (default), 'ignore'
- :type server_cert_validation: string
- :param kerberos_delegation: if True, TGT is sent to target server to
- allow multiple hops
- :type kerberos_delegation: bool
- :param read_timeout_sec: maximum seconds to wait before an HTTP connect/read times out
- (default 30). This value should be slightly higher than operation_timeout_sec,
- as the server can block *at least* that long.
- :type read_timeout_sec: int
- :param operation_timeout_sec: maximum allowed time in seconds for any single wsman
- HTTP operation (default 20). Note that operation timeouts while receiving output
- (the only wsman operation that should take any significant time, and where these
- timeouts are expected) will be silently retried indefinitely.
- :type operation_timeout_sec: int
- :param kerberos_hostname_override: the hostname to use for the kerberos exchange
- (defaults to the hostname in the endpoint URL)
- :type kerberos_hostname_override: string
- :param message_encryption_enabled: Will encrypt the WinRM messages if set to True and
- the transport auth supports message encryption (Default True).
- :type message_encryption_enabled: bool
- :param credssp_disable_tlsv1_2: Whether to disable TLSv1.2 support and work with older
- protocols like TLSv1.0, default is False
- :type credssp_disable_tlsv1_2: bool
- :param send_cbt: Will send the channel bindings over a HTTPS channel (Default: True)
- :type send_cbt: bool
- """
- def __init__(self,
- ssh_conn_id=None,
- endpoint=None,
- remote_host=None,
- remote_port=5985,
- transport='plaintext',
- username=None,
- password=None,
- service='HTTP',
- keytab=None,
- ca_trust_path=None,
- cert_pem=None,
- cert_key_pem=None,
- server_cert_validation='validate',
- kerberos_delegation=False,
- read_timeout_sec=30,
- operation_timeout_sec=20,
- kerberos_hostname_override=None,
- message_encryption='auto',
- credssp_disable_tlsv1_2=False,
- send_cbt=True,
- ):
- super(WinRMHook, self).__init__(ssh_conn_id)
- self.ssh_conn_id = ssh_conn_id
- self.endpoint = endpoint
- self.remote_host = remote_host
- self.remote_port = remote_port
- self.transport = transport
- self.username = username
- self.password = password
- self.service = service
- self.keytab = keytab
- self.ca_trust_path = ca_trust_path
- self.cert_pem = cert_pem
- self.cert_key_pem = cert_key_pem
- self.server_cert_validation = server_cert_validation
- self.kerberos_delegation = kerberos_delegation
- self.read_timeout_sec = read_timeout_sec
- self.operation_timeout_sec = operation_timeout_sec
- self.kerberos_hostname_override = kerberos_hostname_override
- self.message_encryption = message_encryption
- self.credssp_disable_tlsv1_2 = credssp_disable_tlsv1_2
- self.send_cbt = send_cbt
-
- self.client = None
- self.winrm_protocol = None
-
- def get_conn(self):
- if self.client:
- return self.client
-
- self.log.debug('Creating WinRM client for conn_id: %s', self.ssh_conn_id)
- if self.ssh_conn_id is not None:
- conn = self.get_connection(self.ssh_conn_id)
-
- if self.username is None:
- self.username = conn.login
- if self.password is None:
- self.password = conn.password
- if self.remote_host is None:
- self.remote_host = conn.host
-
- if conn.extra is not None:
- extra_options = conn.extra_dejson
-
- if "endpoint" in extra_options:
- self.endpoint = str(extra_options["endpoint"])
- if "remote_port" in extra_options:
- self.remote_port = int(extra_options["remote_port"])
- if "transport" in extra_options:
- self.transport = str(extra_options["transport"])
- if "service" in extra_options:
- self.service = str(extra_options["service"])
- if "keytab" in extra_options:
- self.keytab = str(extra_options["keytab"])
- if "ca_trust_path" in extra_options:
- self.ca_trust_path = str(extra_options["ca_trust_path"])
- if "cert_pem" in extra_options:
- self.cert_pem = str(extra_options["cert_pem"])
- if "cert_key_pem" in extra_options:
- self.cert_key_pem = str(extra_options["cert_key_pem"])
- if "server_cert_validation" in extra_options:
- self.server_cert_validation = \
- str(extra_options["server_cert_validation"])
- if "kerberos_delegation" in extra_options:
- self.kerberos_delegation = \
- str(extra_options["kerberos_delegation"]).lower() == 'true'
- if "read_timeout_sec" in extra_options:
- self.read_timeout_sec = int(extra_options["read_timeout_sec"])
- if "operation_timeout_sec" in extra_options:
- self.operation_timeout_sec = \
- int(extra_options["operation_timeout_sec"])
- if "kerberos_hostname_override" in extra_options:
- self.kerberos_hostname_override = \
- str(extra_options["kerberos_hostname_override"])
- if "message_encryption" in extra_options:
- self.message_encryption = str(extra_options["message_encryption"])
- if "credssp_disable_tlsv1_2" in extra_options:
- self.credssp_disable_tlsv1_2 = \
- str(extra_options["credssp_disable_tlsv1_2"]).lower() == 'true'
- if "send_cbt" in extra_options:
- self.send_cbt = str(extra_options["send_cbt"]).lower() == 'true'
-
- if not self.remote_host:
- raise AirflowException("Missing required param: remote_host")
-
- # Auto detecting username values from system
- if not self.username:
- self.log.debug(
- "username to WinRM to host: %s is not specified for connection id"
- " %s. Using system's default provided by getpass.getuser()",
- self.remote_host, self.ssh_conn_id
- )
- self.username = getpass.getuser()
-
- # If endpoint is not set, then build a standard wsman endpoint from host and port.
- if not self.endpoint:
- self.endpoint = 'http://{0}:{1}/wsman'.format(
- self.remote_host,
- self.remote_port
- )
-
- try:
- if self.password and self.password.strip():
- self.winrm_protocol = Protocol(
- endpoint=self.endpoint,
- transport=self.transport,
- username=self.username,
- password=self.password,
- service=self.service,
- keytab=self.keytab,
- ca_trust_path=self.ca_trust_path,
- cert_pem=self.cert_pem,
- cert_key_pem=self.cert_key_pem,
- server_cert_validation=self.server_cert_validation,
- kerberos_delegation=self.kerberos_delegation,
- read_timeout_sec=self.read_timeout_sec,
- operation_timeout_sec=self.operation_timeout_sec,
- kerberos_hostname_override=self.kerberos_hostname_override,
- message_encryption=self.message_encryption,
- credssp_disable_tlsv1_2=self.credssp_disable_tlsv1_2,
- send_cbt=self.send_cbt
- )
-
- self.log.info(
- "Establishing WinRM connection to host: %s",
- self.remote_host
- )
- self.client = self.winrm_protocol.open_shell()
-
- except Exception as error:
- error_msg = "Error connecting to host: {0}, error: {1}".format(
- self.remote_host,
- error
- )
- self.log.error(error_msg)
- raise AirflowException(error_msg)
-
- return self.client
diff --git a/airflow/contrib/operators/__init__.py b/airflow/contrib/operators/__init__.py
deleted file mode 100644
index 3485aaa7fd..0000000000
--- a/airflow/contrib/operators/__init__.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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.
-#
-
-
-# Contrib operators are not imported by default. They should be accessed
-# directly: from airflow.contrib.operators.operator_module import Operator
-
-
-import sys
-import os as _os
-
-# ------------------------------------------------------------------------
-#
-# #TODO #FIXME Airflow 2.0
-#
-# Old import machinary below.
-#
-# This is deprecated but should be kept until Airflow 2.0
-# for compatibility.
-#
-# ------------------------------------------------------------------------
-_operators = {
- 'ssh_operator': ['SSHOperator'],
- 'winrm_operator': ['WinRMOperator'],
- 'vertica_operator': ['VerticaOperator'],
- 'vertica_to_hive': ['VerticaToHiveTransfer'],
- 'qubole_operator': ['QuboleOperator'],
- 'spark_submit_operator': ['SparkSubmitOperator'],
- 'file_to_wasb': ['FileToWasbOperator'],
- 'fs_operator': ['FileSensor'],
- 'hive_to_dynamodb': ['HiveToDynamoDBTransferOperator']
-}
-
-if not _os.environ.get('AIRFLOW_USE_NEW_IMPORTS', False):
- from airflow.utils.helpers import AirflowImporter
- airflow_importer = AirflowImporter(sys.modules[__name__], _operators)
diff --git a/airflow/contrib/operators/awsbatch_operator.py b/airflow/contrib/operators/awsbatch_operator.py
deleted file mode 100644
index 75706aa189..0000000000
--- a/airflow/contrib/operators/awsbatch_operator.py
+++ /dev/null
@@ -1,179 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 sys
-
-from math import pow
-from time import sleep
-
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
-
-from airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class AWSBatchOperator(BaseOperator):
- """
- Execute a job on AWS Batch Service
-
- .. warning: the queue parameter was renamed to job_queue to segreggate the
- internal CeleryExecutor queue from the AWS Batch internal queue.
-
- :param job_name: the name for the job that will run on AWS Batch (templated)
- :type job_name: str
- :param job_definition: the job definition name on AWS Batch
- :type job_definition: str
- :param job_queue: the queue name on AWS Batch
- :type job_queue: str
- :param: overrides: the same parameter that boto3 will receive on
- containerOverrides (templated):
- http://boto3.readthedocs.io/en/latest/reference/services/batch.html#submit_job
- :type: overrides: dict
- :param max_retries: exponential backoff retries while waiter is not merged, 4200 = 48 hours
- :type max_retries: int
- :param aws_conn_id: connection id of AWS credentials / region name. If None,
- credential boto3 strategy will be used
- (http://boto3.readthedocs.io/en/latest/guide/configuration.html).
- :type aws_conn_id: str
- :param region_name: region name to use in AWS Hook.
- Override the region_name in connection (if provided)
- """
-
- ui_color = '#c3dae0'
- client = None
- arn = None
- template_fields = ('job_name', 'overrides',)
-
- @apply_defaults
- def __init__(self, job_name, job_definition, job_queue, overrides, max_retries=4200,
- aws_conn_id=None, region_name=None, **kwargs):
- super(AWSBatchOperator, self).__init__(**kwargs)
-
- self.job_name = job_name
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
- self.job_definition = job_definition
- self.job_queue = job_queue
- self.overrides = overrides
- self.max_retries = max_retries
-
- self.jobId = None
- self.jobName = None
-
- self.hook = self.get_hook()
-
- def execute(self, context):
- self.log.info(
- 'Running AWS Batch Job - Job definition: %s - on queue %s',
- self.job_definition, self.job_queue
- )
- self.log.info('AWSBatchOperator overrides: %s', self.overrides)
-
- self.client = self.hook.get_client_type(
- 'batch',
- region_name=self.region_name
- )
-
- try:
- response = self.client.submit_job(
- jobName=self.job_name,
- jobQueue=self.job_queue,
- jobDefinition=self.job_definition,
- containerOverrides=self.overrides)
-
- self.log.info('AWS Batch Job started: %s', response)
-
- self.jobId = response['jobId']
- self.jobName = response['jobName']
-
- self._wait_for_task_ended()
-
- self._check_success_task()
-
- self.log.info('AWS Batch Job has been successfully executed: %s', response)
- except Exception as e:
- self.log.info('AWS Batch Job has failed executed')
- raise AirflowException(e)
-
- def _wait_for_task_ended(self):
- """
- Try to use a waiter from the below pull request
-
- * https://github.com/boto/botocore/pull/1307
-
- If the waiter is not available apply a exponential backoff
-
- * docs.aws.amazon.com/general/latest/gr/api-retries.html
- """
- try:
- waiter = self.client.get_waiter('job_execution_complete')
- waiter.config.max_attempts = sys.maxsize # timeout is managed by airflow
- waiter.wait(jobs=[self.jobId])
- except ValueError:
- # If waiter not available use expo
- retry = True
- retries = 0
-
- while retries < self.max_retries and retry:
- self.log.info('AWS Batch retry in the next %s seconds', retries)
- response = self.client.describe_jobs(
- jobs=[self.jobId]
- )
- if response['jobs'][-1]['status'] in ['SUCCEEDED', 'FAILED']:
- retry = False
-
- sleep( 1 + pow(retries * 0.1, 2))
- retries += 1
-
- def _check_success_task(self):
- response = self.client.describe_jobs(
- jobs=[self.jobId],
- )
-
- self.log.info('AWS Batch stopped, check status: %s', response)
- if len(response.get('jobs')) < 1:
- raise AirflowException('No job found for {}'.format(response))
-
- for job in response['jobs']:
- job_status = job['status']
- if job_status is 'FAILED':
- reason = job['statusReason']
- raise AirflowException('Job failed with status {}'.format(reason))
- elif job_status in [
- 'SUBMITTED',
- 'PENDING',
- 'RUNNABLE',
- 'STARTING',
- 'RUNNING'
- ]:
- raise AirflowException(
- 'This task is still pending {}'.format(job_status))
-
- def get_hook(self):
- return AwsHook(
- aws_conn_id=self.aws_conn_id
- )
-
- def on_kill(self):
- response = self.client.terminate_job(
- jobId=self.jobId,
- reason='Task killed by the user')
-
- self.log.info(response)
diff --git a/airflow/contrib/operators/bigquery_check_operator.py b/airflow/contrib/operators/bigquery_check_operator.py
deleted file mode 100644
index 59ef5d377d..0000000000
--- a/airflow/contrib/operators/bigquery_check_operator.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.operators.check_operator import \
- CheckOperator, ValueCheckOperator, IntervalCheckOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class BigQueryCheckOperator(CheckOperator):
- """
- Performs checks against BigQuery. The ``BigQueryCheckOperator`` expects
- a sql query that will return a single row. Each value on that
- first row is evaluated using python ``bool`` casting. If any of the
- values return ``False`` the check is failed and errors out.
-
- Note that Python bool casting evals the following as ``False``:
-
- * ``False``
- * ``0``
- * Empty string (``""``)
- * Empty list (``[]``)
- * Empty dictionary or set (``{}``)
-
- Given a query like ``SELECT COUNT(*) FROM foo``, it will fail only if
- the count ``== 0``. You can craft much more complex query that could,
- for instance, check that the table has the same number of rows as
- the source table upstream, or that the count of today's partition is
- greater than yesterday's partition, or that a set of metrics are less
- than 3 standard deviation for the 7 day average.
-
- This operator can be used as a data quality check in your pipeline, and
- depending on where you put it in your DAG, you have the choice to
- stop the critical path, preventing from
- publishing dubious data, or on the side and receive email alterts
- without stopping the progress of the DAG.
-
- :param sql: the sql to be executed
- :type sql: string
- :param bigquery_conn_id: reference to the BigQuery database
- :type bigquery_conn_id: string
- """
-
- @apply_defaults
- def __init__(
- self,
- sql,
- bigquery_conn_id='bigquery_default',
- *args,
- **kwargs):
- super(BigQueryCheckOperator, self).__init__(sql=sql, *args, **kwargs)
- self.bigquery_conn_id = bigquery_conn_id
- self.sql = sql
-
- def get_db_hook(self):
- return BigQueryHook(bigquery_conn_id=self.bigquery_conn_id)
-
-
-class BigQueryValueCheckOperator(ValueCheckOperator):
- """
- Performs a simple value check using sql code.
-
- :param sql: the sql to be executed
- :type sql: string
- """
-
- @apply_defaults
- def __init__(
- self, sql, pass_value, tolerance=None,
- bigquery_conn_id='bigquery_default',
- *args, **kwargs):
- super(BigQueryValueCheckOperator, self).__init__(
- sql=sql, pass_value=pass_value, tolerance=tolerance,
- *args, **kwargs)
- self.bigquery_conn_id = bigquery_conn_id
-
- def get_db_hook(self):
- return BigQueryHook(bigquery_conn_id=self.bigquery_conn_id)
-
-
-class BigQueryIntervalCheckOperator(IntervalCheckOperator):
- """
- Checks that the values of metrics given as SQL expressions are within
- a certain tolerance of the ones from days_back before.
-
- This method constructs a query like so ::
-
- SELECT {metrics_threshold_dict_key} FROM {table}
- WHERE {date_filter_column}=<date>
-
- :param table: the table name
- :type table: str
- :param days_back: number of days between ds and the ds we want to check
- against. Defaults to 7 days
- :type days_back: int
- :param metrics_threshold: a dictionary of ratios indexed by metrics, for
- example 'COUNT(*)': 1.5 would require a 50 percent or less difference
- between the current day, and the prior days_back.
- :type metrics_threshold: dict
- """
-
- @apply_defaults
- def __init__(
- self, table, metrics_thresholds,
- date_filter_column='ds', days_back=-7,
- bigquery_conn_id='bigquery_default',
- *args, **kwargs):
- super(BigQueryIntervalCheckOperator, self).__init__(
- table=table, metrics_thresholds=metrics_thresholds,
- date_filter_column=date_filter_column, days_back=days_back,
- *args, **kwargs)
- self.bigquery_conn_id = bigquery_conn_id
-
- def get_db_hook(self):
- return BigQueryHook(bigquery_conn_id=self.bigquery_conn_id)
diff --git a/airflow/contrib/operators/bigquery_get_data.py b/airflow/contrib/operators/bigquery_get_data.py
deleted file mode 100644
index ab8f71b717..0000000000
--- a/airflow/contrib/operators/bigquery_get_data.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class BigQueryGetDataOperator(BaseOperator):
- """
- Fetches the data from a BigQuery table (alternatively fetch data for selected columns)
- and returns data in a python list. The number of elements in the returned list will
- be equal to the number of rows fetched. Each element in the list will again be a list
- where element would represent the columns values for that row.
-
- **Example Result**: ``[['Tony', '10'], ['Mike', '20'], ['Steve', '15']]``
-
- .. note::
- If you pass fields to ``selected_fields`` which are in different order than the
- order of columns already in
- BQ table, the data will still be in the order of BQ table.
- For example if the BQ table has 3 columns as
- ``[A,B,C]`` and you pass 'B,A' in the ``selected_fields``
- the data would still be of the form ``'A,B'``.
-
- **Example**: ::
-
- get_data = BigQueryGetDataOperator(
- task_id='get_data_from_bq',
- dataset_id='test_dataset',
- table_id='Transaction_partitions',
- max_results='100',
- selected_fields='DATE',
- bigquery_conn_id='airflow-service-account'
- )
-
- :param dataset_id: The dataset ID of the requested table. (templated)
- :type destination_dataset_table: string
- :param table_id: The table ID of the requested table. (templated)
- :type table_id: string
- :param max_results: The maximum number of records (rows) to be fetched
- from the table. (templated)
- :type max_results: string
- :param selected_fields: List of fields to return (comma-separated). If
- unspecified, all fields are returned.
- :type selected_fields: string
- :param bigquery_conn_id: reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
- template_fields = ('dataset_id', 'table_id', 'max_results')
- ui_color = '#e4f0e8'
-
- @apply_defaults
- def __init__(self,
- dataset_id,
- table_id,
- max_results='100',
- selected_fields=None,
- bigquery_conn_id='bigquery_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(BigQueryGetDataOperator, self).__init__(*args, **kwargs)
- self.dataset_id = dataset_id
- self.table_id = table_id
- self.max_results = max_results
- self.selected_fields = selected_fields
- self.bigquery_conn_id = bigquery_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- self.log.info('Fetching Data from:')
- self.log.info('Dataset: %s ; Table: %s ; Max Results: %s',
- self.dataset_id, self.table_id, self.max_results)
-
- hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
-
- conn = hook.get_conn()
- cursor = conn.cursor()
- response = cursor.get_tabledata(dataset_id=self.dataset_id,
- table_id=self.table_id,
- max_results=self.max_results,
- selected_fields=self.selected_fields)
-
- self.log.info('Total Extracted rows: %s', response['totalRows'])
- rows = response['rows']
-
- table_data = []
- for dict_row in rows:
- single_row = []
- for fields in dict_row['f']:
- single_row.append(fields['v'])
- table_data.append(single_row)
-
- return table_data
diff --git a/airflow/contrib/operators/bigquery_operator.py b/airflow/contrib/operators/bigquery_operator.py
deleted file mode 100644
index b36efbd6bf..0000000000
--- a/airflow/contrib/operators/bigquery_operator.py
+++ /dev/null
@@ -1,483 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-
-from airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook, _parse_gcs_url
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class BigQueryOperator(BaseOperator):
- """
- Executes BigQuery SQL queries in a specific BigQuery database
-
- :param bql: (Deprecated. Use `sql` parameter instead) the sql code to be
- executed (templated)
- :type bql: Can receive a str representing a sql statement,
- a list of str (sql statements), or reference to a template file.
- Template reference are recognized by str ending in '.sql'.
- :param sql: the sql code to be executed (templated)
- :type sql: Can receive a str representing a sql statement,
- a list of str (sql statements), or reference to a template file.
- Template reference are recognized by str ending in '.sql'.
- :param destination_dataset_table: A dotted
- (<project>.|<project>:)<dataset>.<table> that, if set, will store the results
- of the query. (templated)
- :type destination_dataset_table: string
- :param write_disposition: Specifies the action that occurs if the destination table
- already exists. (default: 'WRITE_EMPTY')
- :type write_disposition: string
- :param create_disposition: Specifies whether the job is allowed to create new tables.
- (default: 'CREATE_IF_NEEDED')
- :type create_disposition: string
- :param allow_large_results: Whether to allow large results.
- :type allow_large_results: boolean
- :param flatten_results: If true and query uses legacy SQL dialect, flattens
- all nested and repeated fields in the query results. ``allow_large_results``
- must be ``true`` if this is set to ``false``. For standard SQL queries, this
- flag is ignored and results are never flattened.
- :type flatten_results: boolean
- :param bigquery_conn_id: reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param udf_config: The User Defined Function configuration for the query.
- See https://cloud.google.com/bigquery/user-defined-functions for details.
- :type udf_config: list
- :param use_legacy_sql: Whether to use legacy SQL (true) or standard SQL (false).
- :type use_legacy_sql: boolean
- :param maximum_billing_tier: Positive integer that serves as a multiplier
- of the basic price.
- Defaults to None, in which case it uses the value set in the project.
- :type maximum_billing_tier: integer
- :param maximum_bytes_billed: Limits the bytes billed for this job.
- Queries that will have bytes billed beyond this limit will fail
- (without incurring a charge). If unspecified, this will be
- set to your project default.
- :type maximum_bytes_billed: float
- :param schema_update_options: Allows the schema of the destination
- table to be updated as a side effect of the load job.
- :type schema_update_options: tuple
- :param query_params: a dictionary containing query parameter types and
- values, passed to BigQuery.
- :type query_params: dict
- :param priority: Specifies a priority for the query.
- Possible values include INTERACTIVE and BATCH.
- The default value is INTERACTIVE.
- :type priority: string
- :param time_partitioning: configure optional time partitioning fields i.e.
- partition by field, type and
- expiration as per API specifications. Note that 'field' is not available in
- conjunction with dataset.table$partition.
- :type time_partitioning: dict
- """
-
- template_fields = ('bql', 'sql', 'destination_dataset_table')
- template_ext = ('.sql', )
- ui_color = '#e4f0e8'
-
- @apply_defaults
- def __init__(self,
- bql=None,
- sql=None,
- destination_dataset_table=False,
- write_disposition='WRITE_EMPTY',
- allow_large_results=False,
- flatten_results=False,
- bigquery_conn_id='bigquery_default',
- delegate_to=None,
- udf_config=False,
- use_legacy_sql=True,
- maximum_billing_tier=None,
- maximum_bytes_billed=None,
- create_disposition='CREATE_IF_NEEDED',
- schema_update_options=(),
- query_params=None,
- priority='INTERACTIVE',
- time_partitioning={},
- *args,
- **kwargs):
- super(BigQueryOperator, self).__init__(*args, **kwargs)
- self.bql = bql
- self.sql = sql if sql else bql
- self.destination_dataset_table = destination_dataset_table
- self.write_disposition = write_disposition
- self.create_disposition = create_disposition
- self.allow_large_results = allow_large_results
- self.flatten_results = flatten_results
- self.bigquery_conn_id = bigquery_conn_id
- self.delegate_to = delegate_to
- self.udf_config = udf_config
- self.use_legacy_sql = use_legacy_sql
- self.maximum_billing_tier = maximum_billing_tier
- self.maximum_bytes_billed = maximum_bytes_billed
- self.schema_update_options = schema_update_options
- self.query_params = query_params
- self.bq_cursor = None
- self.priority = priority
- self.time_partitioning = time_partitioning
-
- # TODO remove `bql` in Airflow 2.0
- if self.bql:
- import warnings
- warnings.warn('Deprecated parameter `bql` used in Task id: {}. '
- 'Use `sql` parameter instead to pass the sql to be '
- 'executed. `bql` parameter is deprecated and '
- 'will be removed in a future version of '
- 'Airflow.'.format(self.task_id),
- category=DeprecationWarning)
-
- if self.sql is None:
- raise TypeError('{} missing 1 required positional '
- 'argument: `sql`'.format(self.task_id))
-
- def execute(self, context):
- if self.bq_cursor is None:
- self.log.info('Executing: %s', self.sql)
- hook = BigQueryHook(
- bigquery_conn_id=self.bigquery_conn_id,
- use_legacy_sql=self.use_legacy_sql,
- delegate_to=self.delegate_to)
- conn = hook.get_conn()
- self.bq_cursor = conn.cursor()
- self.bq_cursor.run_query(
- self.sql,
- destination_dataset_table=self.destination_dataset_table,
- write_disposition=self.write_disposition,
- allow_large_results=self.allow_large_results,
- flatten_results=self.flatten_results,
- udf_config=self.udf_config,
- maximum_billing_tier=self.maximum_billing_tier,
- maximum_bytes_billed=self.maximum_bytes_billed,
- create_disposition=self.create_disposition,
- query_params=self.query_params,
- schema_update_options=self.schema_update_options,
- priority=self.priority,
- time_partitioning=self.time_partitioning
- )
-
- def on_kill(self):
- super(BigQueryOperator, self).on_kill()
- if self.bq_cursor is not None:
- self.log.info('Canceling running query due to execution timeout')
- self.bq_cursor.cancel_query()
-
-
-class BigQueryCreateEmptyTableOperator(BaseOperator):
- """
- Creates a new, empty table in the specified BigQuery dataset,
- optionally with schema.
-
- The schema to be used for the BigQuery table may be specified in one of
- two ways. You may either directly pass the schema fields in, or you may
- point the operator to a Google cloud storage object name. The object in
- Google cloud storage must be a JSON file with the schema fields in it.
- You can also create a table without schema.
-
- :param project_id: The project to create the table into. (templated)
- :type project_id: string
- :param dataset_id: The dataset to create the table into. (templated)
- :type dataset_id: string
- :param table_id: The Name of the table to be created. (templated)
- :type table_id: string
- :param schema_fields: If set, the schema field list as defined here:
- https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schema
-
- **Example**: ::
-
- schema_fields=[{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"},
- {"name": "salary", "type": "INTEGER", "mode": "NULLABLE"}]
-
- :type schema_fields: list
- :param gcs_schema_object: Full path to the JSON file containing
- schema (templated). For
- example: ``gs://test-bucket/dir1/dir2/employee_schema.json``
- :type gcs_schema_object: string
- :param time_partitioning: configure optional time partitioning fields i.e.
- partition by field, type and expiration as per API specifications.
-
- .. seealso::
- https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#timePartitioning
- :type time_partitioning: dict
- :param bigquery_conn_id: Reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param google_cloud_storage_conn_id: Reference to a specific Google
- cloud storage hook.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any. For this to
- work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
-
- **Example (with schema JSON in GCS)**: ::
-
- CreateTable = BigQueryCreateEmptyTableOperator(
- task_id='BigQueryCreateEmptyTableOperator_task',
- dataset_id='ODS',
- table_id='Employees',
- project_id='internal-gcp-project',
- gcs_schema_object='gs://schema-bucket/employee_schema.json',
- bigquery_conn_id='airflow-service-account',
- google_cloud_storage_conn_id='airflow-service-account'
- )
-
- **Corresponding Schema file** (``employee_schema.json``): ::
-
- [
- {
- "mode": "NULLABLE",
- "name": "emp_name",
- "type": "STRING"
- },
- {
- "mode": "REQUIRED",
- "name": "salary",
- "type": "INTEGER"
- }
- ]
-
- **Example (with schema in the DAG)**: ::
-
- CreateTable = BigQueryCreateEmptyTableOperator(
- task_id='BigQueryCreateEmptyTableOperator_task',
- dataset_id='ODS',
- table_id='Employees',
- project_id='internal-gcp-project',
- schema_fields=[{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"},
- {"name": "salary", "type": "INTEGER", "mode": "NULLABLE"}],
- bigquery_conn_id='airflow-service-account',
- google_cloud_storage_conn_id='airflow-service-account'
- )
-
- """
- template_fields = ('dataset_id', 'table_id', 'project_id', 'gcs_schema_object')
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- dataset_id,
- table_id,
- project_id=None,
- schema_fields=None,
- gcs_schema_object=None,
- time_partitioning={},
- bigquery_conn_id='bigquery_default',
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- *args, **kwargs):
-
- super(BigQueryCreateEmptyTableOperator, self).__init__(*args, **kwargs)
-
- self.project_id = project_id
- self.dataset_id = dataset_id
- self.table_id = table_id
- self.schema_fields = schema_fields
- self.gcs_schema_object = gcs_schema_object
- self.bigquery_conn_id = bigquery_conn_id
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
- self.time_partitioning = time_partitioning
-
- def execute(self, context):
- bq_hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
-
- if not self.schema_fields and self.gcs_schema_object:
-
- gcs_bucket, gcs_object = _parse_gcs_url(self.gcs_schema_object)
-
- gcs_hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- schema_fields = json.loads(gcs_hook.download(
- gcs_bucket,
- gcs_object).decode("utf-8"))
- else:
- schema_fields = self.schema_fields
-
- conn = bq_hook.get_conn()
- cursor = conn.cursor()
-
- cursor.create_empty_table(
- project_id=self.project_id,
- dataset_id=self.dataset_id,
- table_id=self.table_id,
- schema_fields=schema_fields,
- time_partitioning=self.time_partitioning
- )
-
-
-class BigQueryCreateExternalTableOperator(BaseOperator):
- """
- Creates a new external table in the dataset with the data in Google Cloud
- Storage.
-
- The schema to be used for the BigQuery table may be specified in one of
- two ways. You may either directly pass the schema fields in, or you may
- point the operator to a Google cloud storage object name. The object in
- Google cloud storage must be a JSON file with the schema fields in it.
-
- :param bucket: The bucket to point the external table to. (templated)
- :type bucket: string
- :param source_objects: List of Google cloud storage URIs to point
- table to. (templated)
- If source_format is 'DATASTORE_BACKUP', the list must only contain a single URI.
- :type object: list
- :param destination_project_dataset_table: The dotted (<project>.)<dataset>.<table>
- BigQuery table to load data into (templated). If <project> is not included,
- project will be the project defined in the connection json.
- :type destination_project_dataset_table: string
- :param schema_fields: If set, the schema field list as defined here:
- https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schema
-
- **Example**: ::
-
- schema_fields=[{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"},
- {"name": "salary", "type": "INTEGER", "mode": "NULLABLE"}]
-
- Should not be set when source_format is 'DATASTORE_BACKUP'.
- :type schema_fields: list
- :param schema_object: If set, a GCS object path pointing to a .json file that
- contains the schema for the table. (templated)
- :param schema_object: string
- :param source_format: File format of the data.
- :type source_format: string
- :param compression: [Optional] The compression type of the data source.
- Possible values include GZIP and NONE.
- The default value is NONE.
- This setting is ignored for Google Cloud Bigtable,
- Google Cloud Datastore backups and Avro formats.
- :type compression: string
- :param skip_leading_rows: Number of rows to skip when loading from a CSV.
- :type skip_leading_rows: int
- :param field_delimiter: The delimiter to use for the CSV.
- :type field_delimiter: string
- :param max_bad_records: The maximum number of bad records that BigQuery can
- ignore when running the job.
- :type max_bad_records: int
- :param quote_character: The value that is used to quote data sections in a CSV file.
- :type quote_character: string
- :param allow_quoted_newlines: Whether to allow quoted newlines (true) or not (false).
- :type allow_quoted_newlines: boolean
- :param allow_jagged_rows: Accept rows that are missing trailing optional columns.
- The missing values are treated as nulls. If false, records with missing trailing
- columns are treated as bad records, and if there are too many bad records, an
- invalid error is returned in the job result. Only applicable to CSV, ignored
- for other formats.
- :type allow_jagged_rows: bool
- :param bigquery_conn_id: Reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param google_cloud_storage_conn_id: Reference to a specific Google
- cloud storage hook.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any. For this to
- work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param src_fmt_configs: configure optional fields specific to the source format
- :type src_fmt_configs: dict
- """
- template_fields = ('bucket', 'source_objects',
- 'schema_object', 'destination_project_dataset_table')
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- bucket,
- source_objects,
- destination_project_dataset_table,
- schema_fields=None,
- schema_object=None,
- source_format='CSV',
- compression='NONE',
- skip_leading_rows=0,
- field_delimiter=',',
- max_bad_records=0,
- quote_character=None,
- allow_quoted_newlines=False,
- allow_jagged_rows=False,
- bigquery_conn_id='bigquery_default',
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- src_fmt_configs={},
- *args, **kwargs):
-
- super(BigQueryCreateExternalTableOperator, self).__init__(*args, **kwargs)
-
- # GCS config
- self.bucket = bucket
- self.source_objects = source_objects
- self.schema_object = schema_object
-
- # BQ config
- self.destination_project_dataset_table = destination_project_dataset_table
- self.schema_fields = schema_fields
- self.source_format = source_format
- self.compression = compression
- self.skip_leading_rows = skip_leading_rows
- self.field_delimiter = field_delimiter
- self.max_bad_records = max_bad_records
- self.quote_character = quote_character
- self.allow_quoted_newlines = allow_quoted_newlines
- self.allow_jagged_rows = allow_jagged_rows
-
- self.bigquery_conn_id = bigquery_conn_id
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
-
- self.src_fmt_configs = src_fmt_configs
-
- def execute(self, context):
- bq_hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
-
- if not self.schema_fields and self.schema_object \
- and self.source_format != 'DATASTORE_BACKUP':
- gcs_hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- schema_fields = json.loads(gcs_hook.download(
- self.bucket,
- self.schema_object).decode("utf-8"))
- else:
- schema_fields = self.schema_fields
-
- source_uris = ['gs://{}/{}'.format(self.bucket, source_object)
- for source_object in self.source_objects]
- conn = bq_hook.get_conn()
- cursor = conn.cursor()
-
- cursor.create_external_table(
- external_project_dataset_table=self.destination_project_dataset_table,
- schema_fields=schema_fields,
- source_uris=source_uris,
- source_format=self.source_format,
- compression=self.compression,
- skip_leading_rows=self.skip_leading_rows,
- field_delimiter=self.field_delimiter,
- max_bad_records=self.max_bad_records,
- quote_character=self.quote_character,
- allow_quoted_newlines=self.allow_quoted_newlines,
- allow_jagged_rows=self.allow_jagged_rows,
- src_fmt_configs=self.src_fmt_configs
- )
diff --git a/airflow/contrib/operators/bigquery_table_delete_operator.py b/airflow/contrib/operators/bigquery_table_delete_operator.py
deleted file mode 100644
index a16107d8c4..0000000000
--- a/airflow/contrib/operators/bigquery_table_delete_operator.py
+++ /dev/null
@@ -1,66 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class BigQueryTableDeleteOperator(BaseOperator):
- """
- Deletes BigQuery tables
-
- :param deletion_dataset_table: A dotted
- (<project>.|<project>:)<dataset>.<table> that indicates which table
- will be deleted. (templated)
- :type deletion_dataset_table: string
- :param bigquery_conn_id: reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param ignore_if_missing: if True, then return success even if the
- requested table does not exist.
- :type ignore_if_missing: boolean
- """
- template_fields = ('deletion_dataset_table',)
- ui_color = '#ffd1dc'
-
- @apply_defaults
- def __init__(self,
- deletion_dataset_table,
- bigquery_conn_id='bigquery_default',
- delegate_to=None,
- ignore_if_missing=False,
- *args,
- **kwargs):
- super(BigQueryTableDeleteOperator, self).__init__(*args, **kwargs)
- self.deletion_dataset_table = deletion_dataset_table
- self.bigquery_conn_id = bigquery_conn_id
- self.delegate_to = delegate_to
- self.ignore_if_missing = ignore_if_missing
-
- def execute(self, context):
- self.log.info('Deleting: %s', self.deletion_dataset_table)
- hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
- conn = hook.get_conn()
- cursor = conn.cursor()
- cursor.run_table_delete(self.deletion_dataset_table, self.ignore_if_missing)
diff --git a/airflow/contrib/operators/bigquery_to_bigquery.py b/airflow/contrib/operators/bigquery_to_bigquery.py
deleted file mode 100644
index 93a52b3102..0000000000
--- a/airflow/contrib/operators/bigquery_to_bigquery.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class BigQueryToBigQueryOperator(BaseOperator):
- """
- Copies data from one BigQuery table to another.
-
- .. seealso::
- For more details about these parameters:
- https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.copy
-
- :param source_project_dataset_tables: One or more
- dotted (project:|project.)<dataset>.<table> BigQuery tables to use as the
- source data. If <project> is not included, project will be the
- project defined in the connection json. Use a list if there are multiple
- source tables. (templated)
- :type source_project_dataset_tables: list|string
- :param destination_project_dataset_table: The destination BigQuery
- table. Format is: (project:|project.)<dataset>.<table> (templated)
- :type destination_project_dataset_table: string
- :param write_disposition: The write disposition if the table already exists.
- :type write_disposition: string
- :param create_disposition: The create disposition if the table doesn't exist.
- :type create_disposition: string
- :param bigquery_conn_id: reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
- template_fields = ('source_project_dataset_tables',
- 'destination_project_dataset_table')
- template_ext = ('.sql',)
- ui_color = '#e6f0e4'
-
- @apply_defaults
- def __init__(self,
- source_project_dataset_tables,
- destination_project_dataset_table,
- write_disposition='WRITE_EMPTY',
- create_disposition='CREATE_IF_NEEDED',
- bigquery_conn_id='bigquery_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(BigQueryToBigQueryOperator, self).__init__(*args, **kwargs)
- self.source_project_dataset_tables = source_project_dataset_tables
- self.destination_project_dataset_table = destination_project_dataset_table
- self.write_disposition = write_disposition
- self.create_disposition = create_disposition
- self.bigquery_conn_id = bigquery_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- self.log.info(
- 'Executing copy of %s into: %s',
- self.source_project_dataset_tables, self.destination_project_dataset_table
- )
- hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
- conn = hook.get_conn()
- cursor = conn.cursor()
- cursor.run_copy(
- self.source_project_dataset_tables,
- self.destination_project_dataset_table,
- self.write_disposition,
- self.create_disposition)
diff --git a/airflow/contrib/operators/bigquery_to_gcs.py b/airflow/contrib/operators/bigquery_to_gcs.py
deleted file mode 100644
index e2ce93068f..0000000000
--- a/airflow/contrib/operators/bigquery_to_gcs.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class BigQueryToCloudStorageOperator(BaseOperator):
- """
- Transfers a BigQuery table to a Google Cloud Storage bucket.
-
- .. seealso::
- For more details about these parameters:
- https://cloud.google.com/bigquery/docs/reference/v2/jobs
-
- :param source_project_dataset_table: The dotted
- (<project>.|<project>:)<dataset>.<table> BigQuery table to use as the source
- data. If <project> is not included, project will be the project
- defined in the connection json. (templated)
- :type source_project_dataset_table: string
- :param destination_cloud_storage_uris: The destination Google Cloud
- Storage URI (e.g. gs://some-bucket/some-file.txt). (templated) Follows
- convention defined here:
- https://cloud.google.com/bigquery/exporting-data-from-bigquery#exportingmultiple
- :type destination_cloud_storage_uris: list
- :param compression: Type of compression to use.
- :type compression: string
- :param export_format: File format to export.
- :type field_delimiter: string
- :param field_delimiter: The delimiter to use when extracting to a CSV.
- :type field_delimiter: string
- :param print_header: Whether to print a header for a CSV file extract.
- :type print_header: boolean
- :param bigquery_conn_id: reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
- template_fields = ('source_project_dataset_table', 'destination_cloud_storage_uris')
- template_ext = ('.sql',)
- ui_color = '#e4e6f0'
-
- @apply_defaults
- def __init__(self,
- source_project_dataset_table,
- destination_cloud_storage_uris,
- compression='NONE',
- export_format='CSV',
- field_delimiter=',',
- print_header=True,
- bigquery_conn_id='bigquery_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(BigQueryToCloudStorageOperator, self).__init__(*args, **kwargs)
- self.source_project_dataset_table = source_project_dataset_table
- self.destination_cloud_storage_uris = destination_cloud_storage_uris
- self.compression = compression
- self.export_format = export_format
- self.field_delimiter = field_delimiter
- self.print_header = print_header
- self.bigquery_conn_id = bigquery_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- self.log.info('Executing extract of %s into: %s',
- self.source_project_dataset_table,
- self.destination_cloud_storage_uris)
- hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
- conn = hook.get_conn()
- cursor = conn.cursor()
- cursor.run_extract(
- self.source_project_dataset_table,
- self.destination_cloud_storage_uris,
- self.compression,
- self.export_format,
- self.field_delimiter,
- self.print_header)
diff --git a/airflow/contrib/operators/cassandra_to_gcs.py b/airflow/contrib/operators/cassandra_to_gcs.py
deleted file mode 100644
index 211444b96d..0000000000
--- a/airflow/contrib/operators/cassandra_to_gcs.py
+++ /dev/null
@@ -1,355 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 __future__ import unicode_literals
-
-import json
-from builtins import str
-from base64 import b64encode
-from cassandra.util import Date, Time, SortedSet, OrderedMapSerializedKey
-from datetime import datetime
-from decimal import Decimal
-from six import text_type, binary_type, PY3
-from tempfile import NamedTemporaryFile
-from uuid import UUID
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.contrib.hooks.cassandra_hook import CassandraHook
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class CassandraToGoogleCloudStorageOperator(BaseOperator):
- """
- Copy data from Cassandra to Google cloud storage in JSON format
-
- Note: Arrays of arrays are not supported.
- """
- template_fields = ('cql', 'bucket', 'filename', 'schema_filename',)
- template_ext = ('.cql',)
- ui_color = '#a0e08c'
-
- @apply_defaults
- def __init__(self,
- cql,
- bucket,
- filename,
- schema_filename=None,
- approx_max_file_size_bytes=1900000000,
- cassandra_conn_id='cassandra_default',
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param cql: The CQL to execute on the Cassandra table.
- :type cql: string
- :param bucket: The bucket to upload to.
- :type bucket: string
- :param filename: The filename to use as the object name when uploading
- to Google cloud storage. A {} should be specified in the filename
- to allow the operator to inject file numbers in cases where the
- file is split due to size.
- :type filename: string
- :param schema_filename: If set, the filename to use as the object name
- when uploading a .json file containing the BigQuery schema fields
- for the table that was dumped from MySQL.
- :type schema_filename: string
- :param approx_max_file_size_bytes: This operator supports the ability
- to split large table dumps into multiple files (see notes in the
- filenamed param docs above). Google cloud storage allows for files
- to be a maximum of 4GB. This param allows developers to specify the
- file size of the splits.
- :type approx_max_file_size_bytes: long
- :param cassandra_conn_id: Reference to a specific Cassandra hook.
- :type cassandra_conn_id: string
- :param google_cloud_storage_conn_id: Reference to a specific Google
- cloud storage hook.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any. For this to
- work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
- super(CassandraToGoogleCloudStorageOperator, self).__init__(*args, **kwargs)
- self.cql = cql
- self.bucket = bucket
- self.filename = filename
- self.schema_filename = schema_filename
- self.approx_max_file_size_bytes = approx_max_file_size_bytes
- self.cassandra_conn_id = cassandra_conn_id
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
-
- self.hook = None
-
- # Default Cassandra to BigQuery type mapping
- CQL_TYPE_MAP = {
- 'BytesType': 'BYTES',
- 'DecimalType': 'FLOAT',
- 'UUIDType': 'BYTES',
- 'BooleanType': 'BOOL',
- 'ByteType': 'INTEGER',
- 'AsciiType': 'STRING',
- 'FloatType': 'FLOAT',
- 'DoubleType': 'FLOAT',
- 'LongType': 'INTEGER',
- 'Int32Type': 'INTEGER',
- 'IntegerType': 'INTEGER',
- 'InetAddressType': 'STRING',
- 'CounterColumnType': 'INTEGER',
- 'DateType': 'TIMESTAMP',
- 'SimpleDateType': 'DATE',
- 'TimestampType': 'TIMESTAMP',
- 'TimeUUIDType': 'BYTES',
- 'ShortType': 'INTEGER',
- 'TimeType': 'TIME',
- 'DurationType': 'INTEGER',
- 'UTF8Type': 'STRING',
- 'VarcharType': 'STRING',
- }
-
- def execute(self, context):
- cursor = self._query_cassandra()
- files_to_upload = self._write_local_data_files(cursor)
-
- # If a schema is set, create a BQ schema JSON file.
- if self.schema_filename:
- files_to_upload.update(self._write_local_schema_file(cursor))
-
- # Flush all files before uploading
- for file_handle in files_to_upload.values():
- file_handle.flush()
-
- self._upload_to_gcs(files_to_upload)
-
- # Close all temp file handles.
- for file_handle in files_to_upload.values():
- file_handle.close()
-
- # Close all sessions and connection associated with this Cassandra cluster
- self.hook.shutdown_cluster()
-
- def _query_cassandra(self):
- """
- Queries cassandra and returns a cursor to the results.
- """
- self.hook = CassandraHook(cassandra_conn_id=self.cassandra_conn_id)
- session = self.hook.get_conn()
- cursor = session.execute(self.cql)
- return cursor
-
- def _write_local_data_files(self, cursor):
- """
- Takes a cursor, and writes results to a local file.
-
- :return: A dictionary where keys are filenames to be used as object
- names in GCS, and values are file handles to local files that
- contain the data for the GCS objects.
- """
- file_no = 0
- tmp_file_handle = NamedTemporaryFile(delete=True)
- tmp_file_handles = {self.filename.format(file_no): tmp_file_handle}
- for row in cursor:
- row_dict = self.generate_data_dict(row._fields, row)
- s = json.dumps(row_dict)
- if PY3:
- s = s.encode('utf-8')
- tmp_file_handle.write(s)
-
- # Append newline to make dumps BigQuery compatible.
- tmp_file_handle.write(b'\n')
-
- if tmp_file_handle.tell() >= self.approx_max_file_size_bytes:
- file_no += 1
- tmp_file_handle = NamedTemporaryFile(delete=True)
- tmp_file_handles[self.filename.format(file_no)] = tmp_file_handle
-
- return tmp_file_handles
-
- def _write_local_schema_file(self, cursor):
- """
- Takes a cursor, and writes the BigQuery schema for the results to a
- local file system.
-
- :return: A dictionary where key is a filename to be used as an object
- name in GCS, and values are file handles to local files that
- contains the BigQuery schema fields in .json format.
- """
- schema = []
- tmp_schema_file_handle = NamedTemporaryFile(delete=True)
-
- for name, type in zip(cursor.column_names, cursor.column_types):
- schema.append(self.generate_schema_dict(name, type))
- json_serialized_schema = json.dumps(schema)
- if PY3:
- json_serialized_schema = json_serialized_schema.encode('utf-8')
-
- tmp_schema_file_handle.write(json_serialized_schema)
- return {self.schema_filename: tmp_schema_file_handle}
-
- def _upload_to_gcs(self, files_to_upload):
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- for object, tmp_file_handle in files_to_upload.items():
- hook.upload(self.bucket, object, tmp_file_handle.name, 'application/json')
-
- @classmethod
- def generate_data_dict(cls, names, values):
- row_dict = {}
- for name, value in zip(names, values):
- row_dict.update({name: cls.convert_value(name, value)})
- return row_dict
-
- @classmethod
- def convert_value(cls, name, value):
- if not value:
- return value
- elif isinstance(value, (text_type, int, float, bool, dict)):
- return value
- elif isinstance(value, binary_type):
- return b64encode(value).decode('ascii')
- elif isinstance(value, UUID):
- return b64encode(value.bytes).decode('ascii')
- elif isinstance(value, (datetime, Date)):
- return str(value)
- elif isinstance(value, Decimal):
- return float(value)
- elif isinstance(value, Time):
- return str(value).split('.')[0]
- elif isinstance(value, (list, SortedSet)):
- return cls.convert_array_types(name, value)
- elif hasattr(value, '_fields'):
- return cls.convert_user_type(name, value)
- elif isinstance(value, tuple):
- return cls.convert_tuple_type(name, value)
- elif isinstance(value, OrderedMapSerializedKey):
- return cls.convert_map_type(name, value)
- else:
- raise AirflowException('unexpected value: ' + str(value))
-
- @classmethod
- def convert_array_types(cls, name, value):
- return [cls.convert_value(name, nested_value) for nested_value in value]
-
- @classmethod
- def convert_user_type(cls, name, value):
- """
- Converts a user type to RECORD that contains n fields, where n is the
- number of attributes. Each element in the user type class will be converted to its
- corresponding data type in BQ.
- """
- names = value._fields
- values = [cls.convert_value(name, getattr(value, name)) for name in names]
- return cls.generate_data_dict(names, values)
-
- @classmethod
- def convert_tuple_type(cls, name, value):
- """
- Converts a tuple to RECORD that contains n fields, each will be converted
- to its corresponding data type in bq and will be named 'field_<index>', where
- index is determined by the order of the tuple elments defined in cassandra.
- """
- names = ['field_' + str(i) for i in range(len(value))]
- values = [cls.convert_value(name, value) for name, value in zip(names, value)]
- return cls.generate_data_dict(names, values)
-
- @classmethod
- def convert_map_type(cls, name, value):
- """
- Converts a map to a repeated RECORD that contains two fields: 'key' and 'value',
- each will be converted to its corresopnding data type in BQ.
- """
- converted_map = []
- for k, v in zip(value.keys(), value.values()):
- converted_map.append({
- 'key': cls.convert_value('key', k),
- 'value': cls.convert_value('value', v)
- })
- return converted_map
-
- @classmethod
- def generate_schema_dict(cls, name, type):
- field_schema = dict()
- field_schema.update({'name': name})
- field_schema.update({'type': cls.get_bq_type(type)})
- field_schema.update({'mode': cls.get_bq_mode(type)})
- fields = cls.get_bq_fields(name, type)
- if fields:
- field_schema.update({'fields': fields})
- return field_schema
-
- @classmethod
- def get_bq_fields(cls, name, type):
- fields = []
-
- if not cls.is_simple_type(type):
- names, types = [], []
-
- if cls.is_array_type(type) and cls.is_record_type(type.subtypes[0]):
- names = type.subtypes[0].fieldnames
- types = type.subtypes[0].subtypes
- elif cls.is_record_type(type):
- names = type.fieldnames
- types = type.subtypes
-
- if types and not names and type.cassname == 'TupleType':
- names = ['field_' + str(i) for i in range(len(types))]
- elif types and not names and type.cassname == 'MapType':
- names = ['key', 'value']
-
- for name, type in zip(names, types):
- field = cls.generate_schema_dict(name, type)
- fields.append(field)
-
- return fields
-
- @classmethod
- def is_simple_type(cls, type):
- return type.cassname in CassandraToGoogleCloudStorageOperator.CQL_TYPE_MAP
-
- @classmethod
- def is_array_type(cls, type):
- return type.cassname in ['ListType', 'SetType']
-
- @classmethod
- def is_record_type(cls, type):
- return type.cassname in ['UserType', 'TupleType', 'MapType']
-
- @classmethod
- def get_bq_type(cls, type):
- if cls.is_simple_type(type):
- return CassandraToGoogleCloudStorageOperator.CQL_TYPE_MAP[type.cassname]
- elif cls.is_record_type(type):
- return 'RECORD'
- elif cls.is_array_type(type):
- return cls.get_bq_type(type.subtypes[0])
- else:
- raise AirflowException('Not a supported type: ' + type.cassname)
-
- @classmethod
- def get_bq_mode(cls, type):
- if cls.is_array_type(type) or type.cassname == 'MapType':
- return 'REPEATED'
- elif cls.is_record_type(type) or cls.is_simple_type(type):
- return 'NULLABLE'
- else:
- raise AirflowException('Not a supported type: ' + type.cassname)
diff --git a/airflow/contrib/operators/databricks_operator.py b/airflow/contrib/operators/databricks_operator.py
deleted file mode 100644
index 7b8d522dba..0000000000
--- a/airflow/contrib/operators/databricks_operator.py
+++ /dev/null
@@ -1,271 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 six
-import time
-
-from airflow.exceptions import AirflowException
-from airflow.contrib.hooks.databricks_hook import DatabricksHook
-from airflow.models import BaseOperator
-
-
-XCOM_RUN_ID_KEY = 'run_id'
-XCOM_RUN_PAGE_URL_KEY = 'run_page_url'
-
-
-class DatabricksSubmitRunOperator(BaseOperator):
- """
- Submits an Spark job run to Databricks using the
- `api/2.0/jobs/runs/submit
- <https://docs.databricks.com/api/latest/jobs.html#runs-submit>`_
- API endpoint.
-
- There are two ways to instantiate this operator.
-
- In the first way, you can take the JSON payload that you typically use
- to call the ``api/2.0/jobs/runs/submit`` endpoint and pass it directly
- to our ``DatabricksSubmitRunOperator`` through the ``json`` parameter.
- For example ::
- json = {
- 'new_cluster': {
- 'spark_version': '2.1.0-db3-scala2.11',
- 'num_workers': 2
- },
- 'notebook_task': {
- 'notebook_path': '/Users/airflow@example.com/PrepareData',
- },
- }
- notebook_run = DatabricksSubmitRunOperator(task_id='notebook_run', json=json)
-
- Another way to accomplish the same thing is to use the named parameters
- of the ``DatabricksSubmitRunOperator`` directly. Note that there is exactly
- one named parameter for each top level parameter in the ``runs/submit``
- endpoint. In this method, your code would look like this: ::
- new_cluster = {
- 'spark_version': '2.1.0-db3-scala2.11',
- 'num_workers': 2
- }
- notebook_task = {
- 'notebook_path': '/Users/airflow@example.com/PrepareData',
- }
- notebook_run = DatabricksSubmitRunOperator(
- task_id='notebook_run',
- new_cluster=new_cluster,
- notebook_task=notebook_task)
-
- In the case where both the json parameter **AND** the named parameters
- are provided, they will be merged together. If there are conflicts during the merge,
- the named parameters will take precedence and override the top level ``json`` keys.
-
- Currently the named parameters that ``DatabricksSubmitRunOperator`` supports are
- - ``spark_jar_task``
- - ``notebook_task``
- - ``new_cluster``
- - ``existing_cluster_id``
- - ``libraries``
- - ``run_name``
- - ``timeout_seconds``
-
- :param json: A JSON object containing API parameters which will be passed
- directly to the ``api/2.0/jobs/runs/submit`` endpoint. The other named parameters
- (i.e. ``spark_jar_task``, ``notebook_task``..) to this operator will
- be merged with this json dictionary if they are provided.
- If there are conflicts during the merge, the named parameters will
- take precedence and override the top level json keys. (templated)
-
- .. seealso::
- For more information about templating see :ref:`jinja-templating`.
- https://docs.databricks.com/api/latest/jobs.html#runs-submit
- :type json: dict
- :param spark_jar_task: The main class and parameters for the JAR task. Note that
- the actual JAR is specified in the ``libraries``.
- *EITHER* ``spark_jar_task`` *OR* ``notebook_task`` should be specified.
- This field will be templated.
-
- .. seealso::
- https://docs.databricks.com/api/latest/jobs.html#jobssparkjartask
- :type spark_jar_task: dict
- :param notebook_task: The notebook path and parameters for the notebook task.
- *EITHER* ``spark_jar_task`` *OR* ``notebook_task`` should be specified.
- This field will be templated.
-
- .. seealso::
- https://docs.databricks.com/api/latest/jobs.html#jobsnotebooktask
- :type notebook_task: dict
- :param new_cluster: Specs for a new cluster on which this task will be run.
- *EITHER* ``new_cluster`` *OR* ``existing_cluster_id`` should be specified.
- This field will be templated.
-
- .. seealso::
- https://docs.databricks.com/api/latest/jobs.html#jobsclusterspecnewcluster
- :type new_cluster: dict
- :param existing_cluster_id: ID for existing cluster on which to run this task.
- *EITHER* ``new_cluster`` *OR* ``existing_cluster_id`` should be specified.
- This field will be templated.
- :type existing_cluster_id: string
- :param libraries: Libraries which this run will use.
- This field will be templated.
-
- .. seealso::
- https://docs.databricks.com/api/latest/libraries.html#managedlibrarieslibrary
- :type libraries: list of dicts
- :param run_name: The run name used for this task.
- By default this will be set to the Airflow ``task_id``. This ``task_id`` is a
- required parameter of the superclass ``BaseOperator``.
- This field will be templated.
- :type run_name: string
- :param timeout_seconds: The timeout for this run. By default a value of 0 is used
- which means to have no timeout.
- This field will be templated.
- :type timeout_seconds: int32
- :param databricks_conn_id: The name of the Airflow connection to use.
- By default and in the common case this will be ``databricks_default``. To use
- token based authentication, provide the key ``token`` in the extra field for the
- connection.
- :type databricks_conn_id: string
- :param polling_period_seconds: Controls the rate which we poll for the result of
- this run. By default the operator will poll every 30 seconds.
- :type polling_period_seconds: int
- :param databricks_retry_limit: Amount of times retry if the Databricks backend is
- unreachable. Its value must be greater than or equal to 1.
- :type databricks_retry_limit: int
- :param do_xcom_push: Whether we should push run_id and run_page_url to xcom.
- :type do_xcom_push: boolean
- """
- # Used in airflow.models.BaseOperator
- template_fields = ('json',)
- # Databricks brand color (blue) under white text
- ui_color = '#1CB1C2'
- ui_fgcolor = '#fff'
-
- def __init__(
- self,
- json=None,
- spark_jar_task=None,
- notebook_task=None,
- new_cluster=None,
- existing_cluster_id=None,
- libraries=None,
- run_name=None,
- timeout_seconds=None,
- databricks_conn_id='databricks_default',
- polling_period_seconds=30,
- databricks_retry_limit=3,
- do_xcom_push=False,
- **kwargs):
- """
- Creates a new ``DatabricksSubmitRunOperator``.
- """
- super(DatabricksSubmitRunOperator, self).__init__(**kwargs)
- self.json = json or {}
- self.databricks_conn_id = databricks_conn_id
- self.polling_period_seconds = polling_period_seconds
- self.databricks_retry_limit = databricks_retry_limit
- if spark_jar_task is not None:
- self.json['spark_jar_task'] = spark_jar_task
- if notebook_task is not None:
- self.json['notebook_task'] = notebook_task
- if new_cluster is not None:
- self.json['new_cluster'] = new_cluster
- if existing_cluster_id is not None:
- self.json['existing_cluster_id'] = existing_cluster_id
- if libraries is not None:
- self.json['libraries'] = libraries
- if run_name is not None:
- self.json['run_name'] = run_name
- if timeout_seconds is not None:
- self.json['timeout_seconds'] = timeout_seconds
- if 'run_name' not in self.json:
- self.json['run_name'] = run_name or kwargs['task_id']
-
- self.json = self._deep_string_coerce(self.json)
- # This variable will be used in case our task gets killed.
- self.run_id = None
- self.do_xcom_push = do_xcom_push
-
- def _deep_string_coerce(self, content, json_path='json'):
- """
- Coerces content or all values of content if it is a dict to a string. The
- function will throw if content contains non-string or non-numeric types.
-
- The reason why we have this function is because the ``self.json`` field must be a
- dict with only string values. This is because ``render_template`` will fail
- for numerical values.
- """
- c = self._deep_string_coerce
- if isinstance(content, six.string_types):
- return content
- elif isinstance(content, six.integer_types + (float,)):
- # Databricks can tolerate either numeric or string types in the API backend.
- return str(content)
- elif isinstance(content, (list, tuple)):
- return [c(e, '{0}[{1}]'.format(json_path, i)) for i, e in enumerate(content)]
- elif isinstance(content, dict):
- return {k: c(v, '{0}[{1}]'.format(json_path, k))
- for k, v in list(content.items())}
- else:
- param_type = type(content)
- msg = 'Type {0} used for parameter {1} is not a number or a string'\
- .format(param_type, json_path)
- raise AirflowException(msg)
-
- def _log_run_page_url(self, url):
- self.log.info('View run status, Spark UI, and logs at %s', url)
-
- def get_hook(self):
- return DatabricksHook(
- self.databricks_conn_id,
- retry_limit=self.databricks_retry_limit)
-
- def execute(self, context):
- hook = self.get_hook()
- self.run_id = hook.submit_run(self.json)
- if self.do_xcom_push:
- context['ti'].xcom_push(key=XCOM_RUN_ID_KEY, value=self.run_id)
- self.log.info('Run submitted with run_id: %s', self.run_id)
- run_page_url = hook.get_run_page_url(self.run_id)
- if self.do_xcom_push:
- context['ti'].xcom_push(key=XCOM_RUN_PAGE_URL_KEY, value=run_page_url)
- self._log_run_page_url(run_page_url)
- while True:
- run_state = hook.get_run_state(self.run_id)
- if run_state.is_terminal:
- if run_state.is_successful:
- self.log.info('%s completed successfully.', self.task_id)
- self._log_run_page_url(run_page_url)
- return
- else:
- error_message = '{t} failed with terminal state: {s}'.format(
- t=self.task_id,
- s=run_state)
- raise AirflowException(error_message)
- else:
- self.log.info('%s in run state: %s', self.task_id, run_state)
- self._log_run_page_url(run_page_url)
- self.log.info('Sleeping for %s seconds.', self.polling_period_seconds)
- time.sleep(self.polling_period_seconds)
-
- def on_kill(self):
- hook = self.get_hook()
- hook.cancel_run(self.run_id)
- self.log.info(
- 'Task: %s with run_id: %s was requested to be cancelled.',
- self.task_id, self.run_id
- )
diff --git a/airflow/contrib/operators/dataflow_operator.py b/airflow/contrib/operators/dataflow_operator.py
deleted file mode 100644
index e3c8c1fff1..0000000000
--- a/airflow/contrib/operators/dataflow_operator.py
+++ /dev/null
@@ -1,376 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 re
-import uuid
-import copy
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.contrib.hooks.gcp_dataflow_hook import DataFlowHook
-from airflow.models import BaseOperator
-from airflow.version import version
-from airflow.utils.decorators import apply_defaults
-
-
-class DataFlowJavaOperator(BaseOperator):
- """
- Start a Java Cloud DataFlow batch job. The parameters of the operation
- will be passed to the job.
-
- It's a good practice to define dataflow_* parameters in the default_args of the dag
- like the project, zone and staging location.
-
- .. code-block:: python
-
- default_args = {
- 'dataflow_default_options': {
- 'project': 'my-gcp-project',
- 'zone': 'europe-west1-d',
- 'stagingLocation': 'gs://my-staging-bucket/staging/'
- }
- }
-
- You need to pass the path to your dataflow as a file reference with the ``jar``
- parameter, the jar needs to be a self executing jar (see documentation here:
- https://beam.apache.org/documentation/runners/dataflow/#self-executing-jar).
- Use ``options`` to pass on options to your job.
-
- .. code-block:: python
-
- t1 = DataFlowOperation(
- task_id='datapflow_example',
- jar='{{var.value.gcp_dataflow_base}}pipeline/build/libs/pipeline-example-1.0.jar',
- options={
- 'autoscalingAlgorithm': 'BASIC',
- 'maxNumWorkers': '50',
- 'start': '{{ds}}',
- 'partitionType': 'DAY',
- 'labels': {'foo' : 'bar'}
- },
- gcp_conn_id='gcp-airflow-service-account',
- dag=my-dag)
-
- Both ``jar`` and ``options`` are templated so you can use variables in them.
- """
- template_fields = ['options', 'jar']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- jar,
- dataflow_default_options=None,
- options=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- poll_sleep=10,
- job_class=None,
- *args,
- **kwargs):
- """
- Create a new DataFlowJavaOperator. Note that both
- dataflow_default_options and options will be merged to specify pipeline
- execution parameter, and dataflow_default_options is expected to save
- high-level options, for instances, project and zone information, which
- apply to all dataflow operators in the DAG.
-
-
- .. seealso::
- For more detail on job submission have a look at the reference:
- https://cloud.google.com/dataflow/pipelines/specifying-exec-params
-
- :param jar: The reference to a self executing DataFlow jar.
- :type jar: string
- :param dataflow_default_options: Map of default job options.
- :type dataflow_default_options: dict
- :param options: Map of job specific options.
- :type options: dict
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud
- Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- :param poll_sleep: The time in seconds to sleep between polling Google
- Cloud Platform for the dataflow job status while the job is in the
- JOB_STATE_RUNNING state.
- :type poll_sleep: int
- :param job_class: The name of the dataflow job class to be executued, it
- is often not the main class configured in the dataflow jar file.
- :type job_class: string
- """
- super(DataFlowJavaOperator, self).__init__(*args, **kwargs)
-
- dataflow_default_options = dataflow_default_options or {}
- options = options or {}
- options.setdefault('labels', {}).update(
- {'airflow-version': 'v' + version.replace('.', '-').replace('+', '-')})
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.jar = jar
- self.dataflow_default_options = dataflow_default_options
- self.options = options
- self.poll_sleep = poll_sleep
- self.job_class = job_class
-
- def execute(self, context):
- bucket_helper = GoogleCloudBucketHelper(
- self.gcp_conn_id, self.delegate_to)
- self.jar = bucket_helper.google_cloud_to_local(self.jar)
- hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- poll_sleep=self.poll_sleep)
-
- dataflow_options = copy.copy(self.dataflow_default_options)
- dataflow_options.update(self.options)
-
- hook.start_java_dataflow(self.task_id, dataflow_options,
- self.jar, self.job_class)
-
-
-class DataflowTemplateOperator(BaseOperator):
- """
- Start a Templated Cloud DataFlow batch job. The parameters of the operation
- will be passed to the job.
- It's a good practice to define dataflow_* parameters in the default_args of the dag
- like the project, zone and staging location.
-
- .. seealso::
- https://cloud.google.com/dataflow/docs/reference/rest/v1b3/LaunchTemplateParameters
- https://cloud.google.com/dataflow/docs/reference/rest/v1b3/RuntimeEnvironment
-
- .. code-block:: python
-
- default_args = {
- 'dataflow_default_options': {
- 'project': 'my-gcp-project'
- 'zone': 'europe-west1-d',
- 'tempLocation': 'gs://my-staging-bucket/staging/'
- }
- }
- }
-
- You need to pass the path to your dataflow template as a file reference with the
- ``template`` parameter. Use ``parameters`` to pass on parameters to your job.
- Use ``environment`` to pass on runtime environment variables to your job.
-
- .. code-block:: python
-
- t1 = DataflowTemplateOperator(
- task_id='datapflow_example',
- template='{{var.value.gcp_dataflow_base}}',
- parameters={
- 'inputFile': "gs://bucket/input/my_input.txt",
- 'outputFile': "gs://bucket/output/my_output.txt"
- },
- gcp_conn_id='gcp-airflow-service-account',
- dag=my-dag)
-
- ``template``, ``dataflow_default_options`` and ``parameters`` are templated so you can
- use variables in them.
- """
- template_fields = ['parameters', 'dataflow_default_options', 'template']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- template,
- dataflow_default_options=None,
- parameters=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- poll_sleep=10,
- *args,
- **kwargs):
- """
- Create a new DataflowTemplateOperator. Note that
- dataflow_default_options is expected to save high-level options
- for project information, which apply to all dataflow operators in the DAG.
-
- .. seealso::
- https://cloud.google.com/dataflow/docs/reference/rest/v1b3
- /LaunchTemplateParameters
- https://cloud.google.com/dataflow/docs/reference/rest/v1b3/RuntimeEnvironment
- For more detail on job template execution have a look at the reference:
- https://cloud.google.com/dataflow/docs/templates/executing-templates
-
- :param template: The reference to the DataFlow template.
- :type template: string
- :param dataflow_default_options: Map of default job environment options.
- :type dataflow_default_options: dict
- :param parameters: Map of job specific parameters for the template.
- :type parameters: dict
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud
- Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- :param poll_sleep: The time in seconds to sleep between polling Google
- Cloud Platform for the dataflow job status while the job is in the
- JOB_STATE_RUNNING state.
- :type poll_sleep: int
- """
- super(DataflowTemplateOperator, self).__init__(*args, **kwargs)
-
- dataflow_default_options = dataflow_default_options or {}
- parameters = parameters or {}
-
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.dataflow_default_options = dataflow_default_options
- self.poll_sleep = poll_sleep
- self.template = template
- self.parameters = parameters
-
- def execute(self, context):
- hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- poll_sleep=self.poll_sleep)
-
- hook.start_template_dataflow(self.task_id, self.dataflow_default_options,
- self.parameters, self.template)
-
-
-class DataFlowPythonOperator(BaseOperator):
-
- template_fields = ['options', 'dataflow_default_options']
-
- @apply_defaults
- def __init__(
- self,
- py_file,
- py_options=None,
- dataflow_default_options=None,
- options=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- poll_sleep=10,
- *args,
- **kwargs):
- """
- Create a new DataFlowPythonOperator. Note that both
- dataflow_default_options and options will be merged to specify pipeline
- execution parameter, and dataflow_default_options is expected to save
- high-level options, for instances, project and zone information, which
- apply to all dataflow operators in the DAG.
-
- .. seealso::
- For more detail on job submission have a look at the reference:
- https://cloud.google.com/dataflow/pipelines/specifying-exec-params
-
- :param py_file: Reference to the python dataflow pipleline file.py, e.g.,
- /some/local/file/path/to/your/python/pipeline/file.
- :type py_file: string
- :param py_options: Additional python options.
- :type pyt_options: list of strings, e.g., ["-m", "-v"].
- :param dataflow_default_options: Map of default job options.
- :type dataflow_default_options: dict
- :param options: Map of job specific options.
- :type options: dict
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud
- Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- :param poll_sleep: The time in seconds to sleep between polling Google
- Cloud Platform for the dataflow job status while the job is in the
- JOB_STATE_RUNNING state.
- :type poll_sleep: int
- """
- super(DataFlowPythonOperator, self).__init__(*args, **kwargs)
-
- self.py_file = py_file
- self.py_options = py_options or []
- self.dataflow_default_options = dataflow_default_options or {}
- self.options = options or {}
- self.options.setdefault('labels', {}).update(
- {'airflow-version': 'v' + version.replace('.', '-').replace('+', '-')})
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.poll_sleep = poll_sleep
-
- def execute(self, context):
- """Execute the python dataflow job."""
- bucket_helper = GoogleCloudBucketHelper(
- self.gcp_conn_id, self.delegate_to)
- self.py_file = bucket_helper.google_cloud_to_local(self.py_file)
- hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- poll_sleep=self.poll_sleep)
- dataflow_options = self.dataflow_default_options.copy()
- dataflow_options.update(self.options)
- # Convert argument names from lowerCamelCase to snake case.
- camel_to_snake = lambda name: re.sub(
- r'[A-Z]', lambda x: '_' + x.group(0).lower(), name)
- formatted_options = {camel_to_snake(key): dataflow_options[key]
- for key in dataflow_options}
- hook.start_python_dataflow(
- self.task_id, formatted_options,
- self.py_file, self.py_options)
-
-
-class GoogleCloudBucketHelper():
- """GoogleCloudStorageHook helper class to download GCS object."""
- GCS_PREFIX_LENGTH = 5
-
- def __init__(self,
- gcp_conn_id='google_cloud_default',
- delegate_to=None):
- self._gcs_hook = GoogleCloudStorageHook(gcp_conn_id, delegate_to)
-
- def google_cloud_to_local(self, file_name):
- """
- Checks whether the file specified by file_name is stored in Google Cloud
- Storage (GCS), if so, downloads the file and saves it locally. The full
- path of the saved file will be returned. Otherwise the local file_name
- will be returned immediately.
-
- :param file_name: The full path of input file.
- :type file_name: string
- :return: The full path of local file.
- :type: string
- """
- if not file_name.startswith('gs://'):
- return file_name
-
- # Extracts bucket_id and object_id by first removing 'gs://' prefix and
- # then split the remaining by path delimiter '/'.
- path_components = file_name[self.GCS_PREFIX_LENGTH:].split('/')
- if path_components < 2:
- raise Exception(
- 'Invalid Google Cloud Storage (GCS) object path: {}.'
- .format(file_name))
-
- bucket_id = path_components[0]
- object_id = '/'.join(path_components[1:])
- local_file = '/tmp/dataflow{}-{}'.format(str(uuid.uuid1())[:8],
- path_components[-1])
- file_size = self._gcs_hook.download(bucket_id, object_id, local_file)
-
- if file_size > 0:
- return local_file
- raise Exception(
- 'Failed to download Google Cloud Storage GCS object: {}'
- .format(file_name))
diff --git a/airflow/contrib/operators/dataproc_operator.py b/airflow/contrib/operators/dataproc_operator.py
deleted file mode 100644
index 5d59f7fb6e..0000000000
--- a/airflow/contrib/operators/dataproc_operator.py
+++ /dev/null
@@ -1,1309 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 ntpath
-import os
-import re
-import time
-import uuid
-from datetime import timedelta
-
-from airflow.contrib.hooks.gcp_dataproc_hook import DataProcHook
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.version import version
-from googleapiclient.errors import HttpError
-from airflow.utils import timezone
-
-
-class DataprocClusterCreateOperator(BaseOperator):
- """
- Create a new cluster on Google Cloud Dataproc. The operator will wait until the
- creation is successful or an error occurs in the creation process.
-
- The parameters allow to configure the cluster. Please refer to
-
- https://cloud.google.com/dataproc/docs/reference/rest/v1/projects.regions.clusters
-
- for a detailed explanation on the different parameters. Most of the configuration
- parameters detailed in the link are available as a parameter to this operator.
-
- :param cluster_name: The name of the DataProc cluster to create. (templated)
- :type cluster_name: string
- :param project_id: The ID of the google cloud project in which
- to create the cluster. (templated)
- :type project_id: string
- :param num_workers: The # of workers to spin up
- :type num_workers: int
- :param storage_bucket: The storage bucket to use, setting to None lets dataproc
- generate a custom one for you
- :type storage_bucket: string
- :param init_actions_uris: List of GCS uri's containing
- dataproc initialization scripts
- :type init_actions_uris: list[string]
- :param init_action_timeout: Amount of time executable scripts in
- init_actions_uris has to complete
- :type init_action_timeout: string
- :param metadata: dict of key-value google compute engine metadata entries
- to add to all instances
- :type metadata: dict
- :param image_version: the version of software inside the Dataproc cluster
- :type image_version: string
- :param properties: dict of properties to set on
- config files (e.g. spark-defaults.conf), see
- https://cloud.google.com/dataproc/docs/reference/rest/v1/ \
- projects.regions.clusters#SoftwareConfig
- :type properties: dict
- :param master_machine_type: Compute engine machine type to use for the master node
- :type master_machine_type: string
- :param master_disk_size: Disk size for the master node
- :type master_disk_size: int
- :param worker_machine_type: Compute engine machine type to use for the worker nodes
- :type worker_machine_type: string
- :param worker_disk_size: Disk size for the worker nodes
- :type worker_disk_size: int
- :param num_preemptible_workers: The # of preemptible worker nodes to spin up
- :type num_preemptible_workers: int
- :param labels: dict of labels to add to the cluster
- :type labels: dict
- :param zone: The zone where the cluster will be located. (templated)
- :type zone: string
- :param network_uri: The network uri to be used for machine communication, cannot be
- specified with subnetwork_uri
- :type network_uri: string
- :param subnetwork_uri: The subnetwork uri to be used for machine communication,
- cannot be specified with network_uri
- :type subnetwork_uri: string
- :param internal_ip_only: If true, all instances in the cluster will only
- have internal IP addresses. This can only be enabled for subnetwork
- enabled networks
- :type internal_ip_only: bool
- :param tags: The GCE tags to add to all instances
- :type tags: list[string]
- :param region: leave as 'global', might become relevant in the future. (templated)
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param service_account: The service account of the dataproc instances.
- :type service_account: string
- :param service_account_scopes: The URIs of service account scopes to be included.
- :type service_account_scopes: list[string]
- :param idle_delete_ttl: The longest duration that cluster would keep alive while
- staying idle. Passing this threshold will cause cluster to be auto-deleted.
- A duration in seconds.
- :type idle_delete_ttl: int
- :param auto_delete_time: The time when cluster will be auto-deleted.
- :type auto_delete_time: datetime.datetime
- :param auto_delete_ttl: The life duration of cluster, the cluster will be
- auto-deleted at the end of this duration.
- A duration in seconds. (If auto_delete_time is set this parameter will be ignored)
- :type auto_delete_ttl: int
- """
-
- template_fields = ['cluster_name', 'project_id', 'zone', 'region']
-
- @apply_defaults
- def __init__(self,
- cluster_name,
- project_id,
- num_workers,
- zone,
- network_uri=None,
- subnetwork_uri=None,
- internal_ip_only=None,
- tags=None,
- storage_bucket=None,
- init_actions_uris=None,
- init_action_timeout="10m",
- metadata=None,
- image_version=None,
- properties=None,
- master_machine_type='n1-standard-4',
- master_disk_size=500,
- worker_machine_type='n1-standard-4',
- worker_disk_size=500,
- num_preemptible_workers=0,
- labels=None,
- region='global',
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- service_account=None,
- service_account_scopes=None,
- idle_delete_ttl=None,
- auto_delete_time=None,
- auto_delete_ttl=None,
- *args,
- **kwargs):
-
- super(DataprocClusterCreateOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.cluster_name = cluster_name
- self.project_id = project_id
- self.num_workers = num_workers
- self.num_preemptible_workers = num_preemptible_workers
- self.storage_bucket = storage_bucket
- self.init_actions_uris = init_actions_uris
- self.init_action_timeout = init_action_timeout
- self.metadata = metadata
- self.image_version = image_version
- self.properties = properties
- self.master_machine_type = master_machine_type
- self.master_disk_size = master_disk_size
- self.worker_machine_type = worker_machine_type
- self.worker_disk_size = worker_disk_size
- self.labels = labels
- self.zone = zone
- self.network_uri = network_uri
- self.subnetwork_uri = subnetwork_uri
- self.internal_ip_only = internal_ip_only
- self.tags = tags
- self.region = region
- self.service_account = service_account
- self.service_account_scopes = service_account_scopes
- self.idle_delete_ttl = idle_delete_ttl
- self.auto_delete_time = auto_delete_time
- self.auto_delete_ttl = auto_delete_ttl
-
- def _get_cluster_list_for_project(self, service):
- result = service.projects().regions().clusters().list(
- projectId=self.project_id,
- region=self.region
- ).execute()
- return result.get('clusters', [])
-
- def _get_cluster(self, service):
- cluster_list = self._get_cluster_list_for_project(service)
- cluster = [c for c in cluster_list if c['clusterName'] == self.cluster_name]
- if cluster:
- return cluster[0]
- return None
-
- def _get_cluster_state(self, service):
- cluster = self._get_cluster(service)
- if 'status' in cluster:
- return cluster['status']['state']
- else:
- return None
-
- def _cluster_ready(self, state, service):
- if state == 'RUNNING':
- return True
- if state == 'ERROR':
- cluster = self._get_cluster(service)
- try:
- error_details = cluster['status']['details']
- except KeyError:
- error_details = 'Unknown error in cluster creation, ' \
- 'check Google Cloud console for details.'
- raise Exception(error_details)
- return False
-
- def _wait_for_done(self, service):
- while True:
- state = self._get_cluster_state(service)
- if state is None:
- self.log.info("No state for cluster '%s'", self.cluster_name)
- time.sleep(15)
- else:
- self.log.info("State for cluster '%s' is %s", self.cluster_name, state)
- if self._cluster_ready(state, service):
- self.log.info(
- "Cluster '%s' successfully created", self.cluster_name
- )
- return
- time.sleep(15)
-
- def _get_init_action_timeout(self):
- match = re.match(r"^(\d+)(s|m)$", self.init_action_timeout)
- if match:
- if match.group(2) == "s":
- return self.init_action_timeout
- elif match.group(2) == "m":
- val = float(match.group(1))
- return "{}s".format(timedelta(minutes=val).seconds)
-
- raise AirflowException(
- "DataprocClusterCreateOperator init_action_timeout"
- " should be expressed in minutes or seconds. i.e. 10m, 30s")
-
- def _build_cluster_data(self):
- zone_uri = \
- 'https://www.googleapis.com/compute/v1/projects/{}/zones/{}'.format(
- self.project_id, self.zone
- )
- master_type_uri = \
- "https://www.googleapis.com/compute/v1/projects/{}/zones/{}/machineTypes/{}"\
- .format(self.project_id, self.zone, self.master_machine_type)
- worker_type_uri = \
- "https://www.googleapis.com/compute/v1/projects/{}/zones/{}/machineTypes/{}"\
- .format(self.project_id, self.zone, self.worker_machine_type)
-
- cluster_data = {
- 'projectId': self.project_id,
- 'clusterName': self.cluster_name,
- 'config': {
- 'gceClusterConfig': {
- 'zoneUri': zone_uri
- },
- 'masterConfig': {
- 'numInstances': 1,
- 'machineTypeUri': master_type_uri,
- 'diskConfig': {
- 'bootDiskSizeGb': self.master_disk_size
- }
- },
- 'workerConfig': {
- 'numInstances': self.num_workers,
- 'machineTypeUri': worker_type_uri,
- 'diskConfig': {
- 'bootDiskSizeGb': self.worker_disk_size
- }
- },
- 'secondaryWorkerConfig': {},
- 'softwareConfig': {},
- 'lifecycleConfig': {}
- }
- }
- if self.num_preemptible_workers > 0:
- cluster_data['config']['secondaryWorkerConfig'] = {
- 'numInstances': self.num_preemptible_workers,
- 'machineTypeUri': worker_type_uri,
- 'diskConfig': {
- 'bootDiskSizeGb': self.worker_disk_size
- },
- 'isPreemptible': True
- }
-
- cluster_data['labels'] = self.labels if self.labels else {}
- # Dataproc labels must conform to the following regex:
- # [a-z]([-a-z0-9]*[a-z0-9])? (current airflow version string follows
- # semantic versioning spec: x.y.z).
- cluster_data['labels'].update({'airflow-version':
- 'v' + version.replace('.', '-').replace('+', '-')})
- if self.storage_bucket:
- cluster_data['config']['configBucket'] = self.storage_bucket
- if self.metadata:
- cluster_data['config']['gceClusterConfig']['metadata'] = self.metadata
- if self.network_uri:
- cluster_data['config']['gceClusterConfig']['networkUri'] = self.network_uri
- if self.subnetwork_uri:
- cluster_data['config']['gceClusterConfig']['subnetworkUri'] = \
- self.subnetwork_uri
- if self.internal_ip_only:
- if not self.subnetwork_uri:
- raise AirflowException("Set internal_ip_only to true only when"
- " you pass a subnetwork_uri.")
- cluster_data['config']['gceClusterConfig']['internalIpOnly'] = True
- if self.tags:
- cluster_data['config']['gceClusterConfig']['tags'] = self.tags
- if self.image_version:
- cluster_data['config']['softwareConfig']['imageVersion'] = self.image_version
- if self.properties:
- cluster_data['config']['softwareConfig']['properties'] = self.properties
- if self.idle_delete_ttl:
- cluster_data['config']['lifecycleConfig']['idleDeleteTtl'] = \
- "{}s".format(self.idle_delete_ttl)
- if self.auto_delete_time:
- utc_auto_delete_time = timezone.convert_to_utc(self.auto_delete_time)
- cluster_data['config']['lifecycleConfig']['autoDeleteTime'] = \
- utc_auto_delete_time.format('%Y-%m-%dT%H:%M:%S.%fZ', formatter='classic')
- elif self.auto_delete_ttl:
- cluster_data['config']['lifecycleConfig']['autoDeleteTtl'] = \
- "{}s".format(self.auto_delete_ttl)
- if self.init_actions_uris:
- init_actions_dict = [
- {
- 'executableFile': uri,
- 'executionTimeout': self._get_init_action_timeout()
- } for uri in self.init_actions_uris
- ]
- cluster_data['config']['initializationActions'] = init_actions_dict
- if self.service_account:
- cluster_data['config']['gceClusterConfig']['serviceAccount'] =\
- self.service_account
- if self.service_account_scopes:
- cluster_data['config']['gceClusterConfig']['serviceAccountScopes'] =\
- self.service_account_scopes
- return cluster_data
-
- def execute(self, context):
- self.log.info('Creating cluster: %s', self.cluster_name)
- hook = DataProcHook(
- gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to
- )
- service = hook.get_conn()
-
- if self._get_cluster(service):
- self.log.info(
- 'Cluster %s already exists... Checking status...',
- self.cluster_name
- )
- self._wait_for_done(service)
- return True
-
- cluster_data = self._build_cluster_data()
- try:
- service.projects().regions().clusters().create(
- projectId=self.project_id,
- region=self.region,
- body=cluster_data
- ).execute()
- except HttpError as e:
- # probably two cluster start commands at the same time
- time.sleep(10)
- if self._get_cluster(service):
- self.log.info(
- 'Cluster {} already exists... Checking status...',
- self.cluster_name
- )
- self._wait_for_done(service)
- return True
- else:
- raise e
-
- self._wait_for_done(service)
-
-
-class DataprocClusterScaleOperator(BaseOperator):
- """
- Scale, up or down, a cluster on Google Cloud Dataproc.
- The operator will wait until the cluster is re-scaled.
-
- **Example**: ::
-
- t1 = DataprocClusterScaleOperator(
- task_id='dataproc_scale',
- project_id='my-project',
- cluster_name='cluster-1',
- num_workers=10,
- num_preemptible_workers=10,
- graceful_decommission_timeout='1h'
- dag=dag)
-
- .. seealso::
- For more detail on about scaling clusters have a look at the reference:
- https://cloud.google.com/dataproc/docs/concepts/configuring-clusters/scaling-clusters
-
- :param cluster_name: The name of the cluster to scale. (templated)
- :type cluster_name: string
- :param project_id: The ID of the google cloud project in which
- the cluster runs. (templated)
- :type project_id: string
- :param region: The region for the dataproc cluster. (templated)
- :type region: string
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param num_workers: The new number of workers
- :type num_workers: int
- :param num_preemptible_workers: The new number of preemptible workers
- :type num_preemptible_workers: int
- :param graceful_decommission_timeout: Timeout for graceful YARN decomissioning.
- Maximum value is 1d
- :type graceful_decommission_timeout: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
-
- template_fields = ['cluster_name', 'project_id', 'region']
-
- @apply_defaults
- def __init__(self,
- cluster_name,
- project_id,
- region='global',
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- num_workers=2,
- num_preemptible_workers=0,
- graceful_decommission_timeout=None,
- *args,
- **kwargs):
- super(DataprocClusterScaleOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.cluster_name = cluster_name
- self.project_id = project_id
- self.region = region
- self.num_workers = num_workers
- self.num_preemptible_workers = num_preemptible_workers
-
- # Optional
- self.optional_arguments = {}
- if graceful_decommission_timeout:
- self.optional_arguments['gracefulDecommissionTimeout'] = \
- self._get_graceful_decommission_timeout(
- graceful_decommission_timeout)
-
- def _wait_for_done(self, service, operation_name):
- time.sleep(15)
- while True:
- try:
- response = service.projects().regions().operations().get(
- name=operation_name
- ).execute()
-
- if 'done' in response and response['done']:
- if 'error' in response:
- raise Exception(str(response['error']))
- else:
- return
- time.sleep(15)
- except HttpError as e:
- self.log.error("Operation not found.")
- raise e
-
- def _build_scale_cluster_data(self):
- scale_data = {
- 'config': {
- 'workerConfig': {
- 'numInstances': self.num_workers
- },
- 'secondaryWorkerConfig': {
- 'numInstances': self.num_preemptible_workers
- }
- }
- }
- return scale_data
-
- def _get_graceful_decommission_timeout(self, timeout):
- match = re.match(r"^(\d+)(s|m|h|d)$", timeout)
- if match:
- if match.group(2) == "s":
- return timeout
- elif match.group(2) == "m":
- val = float(match.group(1))
- return "{}s".format(timedelta(minutes=val).seconds)
- elif match.group(2) == "h":
- val = float(match.group(1))
- return "{}s".format(timedelta(hours=val).seconds)
- elif match.group(2) == "d":
- val = float(match.group(1))
- return "{}s".format(timedelta(days=val).seconds)
-
- raise AirflowException(
- "DataprocClusterScaleOperator "
- " should be expressed in day, hours, minutes or seconds. "
- " i.e. 1d, 4h, 10m, 30s")
-
- def execute(self, context):
- self.log.info("Scaling cluster: %s", self.cluster_name)
- hook = DataProcHook(
- gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to
- )
- service = hook.get_conn()
-
- update_mask = "config.worker_config.num_instances," \
- + "config.secondary_worker_config.num_instances"
- scaling_cluster_data = self._build_scale_cluster_data()
-
- response = service.projects().regions().clusters().patch(
- projectId=self.project_id,
- region=self.region,
- clusterName=self.cluster_name,
- updateMask=update_mask,
- body=scaling_cluster_data,
- **self.optional_arguments
- ).execute()
- operation_name = response['name']
- self.log.info("Cluster scale operation name: %s", operation_name)
- self._wait_for_done(service, operation_name)
-
-
-class DataprocClusterDeleteOperator(BaseOperator):
- """
- Delete a cluster on Google Cloud Dataproc. The operator will wait until the
- cluster is destroyed.
-
- :param cluster_name: The name of the cluster to create. (templated)
- :type cluster_name: string
- :param project_id: The ID of the google cloud project in which
- the cluster runs. (templated)
- :type project_id: string
- :param region: leave as 'global', might become relevant in the future. (templated)
- :type region: string
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
-
- template_fields = ['cluster_name', 'project_id', 'region']
-
- @apply_defaults
- def __init__(self,
- cluster_name,
- project_id,
- region='global',
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
-
- super(DataprocClusterDeleteOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.cluster_name = cluster_name
- self.project_id = project_id
- self.region = region
-
- def _wait_for_done(self, service, operation_name):
- time.sleep(15)
- while True:
- response = service.projects().regions().operations().get(
- name=operation_name
- ).execute()
-
- if 'done' in response and response['done']:
- if 'error' in response:
- raise Exception(str(response['error']))
- else:
- return
- time.sleep(15)
-
- def execute(self, context):
- self.log.info('Deleting cluster: %s', self.cluster_name)
- hook = DataProcHook(
- gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to
- )
- service = hook.get_conn()
-
- response = service.projects().regions().clusters().delete(
- projectId=self.project_id,
- region=self.region,
- clusterName=self.cluster_name
- ).execute()
- operation_name = response['name']
- self.log.info("Cluster delete operation name: %s", operation_name)
- self._wait_for_done(service, operation_name)
-
-
-class DataProcPigOperator(BaseOperator):
- """
- Start a Pig query Job on a Cloud DataProc cluster. The parameters of the operation
- will be passed to the cluster.
-
- It's a good practice to define dataproc_* parameters in the default_args of the dag
- like the cluster name and UDFs.
-
- .. code-block:: python
-
- default_args = {
- 'cluster_name': 'cluster-1',
- 'dataproc_pig_jars': [
- 'gs://example/udf/jar/datafu/1.2.0/datafu.jar',
- 'gs://example/udf/jar/gpig/1.2/gpig.jar'
- ]
- }
-
- You can pass a pig script as string or file reference. Use variables to pass on
- variables for the pig script to be resolved on the cluster or use the parameters to
- be resolved in the script as template parameters.
-
- **Example**: ::
-
- t1 = DataProcPigOperator(
- task_id='dataproc_pig',
- query='a_pig_script.pig',
- variables={'out': 'gs://example/output/{{ds}}'},
- dag=dag)
-
- .. seealso::
- For more detail on about job submission have a look at the reference:
- https://cloud.google.com/dataproc/reference/rest/v1/projects.regions.jobs
-
- :param query: The query or reference to the query
- file (pg or pig extension). (templated)
- :type query: string
- :param query_uri: The uri of a pig script on Cloud Storage.
- :type query_uri: string
- :param variables: Map of named parameters for the query. (templated)
- :type variables: dict
- :param job_name: The job name used in the DataProc cluster. This
- name by default is the task_id appended with the execution data, but can
- be templated. The name will always be appended with a random number to
- avoid name clashes. (templated)
- :type job_name: string
- :param cluster_name: The name of the DataProc cluster. (templated)
- :type cluster_name: string
- :param dataproc_pig_properties: Map for the Pig properties. Ideal to put in
- default arguments
- :type dataproc_pig_properties: dict
- :param dataproc_pig_jars: URIs to jars provisioned in Cloud Storage (example: for
- UDFs and libs) and are ideal to put in default arguments.
- :type dataproc_pig_jars: list
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param region: The specified region where the dataproc cluster is created.
- :type region: string
- """
- template_fields = ['query', 'variables', 'job_name', 'cluster_name', 'dataproc_jars']
- template_ext = ('.pg', '.pig',)
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- query=None,
- query_uri=None,
- variables=None,
- job_name='{{task.task_id}}_{{ds_nodash}}',
- cluster_name='cluster-1',
- dataproc_pig_properties=None,
- dataproc_pig_jars=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- region='global',
- *args,
- **kwargs):
-
- super(DataProcPigOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.query = query
- self.query_uri = query_uri
- self.variables = variables
- self.job_name = job_name
- self.cluster_name = cluster_name
- self.dataproc_properties = dataproc_pig_properties
- self.dataproc_jars = dataproc_pig_jars
- self.region = region
-
- def execute(self, context):
- hook = DataProcHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
- job = hook.create_job_template(self.task_id, self.cluster_name, "pigJob",
- self.dataproc_properties)
-
- if self.query is None:
- job.add_query_uri(self.query_uri)
- else:
- job.add_query(self.query)
- job.add_variables(self.variables)
- job.add_jar_file_uris(self.dataproc_jars)
- job.set_job_name(self.job_name)
-
- hook.submit(hook.project_id, job.build(), self.region)
-
-
-class DataProcHiveOperator(BaseOperator):
- """
- Start a Hive query Job on a Cloud DataProc cluster.
-
- :param query: The query or reference to the query file (q extension).
- :type query: string
- :param query_uri: The uri of a hive script on Cloud Storage.
- :type query_uri: string
- :param variables: Map of named parameters for the query.
- :type variables: dict
- :param job_name: The job name used in the DataProc cluster. This name by default
- is the task_id appended with the execution data, but can be templated. The
- name will always be appended with a random number to avoid name clashes.
- :type job_name: string
- :param cluster_name: The name of the DataProc cluster.
- :type cluster_name: string
- :param dataproc_hive_properties: Map for the Pig properties. Ideal to put in
- default arguments
- :type dataproc_hive_properties: dict
- :param dataproc_hive_jars: URIs to jars provisioned in Cloud Storage (example: for
- UDFs and libs) and are ideal to put in default arguments.
- :type dataproc_hive_jars: list
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param region: The specified region where the dataproc cluster is created.
- :type region: string
- """
- template_fields = ['query', 'variables', 'job_name', 'cluster_name', 'dataproc_jars']
- template_ext = ('.q',)
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- query=None,
- query_uri=None,
- variables=None,
- job_name='{{task.task_id}}_{{ds_nodash}}',
- cluster_name='cluster-1',
- dataproc_hive_properties=None,
- dataproc_hive_jars=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- region='global',
- *args,
- **kwargs):
-
- super(DataProcHiveOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.query = query
- self.query_uri = query_uri
- self.variables = variables
- self.job_name = job_name
- self.cluster_name = cluster_name
- self.dataproc_properties = dataproc_hive_properties
- self.dataproc_jars = dataproc_hive_jars
- self.region = region
-
- def execute(self, context):
- hook = DataProcHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- job = hook.create_job_template(self.task_id, self.cluster_name, "hiveJob",
- self.dataproc_properties)
-
- if self.query is None:
- job.add_query_uri(self.query_uri)
- else:
- job.add_query(self.query)
- job.add_variables(self.variables)
- job.add_jar_file_uris(self.dataproc_jars)
- job.set_job_name(self.job_name)
-
- hook.submit(hook.project_id, job.build(), self.region)
-
-
-class DataProcSparkSqlOperator(BaseOperator):
- """
- Start a Spark SQL query Job on a Cloud DataProc cluster.
-
- :param query: The query or reference to the query file (q extension). (templated)
- :type query: string
- :param query_uri: The uri of a spark sql script on Cloud Storage.
- :type query_uri: string
- :param variables: Map of named parameters for the query. (templated)
- :type variables: dict
- :param job_name: The job name used in the DataProc cluster. This
- name by default is the task_id appended with the execution data, but can
- be templated. The name will always be appended with a random number to
- avoid name clashes. (templated)
- :type job_name: string
- :param cluster_name: The name of the DataProc cluster. (templated)
- :type cluster_name: string
- :param dataproc_spark_properties: Map for the Pig properties. Ideal to put in
- default arguments
- :type dataproc_spark_properties: dict
- :param dataproc_spark_jars: URIs to jars provisioned in Cloud Storage (example:
- for UDFs and libs) and are ideal to put in default arguments.
- :type dataproc_spark_jars: list
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param region: The specified region where the dataproc cluster is created.
- :type region: string
- """
- template_fields = ['query', 'variables', 'job_name', 'cluster_name', 'dataproc_jars']
- template_ext = ('.q',)
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- query=None,
- query_uri=None,
- variables=None,
- job_name='{{task.task_id}}_{{ds_nodash}}',
- cluster_name='cluster-1',
- dataproc_spark_properties=None,
- dataproc_spark_jars=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- region='global',
- *args,
- **kwargs):
-
- super(DataProcSparkSqlOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.query = query
- self.query_uri = query_uri
- self.variables = variables
- self.job_name = job_name
- self.cluster_name = cluster_name
- self.dataproc_properties = dataproc_spark_properties
- self.dataproc_jars = dataproc_spark_jars
- self.region = region
-
- def execute(self, context):
- hook = DataProcHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- job = hook.create_job_template(self.task_id, self.cluster_name, "sparkSqlJob",
- self.dataproc_properties)
-
- if self.query is None:
- job.add_query_uri(self.query_uri)
- else:
- job.add_query(self.query)
- job.add_variables(self.variables)
- job.add_jar_file_uris(self.dataproc_jars)
- job.set_job_name(self.job_name)
-
- hook.submit(hook.project_id, job.build(), self.region)
-
-
-class DataProcSparkOperator(BaseOperator):
- """
- Start a Spark Job on a Cloud DataProc cluster.
-
- :param main_jar: URI of the job jar provisioned on Cloud Storage. (use this or
- the main_class, not both together).
- :type main_jar: string
- :param main_class: Name of the job class. (use this or the main_jar, not both
- together).
- :type main_class: string
- :param arguments: Arguments for the job. (templated)
- :type arguments: list
- :param archives: List of archived files that will be unpacked in the work
- directory. Should be stored in Cloud Storage.
- :type archives: list
- :param files: List of files to be copied to the working directory
- :type files: list
- :param job_name: The job name used in the DataProc cluster. This
- name by default is the task_id appended with the execution data, but can
- be templated. The name will always be appended with a random number to
- avoid name clashes. (templated)
- :type job_name: string
- :param cluster_name: The name of the DataProc cluster. (templated)
- :type cluster_name: string
- :param dataproc_spark_properties: Map for the Pig properties. Ideal to put in
- default arguments
- :type dataproc_spark_properties: dict
- :param dataproc_spark_jars: URIs to jars provisioned in Cloud Storage (example:
- for UDFs and libs) and are ideal to put in default arguments.
- :type dataproc_spark_jars: list
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param region: The specified region where the dataproc cluster is created.
- :type region: string
- """
-
- template_fields = ['arguments', 'job_name', 'cluster_name', 'dataproc_jars']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- main_jar=None,
- main_class=None,
- arguments=None,
- archives=None,
- files=None,
- job_name='{{task.task_id}}_{{ds_nodash}}',
- cluster_name='cluster-1',
- dataproc_spark_properties=None,
- dataproc_spark_jars=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- region='global',
- *args,
- **kwargs):
-
- super(DataProcSparkOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.main_jar = main_jar
- self.main_class = main_class
- self.arguments = arguments
- self.archives = archives
- self.files = files
- self.job_name = job_name
- self.cluster_name = cluster_name
- self.dataproc_properties = dataproc_spark_properties
- self.dataproc_jars = dataproc_spark_jars
- self.region = region
-
- def execute(self, context):
- hook = DataProcHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
- job = hook.create_job_template(self.task_id, self.cluster_name, "sparkJob",
- self.dataproc_properties)
-
- job.set_main(self.main_jar, self.main_class)
- job.add_args(self.arguments)
- job.add_jar_file_uris(self.dataproc_jars)
- job.add_archive_uris(self.archives)
- job.add_file_uris(self.files)
- job.set_job_name(self.job_name)
-
- hook.submit(hook.project_id, job.build(), self.region)
-
-
-class DataProcHadoopOperator(BaseOperator):
- """
- Start a Hadoop Job on a Cloud DataProc cluster.
-
- :param main_jar: URI of the job jar provisioned on Cloud Storage. (use this or
- the main_class, not both together).
- :type main_jar: string
- :param main_class: Name of the job class. (use this or the main_jar, not both
- together).
- :type main_class: string
- :param arguments: Arguments for the job. (templated)
- :type arguments: list
- :param archives: List of archived files that will be unpacked in the work
- directory. Should be stored in Cloud Storage.
- :type archives: list
- :param files: List of files to be copied to the working directory
- :type files: list
- :param job_name: The job name used in the DataProc cluster. This
- name by default is the task_id appended with the execution data, but can
- be templated. The name will always be appended with a random number to
- avoid name clashes. (templated)
- :type job_name: string
- :param cluster_name: The name of the DataProc cluster. (templated)
- :type cluster_name: string
- :param dataproc_hadoop_properties: Map for the Pig properties. Ideal to put in
- default arguments
- :type dataproc_hadoop_properties: dict
- :param dataproc_hadoop_jars: URIs to jars provisioned in Cloud Storage (example:
- for UDFs and libs) and are ideal to put in default arguments.
- :type dataproc_hadoop_jars: list
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param region: The specified region where the dataproc cluster is created.
- :type region: string
- """
-
- template_fields = ['arguments', 'job_name', 'cluster_name', 'dataproc_jars']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- main_jar=None,
- main_class=None,
- arguments=None,
- archives=None,
- files=None,
- job_name='{{task.task_id}}_{{ds_nodash}}',
- cluster_name='cluster-1',
- dataproc_hadoop_properties=None,
- dataproc_hadoop_jars=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- region='global',
- *args,
- **kwargs):
-
- super(DataProcHadoopOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.main_jar = main_jar
- self.main_class = main_class
- self.arguments = arguments
- self.archives = archives
- self.files = files
- self.job_name = job_name
- self.cluster_name = cluster_name
- self.dataproc_properties = dataproc_hadoop_properties
- self.dataproc_jars = dataproc_hadoop_jars
- self.region = region
-
- def execute(self, context):
- hook = DataProcHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
- job = hook.create_job_template(self.task_id, self.cluster_name, "hadoopJob",
- self.dataproc_properties)
-
- job.set_main(self.main_jar, self.main_class)
- job.add_args(self.arguments)
- job.add_jar_file_uris(self.dataproc_jars)
- job.add_archive_uris(self.archives)
- job.add_file_uris(self.files)
- job.set_job_name(self.job_name)
-
- hook.submit(hook.project_id, job.build(), self.region)
-
-
-class DataProcPySparkOperator(BaseOperator):
- """
- Start a PySpark Job on a Cloud DataProc cluster.
-
- :param main: [Required] The Hadoop Compatible Filesystem (HCFS) URI of the main
- Python file to use as the driver. Must be a .py file.
- :type main: string
- :param arguments: Arguments for the job. (templated)
- :type arguments: list
- :param archives: List of archived files that will be unpacked in the work
- directory. Should be stored in Cloud Storage.
- :type archives: list
- :param files: List of files to be copied to the working directory
- :type files: list
- :param pyfiles: List of Python files to pass to the PySpark framework.
- Supported file types: .py, .egg, and .zip
- :type pyfiles: list
- :param job_name: The job name used in the DataProc cluster. This
- name by default is the task_id appended with the execution data, but can
- be templated. The name will always be appended with a random number to
- avoid name clashes. (templated)
- :type job_name: string
- :param cluster_name: The name of the DataProc cluster.
- :type cluster_name: string
- :param dataproc_pyspark_properties: Map for the Pig properties. Ideal to put in
- default arguments
- :type dataproc_pyspark_properties: dict
- :param dataproc_pyspark_jars: URIs to jars provisioned in Cloud Storage (example:
- for UDFs and libs) and are ideal to put in default arguments.
- :type dataproc_pyspark_jars: list
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- :param region: The specified region where the dataproc cluster is created.
- :type region: string
- """
-
- template_fields = ['arguments', 'job_name', 'cluster_name', 'dataproc_jars']
- ui_color = '#0273d4'
-
- @staticmethod
- def _generate_temp_filename(filename):
- dt = time.strftime('%Y%m%d%H%M%S')
- return "{}_{}_{}".format(dt, str(uuid.uuid1())[:8], ntpath.basename(filename))
-
- """
- Upload a local file to a Google Cloud Storage bucket
- """
- def _upload_file_temp(self, bucket, local_file):
- temp_filename = self._generate_temp_filename(local_file)
- if not bucket:
- raise AirflowException(
- "If you want Airflow to upload the local file to a temporary bucket, set "
- "the 'temp_bucket' key in the connection string")
-
- self.log.info("Uploading %s to %s", local_file, temp_filename)
-
- GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.gcp_conn_id
- ).upload(
- bucket=bucket,
- object=temp_filename,
- mime_type='application/x-python',
- filename=local_file
- )
- return "gs://{}/{}".format(bucket, temp_filename)
-
- @apply_defaults
- def __init__(
- self,
- main,
- arguments=None,
- archives=None,
- pyfiles=None,
- files=None,
- job_name='{{task.task_id}}_{{ds_nodash}}',
- cluster_name='cluster-1',
- dataproc_pyspark_properties=None,
- dataproc_pyspark_jars=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- region='global',
- *args,
- **kwargs):
-
- super(DataProcPySparkOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.main = main
- self.arguments = arguments
- self.archives = archives
- self.files = files
- self.pyfiles = pyfiles
- self.job_name = job_name
- self.cluster_name = cluster_name
- self.dataproc_properties = dataproc_pyspark_properties
- self.dataproc_jars = dataproc_pyspark_jars
- self.region = region
-
- def execute(self, context):
- hook = DataProcHook(
- gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to
- )
- job = hook.create_job_template(
- self.task_id, self.cluster_name, "pysparkJob", self.dataproc_properties)
-
- # Check if the file is local, if that is the case, upload it to a bucket
- if os.path.isfile(self.main):
- cluster_info = hook.get_cluster(
- project_id=hook.project_id,
- region=self.region,
- cluster_name=self.cluster_name
- )
- bucket = cluster_info['config']['configBucket']
- self.main = self._upload_file_temp(bucket, self.main)
- job.set_python_main(self.main)
-
- job.add_args(self.arguments)
- job.add_jar_file_uris(self.dataproc_jars)
- job.add_archive_uris(self.archives)
- job.add_file_uris(self.files)
- job.add_python_file_uris(self.pyfiles)
- job.set_job_name(self.job_name)
-
- hook.submit(hook.project_id, job.build(), self.region)
-
-
-class DataprocWorkflowTemplateBaseOperator(BaseOperator):
- @apply_defaults
- def __init__(self,
- project_id,
- region='global',
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(DataprocWorkflowTemplateBaseOperator, self).__init__(*args, **kwargs)
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.project_id = project_id
- self.region = region
- self.hook = DataProcHook(
- gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- api_version='v1beta2'
- )
-
- def execute(self, context):
- self.hook.await(self.start())
-
- def start(self, context):
- raise AirflowException('plese start a workflow operation')
-
-
-class DataprocWorkflowTemplateInstantiateOperator(DataprocWorkflowTemplateBaseOperator):
- """
- Instantiate a WorkflowTemplate on Google Cloud Dataproc. The operator will wait
- until the WorkflowTemplate is finished executing.
-
- .. seealso::
- Please refer to:
- https://cloud.google.com/dataproc/docs/reference/rest/v1beta2/projects.regions.workflowTemplates/instantiate
-
- :param template_id: The id of the template. (templated)
- :type template_id: string
- :param project_id: The ID of the google cloud project in which
- the template runs
- :type project_id: string
- :param region: leave as 'global', might become relevant in the future
- :type region: string
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
-
- template_fields = ['template_id']
-
- @apply_defaults
- def __init__(self, template_id, *args, **kwargs):
- (super(DataprocWorkflowTemplateInstantiateOperator, self)
- .__init__(*args, **kwargs))
- self.template_id = template_id
-
- def start(self):
- self.log.info('Instantiating Template: %s', self.template_id)
- return (
- self.hook.get_conn().projects().regions().workflowTemplates()
- .instantiate(
- name=('projects/%s/regions/%s/workflowTemplates/%s' %
- (self.project_id, self.region, self.template_id)),
- body={'instanceId': str(uuid.uuid1())})
- .execute())
-
-
-class DataprocWorkflowTemplateInstantiateInlineOperator(
- DataprocWorkflowTemplateBaseOperator):
- """
- Instantiate a WorkflowTemplate Inline on Google Cloud Dataproc. The operator will
- wait until the WorkflowTemplate is finished executing.
-
- .. seealso::
- Please refer to:
- https://cloud.google.com/dataproc/docs/reference/rest/v1beta2/projects.regions.workflowTemplates/instantiateInline
-
- :param template: The template contents. (templated)
- :type template: map
- :param project_id: The ID of the google cloud project in which
- the template runs
- :type project_id: string
- :param region: leave as 'global', might become relevant in the future
- :type region: string
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- """
-
- template_fields = ['template']
-
- @apply_defaults
- def __init__(self, template, *args, **kwargs):
- (super(DataprocWorkflowTemplateInstantiateInlineOperator, self)
- .__init__(*args, **kwargs))
- self.template = template
-
- def start(self):
- self.log.info('Instantiating Inline Template')
- return (
- self.hook.get_conn().projects().regions().workflowTemplates()
- .instantiateInline(
- parent='projects/%s/regions/%s' % (self.project_id, self.region),
- instanceId=str(uuid.uuid1()),
- body=self.template)
- .execute())
diff --git a/airflow/contrib/operators/datastore_export_operator.py b/airflow/contrib/operators/datastore_export_operator.py
deleted file mode 100644
index f6dc7cc571..0000000000
--- a/airflow/contrib/operators/datastore_export_operator.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.datastore_hook import DatastoreHook
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class DatastoreExportOperator(BaseOperator):
- """
- Export entities from Google Cloud Datastore to Cloud Storage
-
- :param bucket: name of the cloud storage bucket to backup data
- :type bucket: string
- :param namespace: optional namespace path in the specified Cloud Storage bucket
- to backup data. If this namespace does not exist in GCS, it will be created.
- :type namespace: str
- :param datastore_conn_id: the name of the Datastore connection id to use
- :type datastore_conn_id: string
- :param cloud_storage_conn_id: the name of the cloud storage connection id to
- force-write backup
- :type cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param entity_filter: description of what data from the project is included in the
- export, refer to
- https://cloud.google.com/datastore/docs/reference/rest/Shared.Types/EntityFilter
- :type entity_filter: dict
- :param labels: client-assigned labels for cloud storage
- :type labels: dict
- :param polling_interval_in_seconds: number of seconds to wait before polling for
- execution status again
- :type polling_interval_in_seconds: int
- :param overwrite_existing: if the storage bucket + namespace is not empty, it will be
- emptied prior to exports. This enables overwriting existing backups.
- :type overwrite_existing: bool
- :param xcom_push: push operation name to xcom for reference
- :type xcom_push: bool
- """
-
- @apply_defaults
- def __init__(self,
- bucket,
- namespace=None,
- datastore_conn_id='google_cloud_default',
- cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- entity_filter=None,
- labels=None,
- polling_interval_in_seconds=10,
- overwrite_existing=False,
- xcom_push=False,
- *args,
- **kwargs):
- super(DatastoreExportOperator, self).__init__(*args, **kwargs)
- self.datastore_conn_id = datastore_conn_id
- self.cloud_storage_conn_id = cloud_storage_conn_id
- self.delegate_to = delegate_to
- self.bucket = bucket
- self.namespace = namespace
- self.entity_filter = entity_filter
- self.labels = labels
- self.polling_interval_in_seconds = polling_interval_in_seconds
- self.overwrite_existing = overwrite_existing
- self.xcom_push = xcom_push
-
- def execute(self, context):
- self.log.info('Exporting data to Cloud Storage bucket ' + self.bucket)
-
- if self.overwrite_existing and self.namespace:
- gcs_hook = GoogleCloudStorageHook(self.cloud_storage_conn_id)
- objects = gcs_hook.list(self.bucket, prefix=self.namespace)
- for o in objects:
- gcs_hook.delete(self.bucket, o)
-
- ds_hook = DatastoreHook(self.datastore_conn_id, self.delegate_to)
- result = ds_hook.export_to_storage_bucket(bucket=self.bucket,
- namespace=self.namespace,
- entity_filter=self.entity_filter,
- labels=self.labels)
- operation_name = result['name']
- result = ds_hook.poll_operation_until_done(operation_name,
- self.polling_interval_in_seconds)
-
- state = result['metadata']['common']['state']
- if state != 'SUCCESSFUL':
- raise AirflowException('Operation failed: result={}'.format(result))
-
- if self.xcom_push:
- return result
diff --git a/airflow/contrib/operators/datastore_import_operator.py b/airflow/contrib/operators/datastore_import_operator.py
deleted file mode 100644
index 401d36e05b..0000000000
--- a/airflow/contrib/operators/datastore_import_operator.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.datastore_hook import DatastoreHook
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class DatastoreImportOperator(BaseOperator):
- """
- Import entities from Cloud Storage to Google Cloud Datastore
-
- :param bucket: container in Cloud Storage to store data
- :type bucket: string
- :param file: path of the backup metadata file in the specified Cloud Storage bucket.
- It should have the extension .overall_export_metadata
- :type file: string
- :param namespace: optional namespace of the backup metadata file in
- the specified Cloud Storage bucket.
- :type namespace: str
- :param entity_filter: description of what data from the project is included in
- the export, refer to
- https://cloud.google.com/datastore/docs/reference/rest/Shared.Types/EntityFilter
- :type entity_filter: dict
- :param labels: client-assigned labels for cloud storage
- :type labels: dict
- :param datastore_conn_id: the name of the connection id to use
- :type datastore_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param polling_interval_in_seconds: number of seconds to wait before polling for
- execution status again
- :type polling_interval_in_seconds: int
- :param xcom_push: push operation name to xcom for reference
- :type xcom_push: bool
- """
-
- @apply_defaults
- def __init__(self,
- bucket,
- file,
- namespace=None,
- entity_filter=None,
- labels=None,
- datastore_conn_id='google_cloud_default',
- delegate_to=None,
- polling_interval_in_seconds=10,
- xcom_push=False,
- *args,
- **kwargs):
- super(DatastoreImportOperator, self).__init__(*args, **kwargs)
- self.datastore_conn_id = datastore_conn_id
- self.delegate_to = delegate_to
- self.bucket = bucket
- self.file = file
- self.namespace = namespace
- self.entity_filter = entity_filter
- self.labels = labels
- self.polling_interval_in_seconds = polling_interval_in_seconds
- self.xcom_push = xcom_push
-
- def execute(self, context):
- self.log.info('Importing data from Cloud Storage bucket %s', self.bucket)
- ds_hook = DatastoreHook(self.datastore_conn_id, self.delegate_to)
- result = ds_hook.import_from_storage_bucket(bucket=self.bucket,
- file=self.file,
- namespace=self.namespace,
- entity_filter=self.entity_filter,
- labels=self.labels)
- operation_name = result['name']
- result = ds_hook.poll_operation_until_done(operation_name,
- self.polling_interval_in_seconds)
-
- state = result['metadata']['common']['state']
- if state != 'SUCCESSFUL':
- raise AirflowException('Operation failed: result={}'.format(result))
-
- if self.xcom_push:
- return result
diff --git a/airflow/contrib/operators/discord_webhook_operator.py b/airflow/contrib/operators/discord_webhook_operator.py
deleted file mode 100644
index 6b44eacbf5..0000000000
--- a/airflow/contrib/operators/discord_webhook_operator.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.discord_webhook_hook import DiscordWebhookHook
-from airflow.exceptions import AirflowException
-from airflow.operators.http_operator import SimpleHttpOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class DiscordWebhookOperator(SimpleHttpOperator):
- """
- This operator allows you to post messages to Discord using incoming webhooks.
- Takes a Discord connection ID with a default relative webhook endpoint. The
- default endpoint can be overridden using the webhook_endpoint parameter
- (https://discordapp.com/developers/docs/resources/webhook).
-
- Each Discord webhook can be pre-configured to use a specific username and
- avatar_url. You can override these defaults in this operator.
-
- :param http_conn_id: Http connection ID with host as "https://discord.com/api/" and
- default webhook endpoint in the extra field in the form of
- {"webhook_endpoint": "webhooks/{webhook.id}/{webhook.token}"}
- :type http_conn_id: str
- :param webhook_endpoint: Discord webhook endpoint in the form of
- "webhooks/{webhook.id}/{webhook.token}"
- :type webhook_endpoint: str
- :param message: The message you want to send to your Discord channel
- (max 2000 characters). (templated)
- :type message: str
- :param username: Override the default username of the webhook. (templated)
- :type username: str
- :param avatar_url: Override the default avatar of the webhook
- :type avatar_url: str
- :param tts: Is a text-to-speech message
- :type tts: bool
- :param proxy: Proxy to use to make the Discord webhook call
- :type proxy: str
- """
-
- template_fields = ('username', 'message')
-
- @apply_defaults
- def __init__(self,
- http_conn_id=None,
- webhook_endpoint=None,
- message="",
- username=None,
- avatar_url=None,
- tts=False,
- proxy=None,
- *args,
- **kwargs):
- super(DiscordWebhookOperator, self).__init__(endpoint=webhook_endpoint,
- *args,
- **kwargs)
-
- if not http_conn_id:
- raise AirflowException('No valid Discord http_conn_id supplied.')
-
- self.http_conn_id = http_conn_id
- self.webhook_endpoint = webhook_endpoint
- self.message = message
- self.username = username
- self.avatar_url = avatar_url
- self.tts = tts
- self.proxy = proxy
- self.hook = None
-
- def execute(self, context):
- """
- Call the DiscordWebhookHook to post message
- """
- self.hook = DiscordWebhookHook(
- self.http_conn_id,
- self.webhook_endpoint,
- self.message,
- self.username,
- self.avatar_url,
- self.tts,
- self.proxy
- )
- self.hook.execute()
diff --git a/airflow/contrib/operators/druid_operator.py b/airflow/contrib/operators/druid_operator.py
deleted file mode 100644
index 426393deeb..0000000000
--- a/airflow/contrib/operators/druid_operator.py
+++ /dev/null
@@ -1,62 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-
-from airflow.hooks.druid_hook import DruidHook
-from airflow.models import BaseOperator
-
-
-class DruidOperator(BaseOperator):
- """
- Allows to submit a task directly to druid
-
- :param json_index_file: The filepath to the druid index specification
- :type json_index_file: str
- :param druid_ingest_conn_id: The connection id of the Druid overlord which
- accepts index jobs
- :type druid_ingest_conn_id: str
- """
- template_fields = ('index_spec_str',)
- template_ext = ('.json',)
-
- def __init__(self, json_index_file,
- druid_ingest_conn_id='druid_ingest_default',
- max_ingestion_time=None,
- *args, **kwargs):
- super(DruidOperator, self).__init__(*args, **kwargs)
- self.conn_id = druid_ingest_conn_id
- self.max_ingestion_time = max_ingestion_time
-
- with open(json_index_file) as data_file:
- index_spec = json.load(data_file)
- self.index_spec_str = json.dumps(
- index_spec,
- sort_keys=True,
- indent=4,
- separators=(',', ': ')
- )
-
- def execute(self, context):
- hook = DruidHook(
- druid_ingest_conn_id=self.conn_id,
- max_ingestion_time=self.max_ingestion_time
- )
- self.log.info("Sumitting %s", self.index_spec_str)
- hook.submit_indexing_job(self.index_spec_str)
diff --git a/airflow/contrib/operators/ecs_operator.py b/airflow/contrib/operators/ecs_operator.py
deleted file mode 100644
index 60540f5e28..0000000000
--- a/airflow/contrib/operators/ecs_operator.py
+++ /dev/null
@@ -1,141 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 sys
-
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
-
-from airflow.contrib.hooks.aws_hook import AwsHook
-
-
-class ECSOperator(BaseOperator):
- """
- Execute a task on AWS EC2 Container Service
-
- :param task_definition: the task definition name on EC2 Container Service
- :type task_definition: str
- :param cluster: the cluster name on EC2 Container Service
- :type cluster: str
- :param: overrides: the same parameter that boto3 will receive (templated):
- http://boto3.readthedocs.org/en/latest/reference/services/ecs.html#ECS.Client.run_task
- :type: overrides: dict
- :param aws_conn_id: connection id of AWS credentials / region name. If None,
- credential boto3 strategy will be used
- (http://boto3.readthedocs.io/en/latest/guide/configuration.html).
- :type aws_conn_id: str
- :param region_name: region name to use in AWS Hook.
- Override the region_name in connection (if provided)
- :param launch_type: the launch type on which to run your task ('EC2' or 'FARGATE')
- :type: launch_type: str
- """
-
- ui_color = '#f0ede4'
- client = None
- arn = None
- template_fields = ('overrides',)
-
- @apply_defaults
- def __init__(self, task_definition, cluster, overrides,
- aws_conn_id=None, region_name=None, launch_type='EC2', **kwargs):
- super(ECSOperator, self).__init__(**kwargs)
-
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
- self.task_definition = task_definition
- self.cluster = cluster
- self.overrides = overrides
- self.launch_type = launch_type
-
- self.hook = self.get_hook()
-
- def execute(self, context):
- self.log.info(
- 'Running ECS Task - Task definition: %s - on cluster %s',
- self.task_definition, self.cluster
- )
- self.log.info('ECSOperator overrides: %s', self.overrides)
-
- self.client = self.hook.get_client_type(
- 'ecs',
- region_name=self.region_name
- )
-
- response = self.client.run_task(
- cluster=self.cluster,
- taskDefinition=self.task_definition,
- overrides=self.overrides,
- startedBy=self.owner,
- launchType=self.launch_type
- )
-
- failures = response['failures']
- if len(failures) > 0:
- raise AirflowException(response)
- self.log.info('ECS Task started: %s', response)
-
- self.arn = response['tasks'][0]['taskArn']
- self._wait_for_task_ended()
-
- self._check_success_task()
- self.log.info('ECS Task has been successfully executed: %s', response)
-
- def _wait_for_task_ended(self):
- waiter = self.client.get_waiter('tasks_stopped')
- waiter.config.max_attempts = sys.maxsize # timeout is managed by airflow
- waiter.wait(
- cluster=self.cluster,
- tasks=[self.arn]
- )
-
- def _check_success_task(self):
- response = self.client.describe_tasks(
- cluster=self.cluster,
- tasks=[self.arn]
- )
- self.log.info('ECS Task stopped, check status: %s', response)
-
- if len(response.get('failures', [])) > 0:
- raise AirflowException(response)
-
- for task in response['tasks']:
- containers = task['containers']
- for container in containers:
- if container.get('lastStatus') == 'STOPPED' and \
- container['exitCode'] != 0:
- raise AirflowException(
- 'This task is not in success state {}'.format(task))
- elif container.get('lastStatus') == 'PENDING':
- raise AirflowException('This task is still pending {}'.format(task))
- elif 'error' in container.get('reason', '').lower():
- raise AirflowException(
- 'This containers encounter an error during launching : {}'.
- format(container.get('reason', '').lower()))
-
- def get_hook(self):
- return AwsHook(
- aws_conn_id=self.aws_conn_id
- )
-
- def on_kill(self):
- response = self.client.stop_task(
- cluster=self.cluster,
- task=self.arn,
- reason='Task killed by the user')
- self.log.info(response)
diff --git a/airflow/contrib/operators/emr_add_steps_operator.py b/airflow/contrib/operators/emr_add_steps_operator.py
deleted file mode 100644
index 643ffe9c1b..0000000000
--- a/airflow/contrib/operators/emr_add_steps_operator.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.models import BaseOperator
-from airflow.utils import apply_defaults
-from airflow.exceptions import AirflowException
-from airflow.contrib.hooks.emr_hook import EmrHook
-
-
-class EmrAddStepsOperator(BaseOperator):
- """
- An operator that adds steps to an existing EMR job_flow.
-
- :param job_flow_id: id of the JobFlow to add steps to. (templated)
- :type job_flow_name: str
- :param aws_conn_id: aws connection to uses
- :type aws_conn_id: str
- :param steps: boto3 style steps to be added to the jobflow. (templated)
- :type steps: list
- """
- template_fields = ['job_flow_id', 'steps']
- template_ext = ()
- ui_color = '#f9c915'
-
- @apply_defaults
- def __init__(
- self,
- job_flow_id,
- aws_conn_id='s3_default',
- steps=None,
- *args, **kwargs):
- super(EmrAddStepsOperator, self).__init__(*args, **kwargs)
- steps = steps or []
- self.job_flow_id = job_flow_id
- self.aws_conn_id = aws_conn_id
- self.steps = steps
-
- def execute(self, context):
- emr = EmrHook(aws_conn_id=self.aws_conn_id).get_conn()
-
- self.log.info('Adding steps to %s', self.job_flow_id)
- response = emr.add_job_flow_steps(JobFlowId=self.job_flow_id, Steps=self.steps)
-
- if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
- raise AirflowException('Adding steps failed: %s' % response)
- else:
- self.log.info('Steps %s added to JobFlow', response['StepIds'])
- return response['StepIds']
diff --git a/airflow/contrib/operators/emr_create_job_flow_operator.py b/airflow/contrib/operators/emr_create_job_flow_operator.py
deleted file mode 100644
index 89be12f065..0000000000
--- a/airflow/contrib/operators/emr_create_job_flow_operator.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.emr_hook import EmrHook
-from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
-from airflow.exceptions import AirflowException
-
-
-class EmrCreateJobFlowOperator(BaseOperator):
- """
- Creates an EMR JobFlow, reading the config from the EMR connection.
- A dictionary of JobFlow overrides can be passed that override
- the config from the connection.
-
- :param aws_conn_id: aws connection to uses
- :type aws_conn_id: str
- :param emr_conn_id: emr connection to use
- :type emr_conn_id: str
- :param job_flow_overrides: boto3 style arguments to override
- emr_connection extra. (templated)
- :type steps: dict
- """
- template_fields = ['job_flow_overrides']
- template_ext = ()
- ui_color = '#f9c915'
-
- @apply_defaults
- def __init__(
- self,
- aws_conn_id='s3_default',
- emr_conn_id='emr_default',
- job_flow_overrides=None,
- *args, **kwargs):
- super(EmrCreateJobFlowOperator, self).__init__(*args, **kwargs)
- self.aws_conn_id = aws_conn_id
- self.emr_conn_id = emr_conn_id
- if job_flow_overrides is None:
- job_flow_overrides = {}
- self.job_flow_overrides = job_flow_overrides
-
- def execute(self, context):
- emr = EmrHook(aws_conn_id=self.aws_conn_id, emr_conn_id=self.emr_conn_id)
-
- self.log.info(
- 'Creating JobFlow using aws-conn-id: %s, emr-conn-id: %s',
- self.aws_conn_id, self.emr_conn_id
- )
- response = emr.create_job_flow(self.job_flow_overrides)
-
- if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
- raise AirflowException('JobFlow creation failed: %s' % response)
- else:
- self.log.info('JobFlow with id %s created', response['JobFlowId'])
- return response['JobFlowId']
diff --git a/airflow/contrib/operators/emr_terminate_job_flow_operator.py b/airflow/contrib/operators/emr_terminate_job_flow_operator.py
deleted file mode 100644
index 50407a14ac..0000000000
--- a/airflow/contrib/operators/emr_terminate_job_flow_operator.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.models import BaseOperator
-from airflow.utils import apply_defaults
-from airflow.exceptions import AirflowException
-from airflow.contrib.hooks.emr_hook import EmrHook
-
-
-class EmrTerminateJobFlowOperator(BaseOperator):
- """
- Operator to terminate EMR JobFlows.
-
- :param job_flow_id: id of the JobFlow to terminate. (templated)
- :type job_flow_name: str
- :param aws_conn_id: aws connection to uses
- :type aws_conn_id: str
- """
- template_fields = ['job_flow_id']
- template_ext = ()
- ui_color = '#f9c915'
-
- @apply_defaults
- def __init__(
- self,
- job_flow_id,
- aws_conn_id='s3_default',
- *args, **kwargs):
- super(EmrTerminateJobFlowOperator, self).__init__(*args, **kwargs)
- self.job_flow_id = job_flow_id
- self.aws_conn_id = aws_conn_id
-
- def execute(self, context):
- emr = EmrHook(aws_conn_id=self.aws_conn_id).get_conn()
-
- self.log.info('Terminating JobFlow %s', self.job_flow_id)
- response = emr.terminate_job_flows(JobFlowIds=[self.job_flow_id])
-
- if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
- raise AirflowException('JobFlow termination failed: %s' % response)
- else:
- self.log.info('JobFlow with id %s terminated', self.job_flow_id)
diff --git a/airflow/contrib/operators/file_to_gcs.py b/airflow/contrib/operators/file_to_gcs.py
deleted file mode 100644
index 807385b43c..0000000000
--- a/airflow/contrib/operators/file_to_gcs.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class FileToGoogleCloudStorageOperator(BaseOperator):
- """
- Uploads a file to Google Cloud Storage
-
- :param src: Path to the local file. (templated)
- :type src: string
- :param dst: Destination path within the specified bucket. (templated)
- :type dst: string
- :param bucket: The bucket to upload to. (templated)
- :type bucket: string
- :param google_cloud_storage_conn_id: The Airflow connection ID to upload with
- :type google_cloud_storage_conn_id: string
- :param mime_type: The mime-type string
- :type mime_type: string
- :param delegate_to: The account to impersonate, if any
- :type delegate_to: string
- """
- template_fields = ('src', 'dst', 'bucket')
-
- @apply_defaults
- def __init__(self,
- src,
- dst,
- bucket,
- google_cloud_storage_conn_id='google_cloud_default',
- mime_type='application/octet-stream',
- delegate_to=None,
- *args,
- **kwargs):
- super(FileToGoogleCloudStorageOperator, self).__init__(*args, **kwargs)
- self.src = src
- self.dst = dst
- self.bucket = bucket
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.mime_type = mime_type
- self.delegate_to = delegate_to
-
- def execute(self, context):
- """
- Uploads the file to Google cloud storage
- """
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
-
- hook.upload(
- bucket=self.bucket,
- object=self.dst,
- mime_type=self.mime_type,
- filename=self.src)
diff --git a/airflow/contrib/operators/file_to_wasb.py b/airflow/contrib/operators/file_to_wasb.py
deleted file mode 100644
index 797960c0e3..0000000000
--- a/airflow/contrib/operators/file_to_wasb.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.wasb_hook import WasbHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class FileToWasbOperator(BaseOperator):
- """
- Uploads a file to Azure Blob Storage.
-
- :param file_path: Path to the file to load. (templated)
- :type file_path: str
- :param container_name: Name of the container. (templated)
- :type container_name: str
- :param blob_name: Name of the blob. (templated)
- :type blob_name: str
- :param wasb_conn_id: Reference to the wasb connection.
- :type wasb_conn_id: str
- :param load_options: Optional keyword arguments that
- `WasbHook.load_file()` takes.
- :type load_options: dict
- """
- template_fields = ('file_path', 'container_name', 'blob_name')
-
- @apply_defaults
- def __init__(self, file_path, container_name, blob_name,
- wasb_conn_id='wasb_default', load_options=None, *args,
- **kwargs):
- super(FileToWasbOperator, self).__init__(*args, **kwargs)
- if load_options is None:
- load_options = {}
- self.file_path = file_path
- self.container_name = container_name
- self.blob_name = blob_name
- self.wasb_conn_id = wasb_conn_id
- self.load_options = load_options
-
- def execute(self, context):
- """Upload a file to Azure Blob Storage."""
- hook = WasbHook(wasb_conn_id=self.wasb_conn_id)
- self.log.info(
- 'Uploading {self.file_path} to wasb://{self.container_name} '
- 'as {self.blob_name}'.format(**locals())
- )
- hook.load_file(self.file_path, self.container_name,
- self.blob_name, **self.load_options)
diff --git a/airflow/contrib/operators/gcp_container_operator.py b/airflow/contrib/operators/gcp_container_operator.py
deleted file mode 100644
index 5648b4d8a0..0000000000
--- a/airflow/contrib/operators/gcp_container_operator.py
+++ /dev/null
@@ -1,172 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow import AirflowException
-from airflow.contrib.hooks.gcp_container_hook import GKEClusterHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class GKEClusterDeleteOperator(BaseOperator):
- template_fields = ['project_id', 'gcp_conn_id', 'name', 'location', 'api_version']
-
- @apply_defaults
- def __init__(self,
- project_id,
- name,
- location,
- gcp_conn_id='google_cloud_default',
- api_version='v2',
- *args,
- **kwargs):
- """
- Deletes the cluster, including the Kubernetes endpoint and all worker nodes.
-
-
- To delete a certain cluster, you must specify the ``project_id``, the ``name``
- of the cluster, the ``location`` that the cluster is in, and the ``task_id``.
-
- **Operator Creation**: ::
-
- operator = GKEClusterDeleteOperator(
- task_id='cluster_delete',
- project_id='my-project',
- location='cluster-location'
- name='cluster-name')
-
- .. seealso::
- For more detail about deleting clusters have a look at the reference:
- https://google-cloud-python.readthedocs.io/en/latest/container/gapic/v1/api.html#google.cloud.container_v1.ClusterManagerClient.delete_cluster
-
- :param project_id: The Google Developers Console [project ID or project number]
- :type project_id: str
- :param name: The name of the resource to delete, in this case cluster name
- :type name: str
- :param location: The name of the Google Compute Engine zone in which the cluster
- resides.
- :type location: str
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: str
- :param api_version: The api version to use
- :type api_version: str
- """
- super(GKEClusterDeleteOperator, self).__init__(*args, **kwargs)
-
- self.project_id = project_id
- self.gcp_conn_id = gcp_conn_id
- self.location = location
- self.api_version = api_version
- self.name = name
-
- def _check_input(self):
- if not all([self.project_id, self.name, self.location]):
- self.log.error(
- 'One of (project_id, name, location) is missing or incorrect')
- raise AirflowException('Operator has incorrect or missing input.')
-
- def execute(self, context):
- self._check_input()
- hook = GKEClusterHook(self.project_id, self.location)
- delete_result = hook.delete_cluster(name=self.name)
- return delete_result
-
-
-class GKEClusterCreateOperator(BaseOperator):
- template_fields = ['project_id', 'gcp_conn_id', 'location', 'api_version', 'body']
-
- @apply_defaults
- def __init__(self,
- project_id,
- location,
- body={},
- gcp_conn_id='google_cloud_default',
- api_version='v2',
- *args,
- **kwargs):
- """
- Create a Google Kubernetes Engine Cluster of specified dimensions
- The operator will wait until the cluster is created.
-
- The **minimum** required to define a cluster to create is:
-
- ``dict()`` ::
- cluster_def = {'name': 'my-cluster-name',
- 'initial_node_count': 1}
-
- or
-
- ``Cluster`` proto ::
- from google.cloud.container_v1.types import Cluster
-
- cluster_def = Cluster(name='my-cluster-name', initial_node_count=1)
-
- **Operator Creation**: ::
-
- operator = GKEClusterCreateOperator(
- task_id='cluster_create',
- project_id='my-project',
- location='my-location'
- body=cluster_def)
-
- .. seealso::
- For more detail on about creating clusters have a look at the reference:
- https://google-cloud-python.readthedocs.io/en/latest/container/gapic/v1/types.html#google.cloud.container_v1.types.Cluster
-
- :param project_id: The Google Developers Console [project ID or project number]
- :type project_id: str
- :param location: The name of the Google Compute Engine zone in which the cluster
- resides.
- :type location: str
- :param body: The Cluster definition to create, can be protobuf or python dict, if
- dict it must match protobuf message Cluster
- :type body: dict or google.cloud.container_v1.types.Cluster
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
- :type gcp_conn_id: str
- :param api_version: The api version to use
- :type api_version: str
- """
- super(GKEClusterCreateOperator, self).__init__(*args, **kwargs)
-
- self.project_id = project_id
- self.gcp_conn_id = gcp_conn_id
- self.location = location
- self.api_version = api_version
- self.body = body
-
- def _check_input(self):
- if all([self.project_id, self.location, self.body]):
- if isinstance(self.body, dict) \
- and 'name' in self.body \
- and 'initial_node_count' in self.body:
- # Don't throw error
- return
- # If not dict, then must
- elif self.body.name and self.body.initial_node_count:
- return
-
- self.log.error(
- 'One of (project_id, location, body, body[\'name\'], '
- 'body[\'initial_node_count\']) is missing or incorrect')
- raise AirflowException('Operator has incorrect or missing input.')
-
- def execute(self, context):
- self._check_input()
- hook = GKEClusterHook(self.project_id, self.location)
- create_op = hook.create_cluster(cluster=self.body)
- return create_op
diff --git a/airflow/contrib/operators/gcs_download_operator.py b/airflow/contrib/operators/gcs_download_operator.py
deleted file mode 100644
index ce272aedd5..0000000000
--- a/airflow/contrib/operators/gcs_download_operator.py
+++ /dev/null
@@ -1,91 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 sys
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class GoogleCloudStorageDownloadOperator(BaseOperator):
- """
- Downloads a file from Google Cloud Storage.
-
- :param bucket: The Google cloud storage bucket where the object is. (templated)
- :type bucket: string
- :param object: The name of the object to download in the Google cloud
- storage bucket. (templated)
- :type object: string
- :param filename: The file path on the local file system (where the
- operator is being executed) that the file should be downloaded to. (templated)
- If no filename passed, the downloaded data will not be stored on the local file
- system.
- :type filename: string
- :param store_to_xcom_key: If this param is set, the operator will push
- the contents of the downloaded file to XCom with the key set in this
- parameter. If not set, the downloaded data will not be pushed to XCom. (templated)
- :type store_to_xcom_key: string
- :param google_cloud_storage_conn_id: The connection ID to use when
- connecting to Google cloud storage.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- """
- template_fields = ('bucket', 'object', 'filename', 'store_to_xcom_key',)
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- bucket,
- object,
- filename=None,
- store_to_xcom_key=None,
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(GoogleCloudStorageDownloadOperator, self).__init__(*args, **kwargs)
- self.bucket = bucket
- self.object = object
- self.filename = filename
- self.store_to_xcom_key = store_to_xcom_key
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- self.log.info('Executing download: %s, %s, %s', self.bucket,
- self.object, self.filename)
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to
- )
- file_bytes = hook.download(bucket=self.bucket,
- object=self.object,
- filename=self.filename)
- if self.store_to_xcom_key:
- if sys.getsizeof(file_bytes) < 48000:
- context['ti'].xcom_push(key=self.store_to_xcom_key, value=file_bytes)
- else:
- raise RuntimeError(
- 'The size of the downloaded file is too large to push to XCom!'
- )
- self.log.debug(file_bytes)
diff --git a/airflow/contrib/operators/gcs_list_operator.py b/airflow/contrib/operators/gcs_list_operator.py
deleted file mode 100644
index 6474453afa..0000000000
--- a/airflow/contrib/operators/gcs_list_operator.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class GoogleCloudStorageListOperator(BaseOperator):
- """
- List all objects from the bucket with the give string prefix and delimiter in name.
-
- This operator returns a python list with the name of objects which can be used by
- `xcom` in the downstream task.
-
- :param bucket: The Google cloud storage bucket to find the objects. (templated)
- :type bucket: string
- :param prefix: Prefix string which filters objects whose name begin with
- this prefix. (templated)
- :type prefix: string
- :param delimiter: The delimiter by which you want to filter the objects. (templated)
- For e.g to lists the CSV files from in a directory in GCS you would use
- delimiter='.csv'.
- :type delimiter: string
- :param google_cloud_storage_conn_id: The connection ID to use when
- connecting to Google cloud storage.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
-
- **Example**:
- The following Operator would list all the Avro files from ``sales/sales-2017``
- folder in ``data`` bucket. ::
-
- GCS_Files = GoogleCloudStorageListOperator(
- task_id='GCS_Files',
- bucket='data',
- prefix='sales/sales-2017/',
- delimiter='.avro',
- google_cloud_storage_conn_id=google_cloud_conn_id
- )
- """
- template_fields = ('bucket', 'prefix', 'delimiter')
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- bucket,
- prefix=None,
- delimiter=None,
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(GoogleCloudStorageListOperator, self).__init__(*args, **kwargs)
- self.bucket = bucket
- self.prefix = prefix
- self.delimiter = delimiter
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
-
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to
- )
-
- self.log.info('Getting list of the files. Bucket: %s; Delimiter: %s; Prefix: %s',
- self.bucket, self.delimiter, self.prefix)
-
- return hook.list(bucket=self.bucket,
- prefix=self.prefix,
- delimiter=self.delimiter)
diff --git a/airflow/contrib/operators/gcs_operator.py b/airflow/contrib/operators/gcs_operator.py
deleted file mode 100644
index ef5e8de504..0000000000
--- a/airflow/contrib/operators/gcs_operator.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.version import version
-
-
-class GoogleCloudStorageCreateBucketOperator(BaseOperator):
- """
- Creates a new bucket. Google Cloud Storage uses a flat namespace,
- so you can't create a bucket with a name that is already in use.
-
- .. seealso::
- For more information, see Bucket Naming Guidelines:
- https://cloud.google.com/storage/docs/bucketnaming.html#requirements
-
- :param bucket_name: The name of the bucket. (templated)
- :type bucket_name: string
- :param storage_class: This defines how objects in the bucket are stored
- and determines the SLA and the cost of storage (templated). Values include
-
- - ``MULTI_REGIONAL``
- - ``REGIONAL``
- - ``STANDARD``
- - ``NEARLINE``
- - ``COLDLINE``.
- If this value is not specified when the bucket is
- created, it will default to STANDARD.
- :type storage_class: string
- :param location: The location of the bucket. (templated)
- Object data for objects in the bucket resides in physical storage
- within this region. Defaults to US.
-
- .. seealso::
- https://developers.google.com/storage/docs/bucket-locations
-
- :type location: string
- :param project_id: The ID of the GCP Project. (templated)
- :type project_id: string
- :param labels: User-provided labels, in key/value pairs.
- :type labels: dict
- :param google_cloud_storage_conn_id: The connection ID to use when
- connecting to Google cloud storage.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must
- have domain-wide delegation enabled.
- :type delegate_to: string
-
- **Example**:
- The following Operator would create a new bucket ``test-bucket``
- with ``MULTI_REGIONAL`` storage class in ``EU`` region ::
-
- CreateBucket = GoogleCloudStorageCreateBucketOperator(
- task_id='CreateNewBucket',
- bucket_name='test-bucket',
- storage_class='MULTI_REGIONAL',
- location='EU',
- labels={'env': 'dev', 'team': 'airflow'},
- google_cloud_storage_conn_id='airflow-service-account'
- )
- """
-
- template_fields = ('bucket_name', 'storage_class',
- 'location', 'project_id')
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- bucket_name,
- storage_class='MULTI_REGIONAL',
- location='US',
- project_id=None,
- labels=None,
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(GoogleCloudStorageCreateBucketOperator, self).__init__(*args, **kwargs)
- self.bucket_name = bucket_name
- self.storage_class = storage_class
- self.location = location
- self.project_id = project_id
- self.labels = labels
-
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- if self.labels is not None:
- self.labels.update(
- {'airflow-version': 'v' + version.replace('.', '-').replace('+', '-')}
- )
-
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to
- )
-
- hook.create_bucket(bucket_name=self.bucket_name,
- storage_class=self.storage_class,
- location=self.location,
- project_id=self.project_id,
- labels=self.labels)
diff --git a/airflow/contrib/operators/gcs_to_bq.py b/airflow/contrib/operators/gcs_to_bq.py
deleted file mode 100644
index 3a7798030c..0000000000
--- a/airflow/contrib/operators/gcs_to_bq.py
+++ /dev/null
@@ -1,249 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.contrib.hooks.bigquery_hook import BigQueryHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class GoogleCloudStorageToBigQueryOperator(BaseOperator):
- """
- Loads files from Google cloud storage into BigQuery.
-
- The schema to be used for the BigQuery table may be specified in one of
- two ways. You may either directly pass the schema fields in, or you may
- point the operator to a Google cloud storage object name. The object in
- Google cloud storage must be a JSON file with the schema fields in it.
-
- :param bucket: The bucket to load from. (templated)
- :type bucket: string
- :param source_objects: List of Google cloud storage URIs to load from. (templated)
- If source_format is 'DATASTORE_BACKUP', the list must only contain a single URI.
- :type object: list
- :param destination_project_dataset_table: The dotted (<project>.)<dataset>.<table>
- BigQuery table to load data into. If <project> is not included,
- project will be the project defined in the connection json. (templated)
- :type destination_project_dataset_table: string
- :param schema_fields: If set, the schema field list as defined here:
- https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.load
- Should not be set when source_format is 'DATASTORE_BACKUP'.
- :type schema_fields: list
- :param schema_object: If set, a GCS object path pointing to a .json file that
- contains the schema for the table. (templated)
- :param schema_object: string
- :param source_format: File format to export.
- :type source_format: string
- :param compression: [Optional] The compression type of the data source.
- Possible values include GZIP and NONE.
- The default value is NONE.
- This setting is ignored for Google Cloud Bigtable,
- Google Cloud Datastore backups and Avro formats.
- :type compression: string
- :param create_disposition: The create disposition if the table doesn't exist.
- :type create_disposition: string
- :param skip_leading_rows: Number of rows to skip when loading from a CSV.
- :type skip_leading_rows: int
- :param write_disposition: The write disposition if the table already exists.
- :type write_disposition: string
- :param field_delimiter: The delimiter to use when loading from a CSV.
- :type field_delimiter: string
- :param max_bad_records: The maximum number of bad records that BigQuery can
- ignore when running the job.
- :type max_bad_records: int
- :param quote_character: The value that is used to quote data sections in a CSV file.
- :type quote_character: string
- :param ignore_unknown_values: [Optional] Indicates if BigQuery should allow
- extra values that are not represented in the table schema.
- If true, the extra values are ignored. If false, records with extra columns
- are treated as bad records, and if there are too many bad records, an
- invalid error is returned in the job result.
- :type ignore_unknown_values: bool
- :param allow_quoted_newlines: Whether to allow quoted newlines (true) or not (false).
- :type allow_quoted_newlines: boolean
- :param allow_jagged_rows: Accept rows that are missing trailing optional columns.
- The missing values are treated as nulls. If false, records with missing trailing
- columns are treated as bad records, and if there are too many bad records, an
- invalid error is returned in the job result. Only applicable to CSV, ignored
- for other formats.
- :type allow_jagged_rows: bool
- :param max_id_key: If set, the name of a column in the BigQuery table
- that's to be loaded. Thsi will be used to select the MAX value from
- BigQuery after the load occurs. The results will be returned by the
- execute() command, which in turn gets stored in XCom for future
- operators to use. This can be helpful with incremental loads--during
- future executions, you can pick up from the max ID.
- :type max_id_key: string
- :param bigquery_conn_id: Reference to a specific BigQuery hook.
- :type bigquery_conn_id: string
- :param google_cloud_storage_conn_id: Reference to a specific Google
- cloud storage hook.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any. For this to
- work, the service account making the request must have domain-wide
- delegation enabled.
- :type delegate_to: string
- :param schema_update_options: Allows the schema of the destination
- table to be updated as a side effect of the load job.
- :type schema_update_options: list
- :param src_fmt_configs: configure optional fields specific to the source format
- :type src_fmt_configs: dict
- :param external_table: Flag to specify if the destination table should be
- a BigQuery external table. Default Value is False.
- :type external_table: bool
- :param time_partitioning: configure optional time partitioning fields i.e.
- partition by field, type and expiration as per API specifications.
- Note that 'field' is not available in concurrency with
- dataset.table$partition.
- :type time_partitioning: dict
- """
- template_fields = ('bucket', 'source_objects',
- 'schema_object', 'destination_project_dataset_table')
- template_ext = ('.sql',)
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- bucket,
- source_objects,
- destination_project_dataset_table,
- schema_fields=None,
- schema_object=None,
- source_format='CSV',
- compression='NONE',
- create_disposition='CREATE_IF_NEEDED',
- skip_leading_rows=0,
- write_disposition='WRITE_EMPTY',
- field_delimiter=',',
- max_bad_records=0,
- quote_character=None,
- ignore_unknown_values=False,
- allow_quoted_newlines=False,
- allow_jagged_rows=False,
- max_id_key=None,
- bigquery_conn_id='bigquery_default',
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- schema_update_options=(),
- src_fmt_configs={},
- external_table=False,
- time_partitioning={},
- *args, **kwargs):
-
- super(GoogleCloudStorageToBigQueryOperator, self).__init__(*args, **kwargs)
-
- # GCS config
- self.bucket = bucket
- self.source_objects = source_objects
- self.schema_object = schema_object
-
- # BQ config
- self.destination_project_dataset_table = destination_project_dataset_table
- self.schema_fields = schema_fields
- self.source_format = source_format
- self.compression = compression
- self.create_disposition = create_disposition
- self.skip_leading_rows = skip_leading_rows
- self.write_disposition = write_disposition
- self.field_delimiter = field_delimiter
- self.max_bad_records = max_bad_records
- self.quote_character = quote_character
- self.ignore_unknown_values = ignore_unknown_values
- self.allow_quoted_newlines = allow_quoted_newlines
- self.allow_jagged_rows = allow_jagged_rows
- self.external_table = external_table
-
- self.max_id_key = max_id_key
- self.bigquery_conn_id = bigquery_conn_id
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
-
- self.schema_update_options = schema_update_options
- self.src_fmt_configs = src_fmt_configs
- self.time_partitioning = time_partitioning
-
- def execute(self, context):
- bq_hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to)
-
- if not self.schema_fields and \
- self.schema_object and \
- self.source_format != 'DATASTORE_BACKUP':
- gcs_hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- schema_fields = json.loads(gcs_hook.download(
- self.bucket,
- self.schema_object).decode("utf-8"))
- else:
- schema_fields = self.schema_fields
-
- source_uris = ['gs://{}/{}'.format(self.bucket, source_object)
- for source_object in self.source_objects]
- conn = bq_hook.get_conn()
- cursor = conn.cursor()
-
- if self.external_table:
- cursor.create_external_table(
- external_project_dataset_table=self.destination_project_dataset_table,
- schema_fields=schema_fields,
- source_uris=source_uris,
- source_format=self.source_format,
- compression=self.compression,
- skip_leading_rows=self.skip_leading_rows,
- field_delimiter=self.field_delimiter,
- max_bad_records=self.max_bad_records,
- quote_character=self.quote_character,
- ignore_unknown_values=self.ignore_unknown_values,
- allow_quoted_newlines=self.allow_quoted_newlines,
- allow_jagged_rows=self.allow_jagged_rows,
- src_fmt_configs=self.src_fmt_configs
- )
- else:
- cursor.run_load(
- destination_project_dataset_table=self.destination_project_dataset_table,
- schema_fields=schema_fields,
- source_uris=source_uris,
- source_format=self.source_format,
- create_disposition=self.create_disposition,
- skip_leading_rows=self.skip_leading_rows,
- write_disposition=self.write_disposition,
- field_delimiter=self.field_delimiter,
- max_bad_records=self.max_bad_records,
- quote_character=self.quote_character,
- ignore_unknown_values=self.ignore_unknown_values,
- allow_quoted_newlines=self.allow_quoted_newlines,
- allow_jagged_rows=self.allow_jagged_rows,
- schema_update_options=self.schema_update_options,
- src_fmt_configs=self.src_fmt_configs,
- time_partitioning=self.time_partitioning)
-
- if self.max_id_key:
- cursor.execute('SELECT MAX({}) FROM {}'.format(
- self.max_id_key,
- self.destination_project_dataset_table))
- row = cursor.fetchone()
- max_id = row[0] if row[0] else 0
- self.log.info(
- 'Loaded BQ data with max %s.%s=%s',
- self.destination_project_dataset_table, self.max_id_key, max_id
- )
- return max_id
diff --git a/airflow/contrib/operators/gcs_to_gcs.py b/airflow/contrib/operators/gcs_to_gcs.py
deleted file mode 100644
index 256685f90b..0000000000
--- a/airflow/contrib/operators/gcs_to_gcs.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class GoogleCloudStorageToGoogleCloudStorageOperator(BaseOperator):
- """
- Copies objects from a bucket to another, with renaming if requested.
-
- :param source_bucket: The source Google cloud storage bucket where the
- object is. (templated)
- :type source_bucket: string
- :param source_object: The source name of the object to copy in the Google cloud
- storage bucket. (templated)
- If wildcards are used in this argument:
- You can use only one wildcard for objects (filenames) within your
- bucket. The wildcard can appear inside the object name or at the
- end of the object name. Appending a wildcard to the bucket name is
- unsupported.
- :type source_object: string
- :param destination_bucket: The destination Google cloud storage bucket
- where the object should be. (templated)
- :type destination_bucket: string
- :param destination_object: The destination name of the object in the
- destination Google cloud storage bucket. (templated)
- If a wildcard is supplied in the source_object argument, this is the
- prefix that will be prepended to the final destination objects' paths.
- Note that the source path's part before the wildcard will be removed;
- if it needs to be retained it should be appended to destination_object.
- For example, with prefix ``foo/*`` and destination_object `'blah/``, the
- file ``foo/baz`` will be copied to ``blah/baz``; to retain the prefix write
- the destination_object as e.g. ``blah/foo``, in which case the copied file
- will be named ``blah/foo/baz``.
- :type destination_object: string
- :param move_object: When move object is True, the object is moved instead
- of copied to the new location.
- This is the equivalent of a mv command as opposed to a
- cp command.
- :type move_object: bool
- :param google_cloud_storage_conn_id: The connection ID to use when
- connecting to Google cloud storage.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
-
- **Examples**:
- The following Operator would copy a single file named
- ``sales/sales-2017/january.avro`` in the ``data`` bucket to the file named
- ``copied_sales/2017/january-backup.avro` in the ``data_backup`` bucket ::
- copy_single_file = GoogleCloudStorageToGoogleCloudStorageOperator(
- task_id='copy_single_file',
- source_bucket='data',
- source_object='sales/sales-2017/january.avro',
- destination_bucket='data_backup',
- destination_object='copied_sales/2017/january-backup.avro',
- google_cloud_storage_conn_id=google_cloud_conn_id
- )
-
- The following Operator would copy all the Avro files from ``sales/sales-2017``
- folder (i.e. with names starting with that prefix) in ``data`` bucket to the
- ``copied_sales/2017`` folder in the ``data_backup`` bucket. ::
- copy_files = GoogleCloudStorageToGoogleCloudStorageOperator(
- task_id='copy_files',
- source_bucket='data',
- source_object='sales/sales-2017/*.avro',
- destination_bucket='data_backup',
- destination_object='copied_sales/2017/',
- google_cloud_storage_conn_id=google_cloud_conn_id
- )
-
- The following Operator would move all the Avro files from ``sales/sales-2017``
- folder (i.e. with names starting with that prefix) in ``data`` bucket to the
- same folder in the ``data_backup`` bucket, deleting the original files in the
- process. ::
- move_files = GoogleCloudStorageToGoogleCloudStorageOperator(
- task_id='move_files',
- source_bucket='data',
- source_object='sales/sales-2017/*.avro',
- destination_bucket='data_backup',
- move_object=True,
- google_cloud_storage_conn_id=google_cloud_conn_id
- )
- """
- template_fields = ('source_bucket', 'source_object', 'destination_bucket',
- 'destination_object',)
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- source_bucket,
- source_object,
- destination_bucket=None,
- destination_object=None,
- move_object=False,
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(GoogleCloudStorageToGoogleCloudStorageOperator,
- self).__init__(*args, **kwargs)
- self.source_bucket = source_bucket
- self.source_object = source_object
- self.destination_bucket = destination_bucket
- self.destination_object = destination_object
- self.move_object = move_object
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
- self.wildcard = '*'
-
- def execute(self, context):
-
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to
- )
- log_message = 'Executing copy of gs://{0}/{1} to gs://{2}/{3}'
-
- if self.wildcard in self.source_object:
- prefix, delimiter = self.source_object.split(self.wildcard, 1)
- objects = hook.list(self.source_bucket, prefix=prefix, delimiter=delimiter)
-
- for source_object in objects:
- if self.destination_object is None:
- destination_object = source_object
- else:
- destination_object = source_object.replace(prefix,
- self.destination_object, 1)
- self.log.info(
- log_message.format(self.source_bucket, source_object,
- self.destination_bucket, destination_object)
- )
-
- hook.rewrite(self.source_bucket, source_object,
- self.destination_bucket, destination_object)
- if self.move_object:
- hook.delete(self.source_bucket, source_object)
-
- else:
- self.log.info(
- log_message.format(self.source_bucket, self.source_object,
- self.destination_bucket or self.source_bucket,
- self.destination_object or self.source_object)
- )
- hook.rewrite(self.source_bucket, self.source_object,
- self.destination_bucket, self.destination_object)
-
- if self.move_object:
- hook.delete(self.source_bucket, self.source_object)
diff --git a/airflow/contrib/operators/gcs_to_s3.py b/airflow/contrib/operators/gcs_to_s3.py
deleted file mode 100644
index a87aa3af5c..0000000000
--- a/airflow/contrib/operators/gcs_to_s3.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.contrib.operators.gcs_list_operator import GoogleCloudStorageListOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.hooks.S3_hook import S3Hook
-
-
-class GoogleCloudStorageToS3Operator(GoogleCloudStorageListOperator):
- """
- Synchronizes a Google Cloud Storage bucket with an S3 bucket.
-
- :param bucket: The Google Cloud Storage bucket to find the objects. (templated)
- :type bucket: string
- :param prefix: Prefix string which filters objects whose name begin with
- this prefix. (templated)
- :type prefix: string
- :param delimiter: The delimiter by which you want to filter the objects. (templated)
- For e.g to lists the CSV files from in a directory in GCS you would use
- delimiter='.csv'.
- :type delimiter: string
- :param google_cloud_storage_conn_id: The connection ID to use when
- connecting to Google Cloud Storage.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- :param dest_aws_conn_id: The destination S3 connection
- :type dest_aws_conn_id: str
- :param dest_s3_key: The base S3 key to be used to store the files. (templated)
- :type dest_s3_key: str
- """
- template_fields = ('bucket', 'prefix', 'delimiter', 'dest_s3_key')
- ui_color = '#f0eee4'
-
- @apply_defaults
- def __init__(self,
- bucket,
- prefix=None,
- delimiter=None,
- google_cloud_storage_conn_id='google_cloud_storage_default',
- delegate_to=None,
- dest_aws_conn_id=None,
- dest_s3_key=None,
- replace=False,
- *args,
- **kwargs):
-
- super(GoogleCloudStorageToS3Operator, self).__init__(
- bucket=bucket,
- prefix=prefix,
- delimiter=delimiter,
- google_cloud_storage_conn_id=google_cloud_storage_conn_id,
- delegate_to=delegate_to,
- *args,
- **kwargs
- )
- self.dest_aws_conn_id = dest_aws_conn_id
- self.dest_s3_key = dest_s3_key
- self.replace = replace
-
- def execute(self, context):
- # use the super to list all files in an Google Cloud Storage bucket
- files = super(GoogleCloudStorageToS3Operator, self).execute(context)
- s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id)
-
- if not self.replace:
- # if we are not replacing -> list all files in the S3 bucket
- # and only keep those files which are present in
- # Google Cloud Storage and not in S3
- bucket_name, _ = S3Hook.parse_s3_url(self.dest_s3_key)
- existing_files = s3_hook.list_keys(bucket_name)
- files = set(files) - set(existing_files)
-
- if files:
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to
- )
-
- for file in files:
- file_bytes = hook.download(self.bucket, file)
-
- dest_key = self.dest_s3_key + file
- self.log.info("Saving file to %s", dest_key)
-
- s3_hook.load_bytes(file_bytes,
- key=dest_key,
- replace=self.replace)
-
- self.log.info("All done, uploaded %d files to S3", len(files))
- else:
- self.log.info("In sync, no files needed to be uploaded to S3")
-
- return files
diff --git a/airflow/contrib/operators/hipchat_operator.py b/airflow/contrib/operators/hipchat_operator.py
deleted file mode 100644
index 381cd72cdf..0000000000
--- a/airflow/contrib/operators/hipchat_operator.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 builtins import str
-
-from airflow.utils.decorators import apply_defaults
-from airflow.models import BaseOperator
-from airflow.exceptions import AirflowException
-import requests
-import json
-
-
-class HipChatAPIOperator(BaseOperator):
- """
- Base HipChat Operator.
- All derived HipChat operators reference from HipChat's official REST API documentation
- at https://www.hipchat.com/docs/apiv2. Before using any HipChat API operators you need
- to get an authentication token at https://www.hipchat.com/docs/apiv2/auth.
- In the future additional HipChat operators will be derived from this class as well.
-
- :param token: HipChat REST API authentication token
- :type token: str
- :param base_url: HipChat REST API base url.
- :type base_url: str
- """
- @apply_defaults
- def __init__(self,
- token,
- base_url='https://api.hipchat.com/v2',
- *args,
- **kwargs):
- super(HipChatAPIOperator, self).__init__(*args, **kwargs)
- self.token = token
- self.base_url = base_url
- self.method = None
- self.url = None
- self.body = None
-
- def prepare_request(self):
- """
- Used by the execute function. Set the request method, url, and body of HipChat's
- REST API call.
- Override in child class. Each HipChatAPI child operator is responsible for having
- a prepare_request method call which sets self.method, self.url, and self.body.
- """
- pass
-
- def execute(self, context):
- self.prepare_request()
-
- response = requests.request(self.method,
- self.url,
- headers={
- 'Content-Type': 'application/json',
- 'Authorization': 'Bearer %s' % self.token},
- data=self.body)
- if response.status_code >= 400:
- self.log.error('HipChat API call failed: %s %s',
- response.status_code, response.reason)
- raise AirflowException('HipChat API call failed: %s %s' %
- (response.status_code, response.reason))
-
-
-class HipChatAPISendRoomNotificationOperator(HipChatAPIOperator):
- """
- Send notification to a specific HipChat room.
- More info: https://www.hipchat.com/docs/apiv2/method/send_room_notification
-
- :param room_id: Room in which to send notification on HipChat. (templated)
- :type room_id: str
- :param message: The message body. (templated)
- :type message: str
- :param frm: Label to be shown in addition to sender's name
- :type frm: str
- :param message_format: How the notification is rendered: html or text
- :type message_format: str
- :param color: Background color of the msg: yellow, green, red, purple, gray, or random
- :type color: str
- :param attach_to: The message id to attach this notification to
- :type attach_to: str
- :param notify: Whether this message should trigger a user notification
- :type notify: bool
- :param card: HipChat-defined card object
- :type card: dict
- """
- template_fields = ('token', 'room_id', 'message')
- ui_color = '#2980b9'
-
- @apply_defaults
- def __init__(self, room_id, message, *args, **kwargs):
- super(HipChatAPISendRoomNotificationOperator, self).__init__(*args, **kwargs)
- self.room_id = room_id
- self.message = message
- default_options = {
- 'message_format': 'html',
- 'color': 'yellow',
- 'frm': 'airflow',
- 'attach_to': None,
- 'notify': False,
- 'card': None
- }
- for (prop, default) in default_options.items():
- setattr(self, prop, kwargs.get(prop, default))
-
- def prepare_request(self):
- params = {
- 'message': self.message,
- 'message_format': self.message_format,
- 'color': self.color,
- 'from': self.frm,
- 'attach_to': self.attach_to,
- 'notify': self.notify,
- 'card': self.card
- }
-
- self.method = 'POST'
- self.url = '%s/room/%s/notification' % (self.base_url, self.room_id)
- self.body = json.dumps(dict(
- (str(k), str(v)) for k, v in params.items() if v))
diff --git a/airflow/contrib/operators/hive_to_dynamodb.py b/airflow/contrib/operators/hive_to_dynamodb.py
deleted file mode 100644
index 4a39e40741..0000000000
--- a/airflow/contrib/operators/hive_to_dynamodb.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 json
-
-from airflow.contrib.hooks.aws_dynamodb_hook import AwsDynamoDBHook
-from airflow.hooks.hive_hooks import HiveServer2Hook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class HiveToDynamoDBTransferOperator(BaseOperator):
- """
- Moves data from Hive to DynamoDB, note that for now the data is loaded
- into memory before being pushed to DynamoDB, so this operator should
- be used for smallish amount of data.
-
- :param sql: SQL query to execute against the hive database. (templated)
- :type sql: str
- :param table_name: target DynamoDB table
- :type table_name: str
- :param table_keys: partition key and sort key
- :type table_keys: list
- :param pre_process: implement pre-processing of source data
- :type pre_process: function
- :param pre_process_args: list of pre_process function arguments
- :type pre_process_args: list
- :param pre_process_kwargs: dict of pre_process function arguments
- :type pre_process_kwargs: dict
- :param region_name: aws region name (example: us-east-1)
- :type region_name: str
- :param schema: hive database schema
- :type schema: str
- :param hiveserver2_conn_id: source hive connection
- :type hiveserver2_conn_id: str
- :param aws_conn_id: aws connection
- :type aws_conn_id: str
- """
-
- template_fields = ('sql',)
- template_ext = ('.sql',)
- ui_color = '#a0e08c'
-
- @apply_defaults
- def __init__(
- self,
- sql,
- table_name,
- table_keys,
- pre_process=None,
- pre_process_args=None,
- pre_process_kwargs=None,
- region_name=None,
- schema='default',
- hiveserver2_conn_id='hiveserver2_default',
- aws_conn_id='aws_default',
- *args, **kwargs):
- super(HiveToDynamoDBTransferOperator, self).__init__(*args, **kwargs)
- self.sql = sql
- self.table_name = table_name
- self.table_keys = table_keys
- self.pre_process = pre_process
- self.pre_process_args = pre_process_args
- self.pre_process_kwargs = pre_process_kwargs
- self.region_name = region_name
- self.schema = schema
- self.hiveserver2_conn_id = hiveserver2_conn_id
- self.aws_conn_id = aws_conn_id
-
- def execute(self, context):
- hive = HiveServer2Hook(hiveserver2_conn_id=self.hiveserver2_conn_id)
-
- self.log.info('Extracting data from Hive')
- self.log.info(self.sql)
-
- data = hive.get_pandas_df(self.sql, schema=self.schema)
- dynamodb = AwsDynamoDBHook(aws_conn_id=self.aws_conn_id,
- table_name=self.table_name,
- table_keys=self.table_keys,
- region_name=self.region_name)
-
- self.log.info('Inserting rows into dynamodb')
-
- if self.pre_process is None:
- dynamodb.write_batch_data(
- json.loads(data.to_json(orient='records')))
- else:
- dynamodb.write_batch_data(
- self.pre_process(data=data,
- args=self.pre_process_args,
- kwargs=self.pre_process_kwargs))
-
- self.log.info('Done.')
diff --git a/airflow/contrib/operators/jenkins_job_trigger_operator.py b/airflow/contrib/operators/jenkins_job_trigger_operator.py
deleted file mode 100644
index 37185f3e13..0000000000
--- a/airflow/contrib/operators/jenkins_job_trigger_operator.py
+++ /dev/null
@@ -1,249 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 time
-import socket
-import json
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.contrib.hooks.jenkins_hook import JenkinsHook
-import jenkins
-from jenkins import JenkinsException
-from six.moves.urllib.request import Request, urlopen
-from six.moves.urllib.error import HTTPError, URLError
-
-try:
- basestring
-except NameError:
- basestring = str # For python3 compatibility
-
-
-# TODO Use jenkins_urlopen instead when it will be available
-# in the stable python-jenkins version (> 0.4.15)
-def jenkins_request_with_headers(jenkins_server, req, add_crumb=True):
- """
- We need to get the headers in addition to the body answer
- to get the location from them
- This function is just a copy of the one present in python-jenkins library
- with just the return call changed
- :param jenkins_server: The server to query
- :param req: The request to execute
- :param add_crumb: Boolean to indicate if it should add crumb to the request
- :return:
- """
- try:
- if jenkins_server.auth:
- req.add_header('Authorization', jenkins_server.auth)
- if add_crumb:
- jenkins_server.maybe_add_crumb(req)
- response = urlopen(req, timeout=jenkins_server.timeout)
- response_body = response.read()
- response_headers = response.info()
- if response_body is None:
- raise jenkins.EmptyResponseException(
- "Error communicating with server[%s]: "
- "empty response" % jenkins_server.server)
- return {'body': response_body.decode('utf-8'), 'headers': response_headers}
- except HTTPError as e:
- # Jenkins's funky authentication means its nigh impossible to
- # distinguish errors.
- if e.code in [401, 403, 500]:
- # six.moves.urllib.error.HTTPError provides a 'reason'
- # attribute for all python version except for ver 2.6
- # Falling back to HTTPError.msg since it contains the
- # same info as reason
- raise JenkinsException(
- 'Error in request. ' +
- 'Possibly authentication failed [%s]: %s' % (
- e.code, e.msg)
- )
- elif e.code == 404:
- raise jenkins.NotFoundException('Requested item could not be found')
- else:
- raise
- except socket.timeout as e:
- raise jenkins.TimeoutException('Error in request: %s' % e)
- except URLError as e:
- # python 2.6 compatibility to ensure same exception raised
- # since URLError wraps a socket timeout on python 2.6.
- if str(e.reason) == "timed out":
- raise jenkins.TimeoutException('Error in request: %s' % e.reason)
- raise JenkinsException('Error in request: %s' % e.reason)
-
-
-class JenkinsJobTriggerOperator(BaseOperator):
- """
- Trigger a Jenkins Job and monitor it's execution.
- This operator depend on python-jenkins library,
- version >= 0.4.15 to communicate with jenkins server.
- You'll also need to configure a Jenkins connection in the connections screen.
- :param jenkins_connection_id: The jenkins connection to use for this job
- :type jenkins_connection_id: string
- :param job_name: The name of the job to trigger
- :type job_name: string
- :param parameters: The parameters block to provide to jenkins. (templated)
- :type parameters: string
- :param sleep_time: How long will the operator sleep between each status
- request for the job (min 1, default 10)
- :type sleep_time: int
- :param max_try_before_job_appears: The maximum number of requests to make
- while waiting for the job to appears on jenkins server (default 10)
- :type max_try_before_job_appears: int
- """
- template_fields = ('parameters',)
- template_ext = ('.json',)
- ui_color = '#f9ec86'
-
- @apply_defaults
- def __init__(self,
- jenkins_connection_id,
- job_name,
- parameters="",
- sleep_time=10,
- max_try_before_job_appears=10,
- *args,
- **kwargs):
- super(JenkinsJobTriggerOperator, self).__init__(*args, **kwargs)
- self.job_name = job_name
- self.parameters = parameters
- if sleep_time < 1:
- sleep_time = 1
- self.sleep_time = sleep_time
- self.jenkins_connection_id = jenkins_connection_id
- self.max_try_before_job_appears = max_try_before_job_appears
-
- def build_job(self, jenkins_server):
- """
- This function makes an API call to Jenkins to trigger a build for 'job_name'
- It returned a dict with 2 keys : body and headers.
- headers contains also a dict-like object which can be queried to get
- the location to poll in the queue.
- :param jenkins_server: The jenkins server where the job should be triggered
- :return: Dict containing the response body (key body)
- and the headers coming along (headers)
- """
- # Warning if the parameter is too long, the URL can be longer than
- # the maximum allowed size
- if self.parameters and isinstance(self.parameters, basestring):
- import ast
- self.parameters = ast.literal_eval(self.parameters)
-
- if not self.parameters:
- # We need a None to call the non parametrized jenkins api end point
- self.parameters = None
-
- request = Request(jenkins_server.build_job_url(self.job_name,
- self.parameters, None), b'')
- return jenkins_request_with_headers(jenkins_server, request)
-
- def poll_job_in_queue(self, location, jenkins_server):
- """
- This method poll the jenkins queue until the job is executed.
- When we trigger a job through an API call,
- the job is first put in the queue without having a build number assigned.
- Thus we have to wait the job exit the queue to know its build number.
- To do so, we have to add /api/json (or /api/xml) to the location
- returned by the build_job call and poll this file.
- When a 'executable' block appears in the json, it means the job execution started
- and the field 'number' then contains the build number.
- :param location: Location to poll, returned in the header of the build_job call
- :param jenkins_server: The jenkins server to poll
- :return: The build_number corresponding to the triggered job
- """
- try_count = 0
- location = location + '/api/json'
- # TODO Use get_queue_info instead
- # once it will be available in python-jenkins (v > 0.4.15)
- self.log.info('Polling jenkins queue at the url %s', location)
- while try_count < self.max_try_before_job_appears:
- location_answer = jenkins_request_with_headers(jenkins_server,
- Request(location))
- if location_answer is not None:
- json_response = json.loads(location_answer['body'])
- if 'executable' in json_response:
- build_number = json_response['executable']['number']
- self.log.info('Job executed on Jenkins side with the build number %s',
- build_number)
- return build_number
- try_count += 1
- time.sleep(self.sleep_time)
- raise AirflowException("The job hasn't been executed"
- " after polling the queue %d times",
- self.max_try_before_job_appears)
-
- def get_hook(self):
- return JenkinsHook(self.jenkins_connection_id)
-
- def execute(self, context):
- if not self.jenkins_connection_id:
- self.log.error(
- 'Please specify the jenkins connection id to use.'
- 'You must create a Jenkins connection before'
- ' being able to use this operator')
- raise AirflowException('The jenkins_connection_id parameter is missing,'
- 'impossible to trigger the job')
-
- if not self.job_name:
- self.log.error("Please specify the job name to use in the job_name parameter")
- raise AirflowException('The job_name parameter is missing,'
- 'impossible to trigger the job')
-
- self.log.info(
- 'Triggering the job %s on the jenkins : %s with the parameters : %s',
- self.job_name, self.jenkins_connection_id, self.parameters)
- jenkins_server = self.get_hook().get_jenkins_server()
- jenkins_response = self.build_job(jenkins_server)
- build_number = self.poll_job_in_queue(
- jenkins_response['headers']['Location'], jenkins_server)
-
- time.sleep(self.sleep_time)
- keep_polling_job = True
- build_info = None
- while keep_polling_job:
- try:
- build_info = jenkins_server.get_build_info(name=self.job_name,
- number=build_number)
- if build_info['result'] is not None:
- keep_polling_job = False
- # Check if job had errors.
- if build_info['result'] != 'SUCCESS':
- raise AirflowException(
- 'Jenkins job failed, final state : %s.'
- 'Find more information on job url : %s'
- % (build_info['result'], build_info['url']))
- else:
- self.log.info('Waiting for job to complete : %s , build %s',
- self.job_name, build_number)
- time.sleep(self.sleep_time)
- except jenkins.NotFoundException as err:
- raise AirflowException(
- 'Jenkins job status check failed. Final error was: %s'
- % err.resp.status)
- except jenkins.JenkinsException as err:
- raise AirflowException(
- 'Jenkins call failed with error : %s, if you have parameters '
- 'double check them, jenkins sends back '
- 'this exception for unknown parameters'
- 'You can also check logs for more details on this exception '
- '(jenkins_url/log/rss)', str(err))
- if build_info:
- # If we can we return the url of the job
- # for later use (like retrieving an artifact)
- return build_info['url']
diff --git a/airflow/contrib/operators/jira_operator.py b/airflow/contrib/operators/jira_operator.py
deleted file mode 100644
index 01d78b8645..0000000000
--- a/airflow/contrib/operators/jira_operator.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.jira_hook import JIRAError
-from airflow.contrib.hooks.jira_hook import JiraHook
-from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class JiraOperator(BaseOperator):
- """
- JiraOperator to interact and perform action on Jira issue tracking system.
- This operator is designed to use Jira Python SDK: http://jira.readthedocs.io
-
- :param jira_conn_id: reference to a pre-defined Jira Connection
- :type jira_conn_id: str
- :param jira_method: method name from Jira Python SDK to be called
- :type jira_method: str
- :param jira_method_args: required method parameters for the jira_method. (templated)
- :type jira_method_args: dict
- :param result_processor: function to further process the response from Jira
- :type result_processor: function
- :param get_jira_resource_method: function or operator to get jira resource
- on which the provided jira_method will be executed
- :type get_jira_resource_method: function
- """
-
- template_fields = ("jira_method_args",)
-
- @apply_defaults
- def __init__(self,
- jira_conn_id='jira_default',
- jira_method=None,
- jira_method_args=None,
- result_processor=None,
- get_jira_resource_method=None,
- *args,
- **kwargs):
- super(JiraOperator, self).__init__(*args, **kwargs)
- self.jira_conn_id = jira_conn_id
- self.method_name = jira_method
- self.jira_method_args = jira_method_args
- self.result_processor = result_processor
- self.get_jira_resource_method = get_jira_resource_method
-
- def execute(self, context):
- try:
- if self.get_jira_resource_method is not None:
- # if get_jira_resource_method is provided, jira_method will be executed on
- # resource returned by executing the get_jira_resource_method.
- # This makes all the provided methods of JIRA sdk accessible and usable
- # directly at the JiraOperator without additional wrappers.
- # ref: http://jira.readthedocs.io/en/latest/api.html
- if isinstance(self.get_jira_resource_method, JiraOperator):
- resource = self.get_jira_resource_method.execute(**context)
- else:
- resource = self.get_jira_resource_method(**context)
- else:
- # Default method execution is on the top level jira client resource
- hook = JiraHook(jira_conn_id=self.jira_conn_id)
- resource = hook.client
-
- # Current Jira-Python SDK (1.0.7) has issue with pickling the jira response.
- # ex: self.xcom_push(context, key='operator_response', value=jira_response)
- # This could potentially throw error if jira_result is not picklable
- jira_result = getattr(resource, self.method_name)(**self.jira_method_args)
- if self.result_processor:
- return self.result_processor(context, jira_result)
-
- return jira_result
-
- except JIRAError as jira_error:
- raise AirflowException("Failed to execute jiraOperator, error: %s"
- % str(jira_error))
- except Exception as e:
- raise AirflowException("Jira operator error: %s" % str(e))
diff --git a/airflow/contrib/operators/kubernetes_pod_operator.py b/airflow/contrib/operators/kubernetes_pod_operator.py
deleted file mode 100644
index fb905622d8..0000000000
--- a/airflow/contrib/operators/kubernetes_pod_operator.py
+++ /dev/null
@@ -1,174 +0,0 @@
-# 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 airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.contrib.kubernetes import kube_client, pod_generator, pod_launcher
-from airflow.contrib.kubernetes.pod import Resources
-from airflow.utils.state import State
-from airflow.contrib.kubernetes.volume_mount import VolumeMount # noqa
-from airflow.contrib.kubernetes.volume import Volume # noqa
-from airflow.contrib.kubernetes.secret import Secret # noqa
-
-template_fields = ('templates_dict',)
-template_ext = tuple()
-ui_color = '#ffefeb'
-
-
-class KubernetesPodOperator(BaseOperator):
- """
- Execute a task in a Kubernetes Pod
-
- :param image: Docker image you wish to launch. Defaults to dockerhub.io,
- but fully qualified URLS will point to custom repositories
- :type image: str
- :param: namespace: the namespace to run within kubernetes
- :type: namespace: str
- :param cmds: entrypoint of the container. (templated)
- The docker images's entrypoint is used if this is not provide.
- :type cmds: list of str
- :param arguments: arguments of to the entrypoint. (templated)
- The docker image's CMD is used if this is not provided.
- :type arguments: list of str
- :param volume_mounts: volumeMounts for launched pod
- :type volume_mounts: list of VolumeMount
- :param volumes: volumes for launched pod. Includes ConfigMaps and PersistentVolumes
- :type volumes: list of Volume
- :param labels: labels to apply to the Pod
- :type labels: dict
- :param startup_timeout_seconds: timeout in seconds to startup the pod
- :type startup_timeout_seconds: int
- :param name: name of the task you want to run,
- will be used to generate a pod id
- :type name: str
- :param env_vars: Environment variables initialized in the container. (templated)
- :type env_vars: dict
- :param secrets: Kubernetes secrets to inject in the container,
- They can be exposed as environment vars or files in a volume.
- :type secrets: list of Secret
- :param in_cluster: run kubernetes client with in_cluster configuration
- :type in_cluster: bool
- :param cluster_context: context that points to kubernetes cluster.
- Ignored when in_cluster is True. If None, current-context is used.
- :type cluster_context: string
- :param get_logs: get the stdout of the container as logs of the tasks
- :type get_logs: bool
- :param affinity: A dict containing a group of affinity scheduling rules
- :type affinity: dict
- :param node_selectors: A dict containing a group of scheduling rules
- :type node_selectors: dict
- :param config_file: The path to the Kubernetes config file
- :type config_file: str
- :param xcom_push: If xcom_push is True, the content of the file
- /airflow/xcom/return.json in the container will also be pushed to an
- XCom when the container completes.
- :type xcom_push: bool
- """
- template_fields = ('cmds', 'arguments', 'env_vars', 'config_file')
-
- def execute(self, context):
- try:
- client = kube_client.get_kube_client(in_cluster=self.in_cluster,
- cluster_context=self.cluster_context,
- config_file=self.config_file)
- gen = pod_generator.PodGenerator()
-
- for mount in self.volume_mounts:
- gen.add_mount(mount)
- for volume in self.volumes:
- gen.add_volume(volume)
-
- pod = gen.make_pod(
- namespace=self.namespace,
- image=self.image,
- pod_id=self.name,
- cmds=self.cmds,
- arguments=self.arguments,
- labels=self.labels,
- )
-
- pod.secrets = self.secrets
- pod.envs = self.env_vars
- pod.image_pull_policy = self.image_pull_policy
- pod.annotations = self.annotations
- pod.resources = self.resources
- pod.affinity = self.affinity
- pod.node_selectors = self.node_selectors
-
- launcher = pod_launcher.PodLauncher(kube_client=client,
- extract_xcom=self.xcom_push)
- (final_state, result) = launcher.run_pod(
- pod,
- startup_timeout=self.startup_timeout_seconds,
- get_logs=self.get_logs)
- if final_state != State.SUCCESS:
- raise AirflowException(
- 'Pod returned a failure: {state}'.format(state=final_state)
- )
- if self.xcom_push:
- return result
- except AirflowException as ex:
- raise AirflowException('Pod Launching failed: {error}'.format(error=ex))
-
- @apply_defaults
- def __init__(self,
- namespace,
- image,
- name,
- cmds=None,
- arguments=None,
- volume_mounts=None,
- volumes=None,
- env_vars=None,
- secrets=None,
- in_cluster=False,
- cluster_context=None,
- labels=None,
- startup_timeout_seconds=120,
- get_logs=True,
- image_pull_policy='IfNotPresent',
- annotations=None,
- resources=None,
- affinity=None,
- config_file=None,
- xcom_push=False,
- node_selectors=None,
- *args,
- **kwargs):
- super(KubernetesPodOperator, self).__init__(*args, **kwargs)
- self.image = image
- self.namespace = namespace
- self.cmds = cmds or []
- self.arguments = arguments or []
- self.labels = labels or {}
- self.startup_timeout_seconds = startup_timeout_seconds
- self.name = name
- self.env_vars = env_vars or {}
- self.volume_mounts = volume_mounts or []
- self.volumes = volumes or []
- self.secrets = secrets or []
- self.in_cluster = in_cluster
- self.cluster_context = cluster_context
- self.get_logs = get_logs
- self.image_pull_policy = image_pull_policy
- self.node_selectors = node_selectors or {}
- self.annotations = annotations or {}
- self.affinity = affinity or {}
- self.xcom_push = xcom_push
- self.resources = resources or Resources()
- self.config_file = config_file
diff --git a/airflow/contrib/operators/mlengine_operator.py b/airflow/contrib/operators/mlengine_operator.py
deleted file mode 100644
index 9fe966d387..0000000000
--- a/airflow/contrib/operators/mlengine_operator.py
+++ /dev/null
@@ -1,611 +0,0 @@
-#
-# 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 re
-
-from apiclient import errors
-
-from airflow.contrib.hooks.gcp_mlengine_hook import MLEngineHook
-from airflow.exceptions import AirflowException
-from airflow.operators import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-log = LoggingMixin().log
-
-
-def _normalize_mlengine_job_id(job_id):
- """
- Replaces invalid MLEngine job_id characters with '_'.
-
- This also adds a leading 'z' in case job_id starts with an invalid
- character.
-
- Args:
- job_id: A job_id str that may have invalid characters.
-
- Returns:
- A valid job_id representation.
- """
-
- # Add a prefix when a job_id starts with a digit or a template
- match = re.search(r'\d|\{{2}', job_id)
- if match and match.start() is 0:
- job = 'z_{}'.format(job_id)
- else:
- job = job_id
-
- # Clean up 'bad' characters except templates
- tracker = 0
- cleansed_job_id = ''
- for m in re.finditer(r'\{{2}.+?\}{2}', job):
- cleansed_job_id += re.sub(r'[^0-9a-zA-Z]+', '_',
- job[tracker:m.start()])
- cleansed_job_id += job[m.start():m.end()]
- tracker = m.end()
-
- # Clean up last substring or the full string if no templates
- cleansed_job_id += re.sub(r'[^0-9a-zA-Z]+', '_', job[tracker:])
-
- return cleansed_job_id
-
-
-class MLEngineBatchPredictionOperator(BaseOperator):
- """
- Start a Google Cloud ML Engine prediction job.
-
- NOTE: For model origin, users should consider exactly one from the
- three options below:
- 1. Populate 'uri' field only, which should be a GCS location that
- points to a tensorflow savedModel directory.
- 2. Populate 'model_name' field only, which refers to an existing
- model, and the default version of the model will be used.
- 3. Populate both 'model_name' and 'version_name' fields, which
- refers to a specific version of a specific model.
-
- In options 2 and 3, both model and version name should contain the
- minimal identifier. For instance, call
-
- ::
-
- MLEngineBatchPredictionOperator(
- ...,
- model_name='my_model',
- version_name='my_version',
- ...)
-
- if the desired model version is
- "projects/my_project/models/my_model/versions/my_version".
-
- See https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs
- for further documentation on the parameters.
-
- :param project_id: The Google Cloud project name where the
- prediction job is submitted. (templated)
- :type project_id: string
-
- :param job_id: A unique id for the prediction job on Google Cloud
- ML Engine. (templated)
- :type job_id: string
-
- :param data_format: The format of the input data.
- It will default to 'DATA_FORMAT_UNSPECIFIED' if is not provided
- or is not one of ["TEXT", "TF_RECORD", "TF_RECORD_GZIP"].
- :type data_format: string
-
- :param input_paths: A list of GCS paths of input data for batch
- prediction. Accepting wildcard operator *, but only at the end. (templated)
- :type input_paths: list of string
-
- :param output_path: The GCS path where the prediction results are
- written to. (templated)
- :type output_path: string
-
- :param region: The Google Compute Engine region to run the
- prediction job in. (templated)
- :type region: string
-
- :param model_name: The Google Cloud ML Engine model to use for prediction.
- If version_name is not provided, the default version of this
- model will be used.
- Should not be None if version_name is provided.
- Should be None if uri is provided. (templated)
- :type model_name: string
-
- :param version_name: The Google Cloud ML Engine model version to use for
- prediction.
- Should be None if uri is provided. (templated)
- :type version_name: string
-
- :param uri: The GCS path of the saved model to use for prediction.
- Should be None if model_name is provided.
- It should be a GCS path pointing to a tensorflow SavedModel. (templated)
- :type uri: string
-
- :param max_worker_count: The maximum number of workers to be used
- for parallel processing. Defaults to 10 if not specified.
- :type max_worker_count: int
-
- :param runtime_version: The Google Cloud ML Engine runtime version to use
- for batch prediction.
- :type runtime_version: string
-
- :param gcp_conn_id: The connection ID used for connection to Google
- Cloud Platform.
- :type gcp_conn_id: string
-
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must
- have doamin-wide delegation enabled.
- :type delegate_to: string
-
- Raises:
- ``ValueError``: if a unique model/version origin cannot be determined.
- """
-
- template_fields = [
- '_project_id',
- '_job_id',
- '_region',
- '_input_paths',
- '_output_path',
- '_model_name',
- '_version_name',
- '_uri',
- ]
-
- @apply_defaults
- def __init__(self,
- project_id,
- job_id,
- region,
- data_format,
- input_paths,
- output_path,
- model_name=None,
- version_name=None,
- uri=None,
- max_worker_count=None,
- runtime_version=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(MLEngineBatchPredictionOperator, self).__init__(*args, **kwargs)
-
- self._project_id = project_id
- self._job_id = job_id
- self._region = region
- self._data_format = data_format
- self._input_paths = input_paths
- self._output_path = output_path
- self._model_name = model_name
- self._version_name = version_name
- self._uri = uri
- self._max_worker_count = max_worker_count
- self._runtime_version = runtime_version
- self._gcp_conn_id = gcp_conn_id
- self._delegate_to = delegate_to
-
- if not self._project_id:
- raise AirflowException('Google Cloud project id is required.')
- if not self._job_id:
- raise AirflowException(
- 'An unique job id is required for Google MLEngine prediction '
- 'job.')
-
- if self._uri:
- if self._model_name or self._version_name:
- raise AirflowException('Ambiguous model origin: Both uri and '
- 'model/version name are provided.')
-
- if self._version_name and not self._model_name:
- raise AirflowException(
- 'Missing model: Batch prediction expects '
- 'a model name when a version name is provided.')
-
- if not (self._uri or self._model_name):
- raise AirflowException(
- 'Missing model origin: Batch prediction expects a model, '
- 'a model & version combination, or a URI to a savedModel.')
-
- def execute(self, context):
- job_id = _normalize_mlengine_job_id(self._job_id)
- prediction_request = {
- 'jobId': job_id,
- 'predictionInput': {
- 'dataFormat': self._data_format,
- 'inputPaths': self._input_paths,
- 'outputPath': self._output_path,
- 'region': self._region
- }
- }
-
- if self._uri:
- prediction_request['predictionInput']['uri'] = self._uri
- elif self._model_name:
- origin_name = 'projects/{}/models/{}'.format(
- self._project_id, self._model_name)
- if not self._version_name:
- prediction_request['predictionInput'][
- 'modelName'] = origin_name
- else:
- prediction_request['predictionInput']['versionName'] = \
- origin_name + '/versions/{}'.format(self._version_name)
-
- if self._max_worker_count:
- prediction_request['predictionInput'][
- 'maxWorkerCount'] = self._max_worker_count
-
- if self._runtime_version:
- prediction_request['predictionInput'][
- 'runtimeVersion'] = self._runtime_version
-
- hook = MLEngineHook(self._gcp_conn_id, self._delegate_to)
-
- # Helper method to check if the existing job's prediction input is the
- # same as the request we get here.
- def check_existing_job(existing_job):
- return existing_job.get('predictionInput', None) == \
- prediction_request['predictionInput']
-
- try:
- finished_prediction_job = hook.create_job(
- self._project_id, prediction_request, check_existing_job)
- except errors.HttpError:
- raise
-
- if finished_prediction_job['state'] != 'SUCCEEDED':
- self.log.error('MLEngine batch prediction job failed: {}'.format(
- str(finished_prediction_job)))
- raise RuntimeError(finished_prediction_job['errorMessage'])
-
- return finished_prediction_job['predictionOutput']
-
-
-class MLEngineModelOperator(BaseOperator):
- """
- Operator for managing a Google Cloud ML Engine model.
-
- :param project_id: The Google Cloud project name to which MLEngine
- model belongs. (templated)
- :type project_id: string
-
- :param model: A dictionary containing the information about the model.
- If the `operation` is `create`, then the `model` parameter should
- contain all the information about this model such as `name`.
-
- If the `operation` is `get`, the `model` parameter
- should contain the `name` of the model.
- :type model: dict
-
- :param operation: The operation to perform. Available operations are:
-
- * ``create``: Creates a new model as provided by the `model` parameter.
- * ``get``: Gets a particular model where the name is specified in `model`.
-
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :type gcp_conn_id: string
-
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- """
-
- template_fields = [
- '_model',
- ]
-
- @apply_defaults
- def __init__(self,
- project_id,
- model,
- operation='create',
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- super(MLEngineModelOperator, self).__init__(*args, **kwargs)
- self._project_id = project_id
- self._model = model
- self._operation = operation
- self._gcp_conn_id = gcp_conn_id
- self._delegate_to = delegate_to
-
- def execute(self, context):
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
- if self._operation == 'create':
- return hook.create_model(self._project_id, self._model)
- elif self._operation == 'get':
- return hook.get_model(self._project_id, self._model['name'])
- else:
- raise ValueError('Unknown operation: {}'.format(self._operation))
-
-
-class MLEngineVersionOperator(BaseOperator):
- """
- Operator for managing a Google Cloud ML Engine version.
-
- :param project_id: The Google Cloud project name to which MLEngine
- model belongs.
- :type project_id: string
-
- :param model_name: The name of the Google Cloud ML Engine model that the version
- belongs to. (templated)
- :type model_name: string
-
- :param version_name: A name to use for the version being operated upon.
- If not None and the `version` argument is None or does not have a value for
- the `name` key, then this will be populated in the payload for the
- `name` key. (templated)
- :type version_name: string
-
- :param version: A dictionary containing the information about the version.
- If the `operation` is `create`, `version` should contain all the
- information about this version such as name, and deploymentUrl.
- If the `operation` is `get` or `delete`, the `version` parameter
- should contain the `name` of the version.
- If it is None, the only `operation` possible would be `list`. (templated)
- :type version: dict
-
- :param operation: The operation to perform. Available operations are:
-
- * ``create``: Creates a new version in the model specified by `model_name`,
- in which case the `version` parameter should contain all the
- information to create that version
- (e.g. `name`, `deploymentUrl`).
-
- * ``get``: Gets full information of a particular version in the model
- specified by `model_name`.
- The name of the version should be specified in the `version`
- parameter.
-
- * ``list``: Lists all available versions of the model specified
- by `model_name`.
-
- * ``delete``: Deletes the version specified in `version` parameter from the
- model specified by `model_name`).
- The name of the version should be specified in the `version`
- parameter.
- :type operation: string
-
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :type gcp_conn_id: string
-
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- """
-
- template_fields = [
- '_model_name',
- '_version_name',
- '_version',
- ]
-
- @apply_defaults
- def __init__(self,
- project_id,
- model_name,
- version_name=None,
- version=None,
- operation='create',
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
-
- super(MLEngineVersionOperator, self).__init__(*args, **kwargs)
- self._project_id = project_id
- self._model_name = model_name
- self._version_name = version_name
- self._version = version or {}
- self._operation = operation
- self._gcp_conn_id = gcp_conn_id
- self._delegate_to = delegate_to
-
- def execute(self, context):
- if 'name' not in self._version:
- self._version['name'] = self._version_name
-
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
-
- if self._operation == 'create':
- assert self._version is not None
- return hook.create_version(self._project_id, self._model_name,
- self._version)
- elif self._operation == 'set_default':
- return hook.set_default_version(self._project_id, self._model_name,
- self._version['name'])
- elif self._operation == 'list':
- return hook.list_versions(self._project_id, self._model_name)
- elif self._operation == 'delete':
- return hook.delete_version(self._project_id, self._model_name,
- self._version['name'])
- else:
- raise ValueError('Unknown operation: {}'.format(self._operation))
-
-
-class MLEngineTrainingOperator(BaseOperator):
- """
- Operator for launching a MLEngine training job.
-
- :param project_id: The Google Cloud project name within which MLEngine
- training job should run (templated).
- :type project_id: string
-
- :param job_id: A unique templated id for the submitted Google MLEngine
- training job. (templated)
- :type job_id: string
-
- :param package_uris: A list of package locations for MLEngine training job,
- which should include the main training program + any additional
- dependencies. (templated)
- :type package_uris: string
-
- :param training_python_module: The Python module name to run within MLEngine
- training job after installing 'package_uris' packages. (templated)
- :type training_python_module: string
-
- :param training_args: A list of templated command line arguments to pass to
- the MLEngine training program. (templated)
- :type training_args: string
-
- :param region: The Google Compute Engine region to run the MLEngine training
- job in (templated).
- :type region: string
-
- :param scale_tier: Resource tier for MLEngine training job. (templated)
- :type scale_tier: string
-
- :param runtime_version: The Google Cloud ML runtime version to use for
- training. (templated)
- :type runtime_version: string
-
- :param python_version: The version of Python used in training. (templated)
- :type python_version: string
-
- :param job_dir: A Google Cloud Storage path in which to store training
- outputs and other data needed for training. (templated)
- :type job_dir: string
-
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :type gcp_conn_id: string
-
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
-
- :param mode: Can be one of 'DRY_RUN'/'CLOUD'. In 'DRY_RUN' mode, no real
- training job will be launched, but the MLEngine training job request
- will be printed out. In 'CLOUD' mode, a real MLEngine training job
- creation request will be issued.
- :type mode: string
- """
-
- template_fields = [
- '_project_id',
- '_job_id',
- '_package_uris',
- '_training_python_module',
- '_training_args',
- '_region',
- '_scale_tier',
- '_runtime_version',
- '_python_version',
- '_job_dir'
- ]
-
- @apply_defaults
- def __init__(self,
- project_id,
- job_id,
- package_uris,
- training_python_module,
- training_args,
- region,
- scale_tier=None,
- runtime_version=None,
- python_version=None,
- job_dir=None,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- mode='PRODUCTION',
- *args,
- **kwargs):
- super(MLEngineTrainingOperator, self).__init__(*args, **kwargs)
- self._project_id = project_id
- self._job_id = job_id
- self._package_uris = package_uris
- self._training_python_module = training_python_module
- self._training_args = training_args
- self._region = region
- self._scale_tier = scale_tier
- self._runtime_version = runtime_version
- self._python_version = python_version
- self._job_dir = job_dir
- self._gcp_conn_id = gcp_conn_id
- self._delegate_to = delegate_to
- self._mode = mode
-
- if not self._project_id:
- raise AirflowException('Google Cloud project id is required.')
- if not self._job_id:
- raise AirflowException(
- 'An unique job id is required for Google MLEngine training '
- 'job.')
- if not package_uris:
- raise AirflowException(
- 'At least one python package is required for MLEngine '
- 'Training job.')
- if not training_python_module:
- raise AirflowException(
- 'Python module name to run after installing required '
- 'packages is required.')
- if not self._region:
- raise AirflowException('Google Compute Engine region is required.')
-
- def execute(self, context):
- job_id = _normalize_mlengine_job_id(self._job_id)
- training_request = {
- 'jobId': job_id,
- 'trainingInput': {
- 'scaleTier': self._scale_tier,
- 'packageUris': self._package_uris,
- 'pythonModule': self._training_python_module,
- 'region': self._region,
- 'args': self._training_args,
- }
- }
-
- if self._runtime_version:
- training_request['trainingInput']['runtimeVersion'] = self._runtime_version
-
- if self._python_version:
- training_request['trainingInput']['pythonVersion'] = self._python_version
-
- if self._job_dir:
- training_request['trainingInput']['jobDir'] = self._job_dir
-
- if self._mode == 'DRY_RUN':
- self.log.info('In dry_run mode.')
- self.log.info('MLEngine Training job request is: {}'.format(
- training_request))
- return
-
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
-
- # Helper method to check if the existing job's training input is the
- # same as the request we get here.
- def check_existing_job(existing_job):
- return existing_job.get('trainingInput', None) == \
- training_request['trainingInput']
-
- try:
- finished_training_job = hook.create_job(
- self._project_id, training_request, check_existing_job)
- except errors.HttpError:
- raise
-
- if finished_training_job['state'] != 'SUCCEEDED':
- self.log.error('MLEngine training job failed: {}'.format(
- str(finished_training_job)))
- raise RuntimeError(finished_training_job['errorMessage'])
diff --git a/airflow/contrib/operators/mlengine_operator_utils.py b/airflow/contrib/operators/mlengine_operator_utils.py
deleted file mode 100644
index 7ce784ebb4..0000000000
--- a/airflow/contrib/operators/mlengine_operator_utils.py
+++ /dev/null
@@ -1,246 +0,0 @@
-#
-# 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 base64
-import json
-import os
-import re
-
-import dill
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.contrib.operators.mlengine_operator import MLEngineBatchPredictionOperator
-from airflow.contrib.operators.dataflow_operator import DataFlowPythonOperator
-from airflow.exceptions import AirflowException
-from airflow.operators.python_operator import PythonOperator
-from six.moves.urllib.parse import urlsplit
-
-
-def create_evaluate_ops(task_prefix,
- data_format,
- input_paths,
- prediction_path,
- metric_fn_and_keys,
- validate_fn,
- batch_prediction_job_id=None,
- project_id=None,
- region=None,
- dataflow_options=None,
- model_uri=None,
- model_name=None,
- version_name=None,
- dag=None):
- """
- Creates Operators needed for model evaluation and returns.
-
- It gets prediction over inputs via Cloud ML Engine BatchPrediction API by
- calling MLEngineBatchPredictionOperator, then summarize and validate
- the result via Cloud Dataflow using DataFlowPythonOperator.
-
- For details and pricing about Batch prediction, please refer to the website
- https://cloud.google.com/ml-engine/docs/how-tos/batch-predict
- and for Cloud Dataflow, https://cloud.google.com/dataflow/docs/
-
- It returns three chained operators for prediction, summary, and validation,
- named as <prefix>-prediction, <prefix>-summary, and <prefix>-validation,
- respectively.
- (<prefix> should contain only alphanumeric characters or hyphen.)
-
- The upstream and downstream can be set accordingly like:
- pred, _, val = create_evaluate_ops(...)
- pred.set_upstream(upstream_op)
- ...
- downstream_op.set_upstream(val)
-
- Callers will provide two python callables, metric_fn and validate_fn, in
- order to customize the evaluation behavior as they wish.
- - metric_fn receives a dictionary per instance derived from json in the
- batch prediction result. The keys might vary depending on the model.
- It should return a tuple of metrics.
- - validation_fn receives a dictionary of the averaged metrics that metric_fn
- generated over all instances.
- The key/value of the dictionary matches to what's given by
- metric_fn_and_keys arg.
- The dictionary contains an additional metric, 'count' to represent the
- total number of instances received for evaluation.
- The function would raise an exception to mark the task as failed, in a
- case the validation result is not okay to proceed (i.e. to set the trained
- version as default).
-
- Typical examples are like this:
-
- def get_metric_fn_and_keys():
- import math # imports should be outside of the metric_fn below.
- def error_and_squared_error(inst):
- label = float(inst['input_label'])
- classes = float(inst['classes']) # 0 or 1
- err = abs(classes-label)
- squared_err = math.pow(classes-label, 2)
- return (err, squared_err) # returns a tuple.
- return error_and_squared_error, ['err', 'mse'] # key order must match.
-
- def validate_err_and_count(summary):
- if summary['err'] > 0.2:
- raise ValueError('Too high err>0.2; summary=%s' % summary)
- if summary['mse'] > 0.05:
- raise ValueError('Too high mse>0.05; summary=%s' % summary)
- if summary['count'] < 1000:
- raise ValueError('Too few instances<1000; summary=%s' % summary)
- return summary
-
- For the details on the other BatchPrediction-related arguments (project_id,
- job_id, region, data_format, input_paths, prediction_path, model_uri),
- please refer to MLEngineBatchPredictionOperator too.
-
- :param task_prefix: a prefix for the tasks. Only alphanumeric characters and
- hyphen are allowed (no underscores), since this will be used as dataflow
- job name, which doesn't allow other characters.
- :type task_prefix: string
-
- :param data_format: either of 'TEXT', 'TF_RECORD', 'TF_RECORD_GZIP'
- :type data_format: string
-
- :param input_paths: a list of input paths to be sent to BatchPrediction.
- :type input_paths: list of strings
-
- :param prediction_path: GCS path to put the prediction results in.
- :type prediction_path: string
-
- :param metric_fn_and_keys: a tuple of metric_fn and metric_keys:
- - metric_fn is a function that accepts a dictionary (for an instance),
- and returns a tuple of metric(s) that it calculates.
- - metric_keys is a list of strings to denote the key of each metric.
- :type metric_fn_and_keys: tuple of a function and a list of strings
-
- :param validate_fn: a function to validate whether the averaged metric(s) is
- good enough to push the model.
- :type validate_fn: function
-
- :param batch_prediction_job_id: the id to use for the Cloud ML Batch
- prediction job. Passed directly to the MLEngineBatchPredictionOperator as
- the job_id argument.
- :type batch_prediction_job_id: string
-
- :param project_id: the Google Cloud Platform project id in which to execute
- Cloud ML Batch Prediction and Dataflow jobs. If None, then the `dag`'s
- `default_args['project_id']` will be used.
- :type project_id: string
-
- :param region: the Google Cloud Platform region in which to execute Cloud ML
- Batch Prediction and Dataflow jobs. If None, then the `dag`'s
- `default_args['region']` will be used.
- :type region: string
-
- :param dataflow_options: options to run Dataflow jobs. If None, then the
- `dag`'s `default_args['dataflow_default_options']` will be used.
- :type dataflow_options: dictionary
-
- :param model_uri: GCS path of the model exported by Tensorflow using
- tensorflow.estimator.export_savedmodel(). It cannot be used with
- model_name or version_name below. See MLEngineBatchPredictionOperator for
- more detail.
- :type model_uri: string
-
- :param model_name: Used to indicate a model to use for prediction. Can be
- used in combination with version_name, but cannot be used together with
- model_uri. See MLEngineBatchPredictionOperator for more detail. If None,
- then the `dag`'s `default_args['model_name']` will be used.
- :type model_name: string
-
- :param version_name: Used to indicate a model version to use for prediciton,
- in combination with model_name. Cannot be used together with model_uri.
- See MLEngineBatchPredictionOperator for more detail. If None, then the
- `dag`'s `default_args['version_name']` will be used.
- :type version_name: string
-
- :param dag: The `DAG` to use for all Operators.
- :type dag: airflow.DAG
-
- :returns: a tuple of three operators, (prediction, summary, validation)
- :rtype: tuple(DataFlowPythonOperator, DataFlowPythonOperator,
- PythonOperator)
- """
-
- # Verify that task_prefix doesn't have any special characters except hyphen
- # '-', which is the only allowed non-alphanumeric character by Dataflow.
- if not re.match(r"^[a-zA-Z][-A-Za-z0-9]*$", task_prefix):
- raise AirflowException(
- "Malformed task_id for DataFlowPythonOperator (only alphanumeric "
- "and hyphens are allowed but got: " + task_prefix)
-
- metric_fn, metric_keys = metric_fn_and_keys
- if not callable(metric_fn):
- raise AirflowException("`metric_fn` param must be callable.")
- if not callable(validate_fn):
- raise AirflowException("`validate_fn` param must be callable.")
-
- if dag is not None and dag.default_args is not None:
- default_args = dag.default_args
- project_id = project_id or default_args.get('project_id')
- region = region or default_args.get('region')
- model_name = model_name or default_args.get('model_name')
- version_name = version_name or default_args.get('version_name')
- dataflow_options = dataflow_options or \
- default_args.get('dataflow_default_options')
-
- evaluate_prediction = MLEngineBatchPredictionOperator(
- task_id=(task_prefix + "-prediction"),
- project_id=project_id,
- job_id=batch_prediction_job_id,
- region=region,
- data_format=data_format,
- input_paths=input_paths,
- output_path=prediction_path,
- uri=model_uri,
- model_name=model_name,
- version_name=version_name,
- dag=dag)
-
- metric_fn_encoded = base64.b64encode(dill.dumps(metric_fn, recurse=True))
- evaluate_summary = DataFlowPythonOperator(
- task_id=(task_prefix + "-summary"),
- py_options=["-m"],
- py_file="airflow.contrib.operators.mlengine_prediction_summary",
- dataflow_default_options=dataflow_options,
- options={
- "prediction_path": prediction_path,
- "metric_fn_encoded": metric_fn_encoded,
- "metric_keys": ','.join(metric_keys)
- },
- dag=dag)
- evaluate_summary.set_upstream(evaluate_prediction)
-
- def apply_validate_fn(*args, **kwargs):
- prediction_path = kwargs["templates_dict"]["prediction_path"]
- scheme, bucket, obj, _, _ = urlsplit(prediction_path)
- if scheme != "gs" or not bucket or not obj:
- raise ValueError("Wrong format prediction_path: %s",
- prediction_path)
- summary = os.path.join(obj.strip("/"),
- "prediction.summary.json")
- gcs_hook = GoogleCloudStorageHook()
- summary = json.loads(gcs_hook.download(bucket, summary))
- return validate_fn(summary)
-
- evaluate_validation = PythonOperator(
- task_id=(task_prefix + "-validation"),
- python_callable=apply_validate_fn,
- provide_context=True,
- templates_dict={"prediction_path": prediction_path},
- dag=dag)
- evaluate_validation.set_upstream(evaluate_summary)
-
- return evaluate_prediction, evaluate_summary, evaluate_validation
diff --git a/airflow/contrib/operators/mlengine_prediction_summary.py b/airflow/contrib/operators/mlengine_prediction_summary.py
deleted file mode 100644
index 17fc2c0903..0000000000
--- a/airflow/contrib/operators/mlengine_prediction_summary.py
+++ /dev/null
@@ -1,175 +0,0 @@
-#
-# 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.
-
-"""A template called by DataFlowPythonOperator to summarize BatchPrediction.
-
-It accepts a user function to calculate the metric(s) per instance in
-the prediction results, then aggregates to output as a summary.
-
-Args:
- --prediction_path:
- The GCS folder that contains BatchPrediction results, containing
- prediction.results-NNNNN-of-NNNNN files in the json format.
- Output will be also stored in this folder, as 'prediction.summary.json'.
-
- --metric_fn_encoded:
- An encoded function that calculates and returns a tuple of metric(s)
- for a given instance (as a dictionary). It should be encoded
- via base64.b64encode(dill.dumps(fn, recurse=True)).
-
- --metric_keys:
- A comma-separated key(s) of the aggregated metric(s) in the summary
- output. The order and the size of the keys must match to the output
- of metric_fn.
- The summary will have an additional key, 'count', to represent the
- total number of instances, so the keys shouldn't include 'count'.
-
-# Usage example:
-def get_metric_fn():
- import math # all imports must be outside of the function to be passed.
- def metric_fn(inst):
- label = float(inst["input_label"])
- classes = float(inst["classes"])
- prediction = float(inst["scores"][1])
- log_loss = math.log(1 + math.exp(
- -(label * 2 - 1) * math.log(prediction / (1 - prediction))))
- squared_err = (classes-label)**2
- return (log_loss, squared_err)
- return metric_fn
-metric_fn_encoded = base64.b64encode(dill.dumps(get_metric_fn(), recurse=True))
-
-airflow.contrib.operators.DataFlowPythonOperator(
- task_id="summary-prediction",
- py_options=["-m"],
- py_file="airflow.contrib.operators.mlengine_prediction_summary",
- options={
- "prediction_path": prediction_path,
- "metric_fn_encoded": metric_fn_encoded,
- "metric_keys": "log_loss,mse"
- },
- dataflow_default_options={
- "project": "xxx", "region": "us-east1",
- "staging_location": "gs://yy", "temp_location": "gs://zz",
- })
- >> dag
-
-# When the input file is like the following:
-{"inputs": "1,x,y,z", "classes": 1, "scores": [0.1, 0.9]}
-{"inputs": "0,o,m,g", "classes": 0, "scores": [0.7, 0.3]}
-{"inputs": "1,o,m,w", "classes": 0, "scores": [0.6, 0.4]}
-{"inputs": "1,b,r,b", "classes": 1, "scores": [0.2, 0.8]}
-
-# The output file will be:
-{"log_loss": 0.43890510565304547, "count": 4, "mse": 0.25}
-
-# To test outside of the dag:
-subprocess.check_call(["python",
- "-m",
- "airflow.contrib.operators.mlengine_prediction_summary",
- "--prediction_path=gs://...",
- "--metric_fn_encoded=" + metric_fn_encoded,
- "--metric_keys=log_loss,mse",
- "--runner=DataflowRunner",
- "--staging_location=gs://...",
- "--temp_location=gs://...",
- ])
-
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import argparse
-import base64
-import json
-import os
-
-import apache_beam as beam
-import dill
-
-
-class JsonCoder(object):
- def encode(self, x):
- return json.dumps(x)
-
- def decode(self, x):
- return json.loads(x)
-
-
-@beam.ptransform_fn
-def MakeSummary(pcoll, metric_fn, metric_keys): # pylint: disable=invalid-name
- return (
- pcoll
- | "ApplyMetricFnPerInstance" >> beam.Map(metric_fn)
- | "PairWith1" >> beam.Map(lambda tup: tup + (1,))
- | "SumTuple" >> beam.CombineGlobally(beam.combiners.TupleCombineFn(
- *([sum] * (len(metric_keys) + 1))))
- | "AverageAndMakeDict" >> beam.Map(
- lambda tup: dict(
- [(name, tup[i]/tup[-1]) for i, name in enumerate(metric_keys)] +
- [("count", tup[-1])])))
-
-
-def run(argv=None):
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--prediction_path", required=True,
- help=(
- "The GCS folder that contains BatchPrediction results, containing "
- "prediction.results-NNNNN-of-NNNNN files in the json format. "
- "Output will be also stored in this folder, as a file"
- "'prediction.summary.json'."))
- parser.add_argument(
- "--metric_fn_encoded", required=True,
- help=(
- "An encoded function that calculates and returns a tuple of "
- "metric(s) for a given instance (as a dictionary). It should be "
- "encoded via base64.b64encode(dill.dumps(fn, recurse=True))."))
- parser.add_argument(
- "--metric_keys", required=True,
- help=(
- "A comma-separated keys of the aggregated metric(s) in the summary "
- "output. The order and the size of the keys must match to the "
- "output of metric_fn. The summary will have an additional key, "
- "'count', to represent the total number of instances, so this flag "
- "shouldn't include 'count'."))
- known_args, pipeline_args = parser.parse_known_args(argv)
-
- metric_fn = dill.loads(base64.b64decode(known_args.metric_fn_encoded))
- if not callable(metric_fn):
- raise ValueError("--metric_fn_encoded must be an encoded callable.")
- metric_keys = known_args.metric_keys.split(",")
-
- with beam.Pipeline(
- options=beam.pipeline.PipelineOptions(pipeline_args)) as p:
- # This is apache-beam ptransform's convention
- # pylint: disable=no-value-for-parameter
- _ = (p
- | "ReadPredictionResult" >> beam.io.ReadFromText(
- os.path.join(known_args.prediction_path,
- "prediction.results-*-of-*"),
- coder=JsonCoder())
- | "Summary" >> MakeSummary(metric_fn, metric_keys)
- | "Write" >> beam.io.WriteToText(
- os.path.join(known_args.prediction_path,
- "prediction.summary.json"),
- shard_name_template='', # without trailing -NNNNN-of-NNNNN.
- coder=JsonCoder()))
- # pylint: enable=no-value-for-parameter
-
-
-if __name__ == "__main__":
- run()
diff --git a/airflow/contrib/operators/mongo_to_s3.py b/airflow/contrib/operators/mongo_to_s3.py
deleted file mode 100644
index 43b5d8b6c3..0000000000
--- a/airflow/contrib/operators/mongo_to_s3.py
+++ /dev/null
@@ -1,119 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Licensed 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 json
-
-from airflow.contrib.hooks.mongo_hook import MongoHook
-from airflow.hooks.S3_hook import S3Hook
-from airflow.models import BaseOperator
-from bson import json_util
-
-
-class MongoToS3Operator(BaseOperator):
- """
- Mongo -> S3
- A more specific baseOperator meant to move data
- from mongo via pymongo to s3 via boto
-
- things to note
- .execute() is written to depend on .transform()
- .transform() is meant to be extended by child classes
- to perform transformations unique to those operators needs
- """
-
- template_fields = ['s3_key', 'mongo_query']
- # pylint: disable=too-many-instance-attributes
-
- def __init__(self,
- mongo_conn_id,
- s3_conn_id,
- mongo_collection,
- mongo_query,
- s3_bucket,
- s3_key,
- mongo_db=None,
- *args, **kwargs):
- super(MongoToS3Operator, self).__init__(*args, **kwargs)
- # Conn Ids
- self.mongo_conn_id = mongo_conn_id
- self.s3_conn_id = s3_conn_id
- # Mongo Query Settings
- self.mongo_db = mongo_db
- self.mongo_collection = mongo_collection
- # Grab query and determine if we need to run an aggregate pipeline
- self.mongo_query = mongo_query
- self.is_pipeline = True if isinstance(
- self.mongo_query, list) else False
-
- # S3 Settings
- self.s3_bucket = s3_bucket
- self.s3_key = s3_key
-
- # KWARGS
- self.replace = kwargs.pop('replace', False)
-
- def execute(self, context):
- """
- Executed by task_instance at runtime
- """
- s3_conn = S3Hook(self.s3_conn_id)
-
- # Grab collection and execute query according to whether or not it is a pipeline
- if self.is_pipeline:
- results = MongoHook(self.mongo_conn_id).aggregate(
- mongo_collection=self.mongo_collection,
- aggregate_query=self.mongo_query,
- mongo_db=self.mongo_db
- )
-
- else:
- results = MongoHook(self.mongo_conn_id).find(
- mongo_collection=self.mongo_collection,
- query=self.mongo_query,
- mongo_db=self.mongo_db
- )
-
- # Performs transform then stringifies the docs results into json format
- docs_str = self._stringify(self.transform(results))
-
- # Load Into S3
- s3_conn.load_string(
- string_data=docs_str,
- key=self.s3_key,
- bucket_name=self.s3_bucket,
- replace=self.replace
- )
-
- return True
-
- def _stringify(self, iterable, joinable='\n'):
- """
- Takes an iterable (pymongo Cursor or Array) containing dictionaries and
- returns a stringified version using python join
- """
- return joinable.join(
- [json.dumps(doc, default=json_util.default) for doc in iterable]
- )
-
- def transform(self, docs):
- """
- Processes pyMongo cursor and returns an iterable with each element being
- a JSON serializable dictionary
-
- Base transform() assumes no processing is needed
- ie. docs is a pyMongo cursor of documents and cursor just
- needs to be passed through
-
- Override this method for custom transformations
- """
- return docs
diff --git a/airflow/contrib/operators/mysql_to_gcs.py b/airflow/contrib/operators/mysql_to_gcs.py
deleted file mode 100644
index 4d1bb7b329..0000000000
--- a/airflow/contrib/operators/mysql_to_gcs.py
+++ /dev/null
@@ -1,288 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 sys
-import json
-import time
-import base64
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.hooks.mysql_hook import MySqlHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from datetime import date, datetime
-from decimal import Decimal
-from MySQLdb.constants import FIELD_TYPE
-from tempfile import NamedTemporaryFile
-from six import string_types
-
-PY3 = sys.version_info[0] == 3
-
-
-class MySqlToGoogleCloudStorageOperator(BaseOperator):
- """
- Copy data from MySQL to Google cloud storage in JSON format.
- """
- template_fields = ('sql', 'bucket', 'filename', 'schema_filename', 'schema')
- template_ext = ('.sql',)
- ui_color = '#a0e08c'
-
- @apply_defaults
- def __init__(self,
- sql,
- bucket,
- filename,
- schema_filename=None,
- approx_max_file_size_bytes=1900000000,
- mysql_conn_id='mysql_default',
- google_cloud_storage_conn_id='google_cloud_default',
- schema=None,
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param sql: The SQL to execute on the MySQL table.
- :type sql: string
- :param bucket: The bucket to upload to.
- :type bucket: string
- :param filename: The filename to use as the object name when uploading
- to Google cloud storage. A {} should be specified in the filename
- to allow the operator to inject file numbers in cases where the
- file is split due to size.
- :type filename: string
- :param schema_filename: If set, the filename to use as the object name
- when uploading a .json file containing the BigQuery schema fields
- for the table that was dumped from MySQL.
- :type schema_filename: string
- :param approx_max_file_size_bytes: This operator supports the ability
- to split large table dumps into multiple files (see notes in the
- filenamed param docs above). Google cloud storage allows for files
- to be a maximum of 4GB. This param allows developers to specify the
- file size of the splits.
- :type approx_max_file_size_bytes: long
- :param mysql_conn_id: Reference to a specific MySQL hook.
- :type mysql_conn_id: string
- :param google_cloud_storage_conn_id: Reference to a specific Google
- cloud storage hook.
- :type google_cloud_storage_conn_id: string
- :param schema: The schema to use, if any. Should be a list of dict or
- a str. Pass a string if using Jinja template, otherwise, pass a list of
- dict. Examples could be seen: https://cloud.google.com/bigquery/docs
- /schemas#specifying_a_json_schema_file
- :type schema: str or list
- :param delegate_to: The account to impersonate, if any. For this to
- work, the service account making the request must have domain-wide
- delegation enabled.
- """
- super(MySqlToGoogleCloudStorageOperator, self).__init__(*args, **kwargs)
- self.sql = sql
- self.bucket = bucket
- self.filename = filename
- self.schema_filename = schema_filename
- self.approx_max_file_size_bytes = approx_max_file_size_bytes
- self.mysql_conn_id = mysql_conn_id
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.schema = schema
- self.delegate_to = delegate_to
-
- def execute(self, context):
- cursor = self._query_mysql()
- files_to_upload = self._write_local_data_files(cursor)
-
- # If a schema is set, create a BQ schema JSON file.
- if self.schema_filename:
- files_to_upload.update(self._write_local_schema_file(cursor))
-
- # Flush all files before uploading.
- for file_handle in files_to_upload.values():
- file_handle.flush()
-
- self._upload_to_gcs(files_to_upload)
-
- # Close all temp file handles.
- for file_handle in files_to_upload.values():
- file_handle.close()
-
- def _query_mysql(self):
- """
- Queries mysql and returns a cursor to the results.
- """
- mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id)
- conn = mysql.get_conn()
- cursor = conn.cursor()
- cursor.execute(self.sql)
- return cursor
-
- def _write_local_data_files(self, cursor):
- """
- Takes a cursor, and writes results to a local file.
-
- :return: A dictionary where keys are filenames to be used as object
- names in GCS, and values are file handles to local files that
- contain the data for the GCS objects.
- """
- schema = list(map(lambda schema_tuple: schema_tuple[0], cursor.description))
- col_type_dict = self._get_col_type_dict()
- file_no = 0
- tmp_file_handle = NamedTemporaryFile(delete=True)
- tmp_file_handles = {self.filename.format(file_no): tmp_file_handle}
-
- for row in cursor:
- # Convert datetime objects to utc seconds, and decimals to floats.
- # Convert binary type object to string encoded with base64.
- row = self._convert_types(schema, col_type_dict, row)
- row_dict = dict(zip(schema, row))
-
- # TODO validate that row isn't > 2MB. BQ enforces a hard row size of 2MB.
- s = json.dumps(row_dict)
- if PY3:
- s = s.encode('utf-8')
- tmp_file_handle.write(s)
-
- # Append newline to make dumps BigQuery compatible.
- tmp_file_handle.write(b'\n')
-
- # Stop if the file exceeds the file size limit.
- if tmp_file_handle.tell() >= self.approx_max_file_size_bytes:
- file_no += 1
- tmp_file_handle = NamedTemporaryFile(delete=True)
- tmp_file_handles[self.filename.format(file_no)] = tmp_file_handle
-
- return tmp_file_handles
-
- def _write_local_schema_file(self, cursor):
- """
- Takes a cursor, and writes the BigQuery schema for the results to a
- local file system.
-
- :return: A dictionary where key is a filename to be used as an object
- name in GCS, and values are file handles to local files that
- contains the BigQuery schema fields in .json format.
- """
- schema_str = None
- tmp_schema_file_handle = NamedTemporaryFile(delete=True)
- if self.schema is not None and isinstance(self.schema, string_types):
- schema_str = self.schema
- elif self.schema is not None and isinstance(self.schema, list):
- schema_str = json.dumps(self.schema)
- else:
- schema = []
- for field in cursor.description:
- # See PEP 249 for details about the description tuple.
- field_name = field[0]
- field_type = self.type_map(field[1])
- # Always allow TIMESTAMP to be nullable. MySQLdb returns None types
- # for required fields because some MySQL timestamps can't be
- # represented by Python's datetime (e.g. 0000-00-00 00:00:00).
- if field[6] or field_type == 'TIMESTAMP':
- field_mode = 'NULLABLE'
- else:
- field_mode = 'REQUIRED'
- schema.append({
- 'name': field_name,
- 'type': field_type,
- 'mode': field_mode,
- })
- schema_str = json.dumps(schema)
- if PY3:
- schema_str = schema_str.encode('utf-8')
- tmp_schema_file_handle.write(schema_str)
-
- self.log.info('Using schema for %s: %s', self.schema_filename, schema_str)
- return {self.schema_filename: tmp_schema_file_handle}
-
- def _upload_to_gcs(self, files_to_upload):
- """
- Upload all of the file splits (and optionally the schema .json file) to
- Google cloud storage.
- """
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- for object, tmp_file_handle in files_to_upload.items():
- hook.upload(self.bucket, object, tmp_file_handle.name, 'application/json')
-
- def _convert_types(self, schema, col_type_dict, row):
- """
- Takes a value from MySQLdb, and converts it to a value that's safe for
- JSON/Google cloud storage/BigQuery. Dates are converted to UTC seconds.
- Decimals are converted to floats. Binary type fields are encoded with base64,
- as imported BYTES data must be base64-encoded according to Bigquery SQL
- date type documentation: https://cloud.google.com/bigquery/data-types
- """
- converted_row = []
- for col_name, col_val in zip(schema, row):
- if type(col_val) in (datetime, date):
- col_val = time.mktime(col_val.timetuple())
- elif isinstance(col_val, Decimal):
- col_val = float(col_val)
- elif col_type_dict.get(col_name) == "BYTES":
- col_val = base64.standard_b64encode(col_val)
- if PY3:
- col_val = col_val.decode('ascii')
- else:
- col_val = col_val
- converted_row.append(col_val)
- return converted_row
-
- def _get_col_type_dict(self):
- """
- Return a dict of column name and column type based on self.schema if not None.
- """
- schema = []
- if isinstance(self.schema, string_types):
- schema = json.loads(self.schema)
- elif isinstance(self.schema, list):
- schema = self.schema
- elif self.schema is not None:
- self.log.warn('Using default schema due to unexpected type.'
- 'Should be a string or list.')
-
- col_type_dict = {}
- try:
- col_type_dict = {col['name']: col['type'] for col in schema}
- except KeyError:
- self.log.warn('Using default schema due to missing name or type. Please '
- 'refer to: https://cloud.google.com/bigquery/docs/schemas'
- '#specifying_a_json_schema_file')
- return col_type_dict
-
- @classmethod
- def type_map(cls, mysql_type):
- """
- Helper function that maps from MySQL fields to BigQuery fields. Used
- when a schema_filename is set.
- """
- d = {
- FIELD_TYPE.INT24: 'INTEGER',
- FIELD_TYPE.TINY: 'INTEGER',
- FIELD_TYPE.BIT: 'INTEGER',
- FIELD_TYPE.DATETIME: 'TIMESTAMP',
- FIELD_TYPE.DATE: 'TIMESTAMP',
- FIELD_TYPE.DECIMAL: 'FLOAT',
- FIELD_TYPE.NEWDECIMAL: 'FLOAT',
- FIELD_TYPE.DOUBLE: 'FLOAT',
- FIELD_TYPE.FLOAT: 'FLOAT',
- FIELD_TYPE.LONG: 'INTEGER',
- FIELD_TYPE.LONGLONG: 'INTEGER',
- FIELD_TYPE.SHORT: 'INTEGER',
- FIELD_TYPE.TIMESTAMP: 'TIMESTAMP',
- FIELD_TYPE.YEAR: 'INTEGER',
- }
- return d[mysql_type] if mysql_type in d else 'STRING'
diff --git a/airflow/contrib/operators/postgres_to_gcs_operator.py b/airflow/contrib/operators/postgres_to_gcs_operator.py
deleted file mode 100644
index 88b4d00e39..0000000000
--- a/airflow/contrib/operators/postgres_to_gcs_operator.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 sys
-import json
-import time
-import datetime
-
-from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
-from airflow.hooks.postgres_hook import PostgresHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from decimal import Decimal
-from tempfile import NamedTemporaryFile
-
-PY3 = sys.version_info[0] == 3
-
-
-class PostgresToGoogleCloudStorageOperator(BaseOperator):
- """
- Copy data from Postgres to Google Cloud Storage in JSON format.
- """
- template_fields = ('sql', 'bucket', 'filename', 'schema_filename',
- 'parameters')
- template_ext = ('.sql', )
- ui_color = '#a0e08c'
-
- @apply_defaults
- def __init__(self,
- sql,
- bucket,
- filename,
- schema_filename=None,
- approx_max_file_size_bytes=1900000000,
- postgres_conn_id='postgres_default',
- google_cloud_storage_conn_id='google_cloud_default',
- delegate_to=None,
- parameters=None,
- *args,
- **kwargs):
- """
- :param sql: The SQL to execute on the Postgres table.
- :type sql: string
- :param bucket: The bucket to upload to.
- :type bucket: string
- :param filename: The filename to use as the object name when uploading
- to Google Cloud Storage. A {} should be specified in the filename
- to allow the operator to inject file numbers in cases where the
- file is split due to size.
- :type filename: string
- :param schema_filename: If set, the filename to use as the object name
- when uploading a .json file containing the BigQuery schema fields
- for the table that was dumped from Postgres.
- :type schema_filename: string
- :param approx_max_file_size_bytes: This operator supports the ability
- to split large table dumps into multiple files (see notes in the
- filenamed param docs above). Google Cloud Storage allows for files
- to be a maximum of 4GB. This param allows developers to specify the
- file size of the splits.
- :type approx_max_file_size_bytes: long
- :param postgres_conn_id: Reference to a specific Postgres hook.
- :type postgres_conn_id: string
- :param google_cloud_storage_conn_id: Reference to a specific Google
- cloud storage hook.
- :type google_cloud_storage_conn_id: string
- :param delegate_to: The account to impersonate, if any. For this to
- work, the service account making the request must have domain-wide
- delegation enabled.
- :param parameters: a parameters dict that is substituted at query runtime.
- :type parameters: dict
- """
- super(PostgresToGoogleCloudStorageOperator, self).__init__(*args, **kwargs)
- self.sql = sql
- self.bucket = bucket
- self.filename = filename
- self.schema_filename = schema_filename
- self.approx_max_file_size_bytes = approx_max_file_size_bytes
- self.postgres_conn_id = postgres_conn_id
- self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
- self.delegate_to = delegate_to
- self.parameters = parameters
-
- def execute(self, context):
- cursor = self._query_postgres()
- files_to_upload = self._write_local_data_files(cursor)
-
- # If a schema is set, create a BQ schema JSON file.
- if self.schema_filename:
- files_to_upload.update(self._write_local_schema_file(cursor))
-
- # Flush all files before uploading
- for file_handle in files_to_upload.values():
- file_handle.flush()
-
- self._upload_to_gcs(files_to_upload)
-
- # Close all temp file handles.
- for file_handle in files_to_upload.values():
- file_handle.close()
-
- def _query_postgres(self):
- """
- Queries Postgres and returns a cursor to the results.
- """
- postgres = PostgresHook(postgres_conn_id=self.postgres_conn_id)
- conn = postgres.get_conn()
- cursor = conn.cursor()
- cursor.execute(self.sql, self.parameters)
- return cursor
-
- def _write_local_data_files(self, cursor):
- """
- Takes a cursor, and writes results to a local file.
-
- :return: A dictionary where keys are filenames to be used as object
- names in GCS, and values are file handles to local files that
- contain the data for the GCS objects.
- """
- schema = list(map(lambda schema_tuple: schema_tuple[0], cursor.description))
- file_no = 0
- tmp_file_handle = NamedTemporaryFile(delete=True)
- tmp_file_handles = {self.filename.format(file_no): tmp_file_handle}
-
- for row in cursor:
- # Convert datetime objects to utc seconds, and decimals to floats
- row = map(self.convert_types, row)
- row_dict = dict(zip(schema, row))
-
- s = json.dumps(row_dict, sort_keys=True)
- if PY3:
- s = s.encode('utf-8')
- tmp_file_handle.write(s)
-
- # Append newline to make dumps BigQuery compatible.
- tmp_file_handle.write(b'\n')
-
- # Stop if the file exceeds the file size limit.
- if tmp_file_handle.tell() >= self.approx_max_file_size_bytes:
- file_no += 1
- tmp_file_handle = NamedTemporaryFile(delete=True)
- tmp_file_handles[self.filename.format(file_no)] = tmp_file_handle
-
- return tmp_file_handles
-
- def _write_local_schema_file(self, cursor):
- """
- Takes a cursor, and writes the BigQuery schema for the results to a
- local file system.
-
- :return: A dictionary where key is a filename to be used as an object
- name in GCS, and values are file handles to local files that
- contains the BigQuery schema fields in .json format.
- """
- schema = []
- for field in cursor.description:
- # See PEP 249 for details about the description tuple.
- field_name = field[0]
- field_type = self.type_map(field[1])
- field_mode = 'REPEATED' if field[1] in (1009, 1005, 1007,
- 1016) else 'NULLABLE'
- schema.append({
- 'name': field_name,
- 'type': field_type,
- 'mode': field_mode,
- })
-
- self.log.info('Using schema for %s: %s', self.schema_filename, schema)
- tmp_schema_file_handle = NamedTemporaryFile(delete=True)
- s = json.dumps(schema, sort_keys=True)
- if PY3:
- s = s.encode('utf-8')
- tmp_schema_file_handle.write(s)
- return {self.schema_filename: tmp_schema_file_handle}
-
- def _upload_to_gcs(self, files_to_upload):
- """
- Upload all of the file splits (and optionally the schema .json file) to
- Google Cloud Storage.
- """
- hook = GoogleCloudStorageHook(
- google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- for object, tmp_file_handle in files_to_upload.items():
- hook.upload(self.bucket, object, tmp_file_handle.name,
- 'application/json')
-
- @classmethod
- def convert_types(cls, value):
- """
- Takes a value from Postgres, and converts it to a value that's safe for
- JSON/Google Cloud Storage/BigQuery. Dates are converted to UTC seconds.
- Decimals are converted to floats. Times are converted to seconds.
- """
- if type(value) in (datetime.datetime, datetime.date):
- return time.mktime(value.timetuple())
- elif type(value) == datetime.time:
- formated_time = time.strptime(str(value), "%H:%M:%S")
- return datetime.timedelta(
- hours=formated_time.tm_hour,
- minutes=formated_time.tm_min,
- seconds=formated_time.tm_sec).seconds
- elif isinstance(value, Decimal):
- return float(value)
- else:
- return value
-
- @classmethod
- def type_map(cls, postgres_type):
- """
- Helper function that maps from Postgres fields to BigQuery fields. Used
- when a schema_filename is set.
- """
- d = {
- 1114: 'TIMESTAMP',
- 1184: 'TIMESTAMP',
- 1082: 'TIMESTAMP',
- 1083: 'TIMESTAMP',
- 1005: 'INTEGER',
- 1007: 'INTEGER',
- 1016: 'INTEGER',
- 20: 'INTEGER',
- 21: 'INTEGER',
- 23: 'INTEGER',
- 16: 'BOOLEAN',
- 700: 'FLOAT',
- 701: 'FLOAT',
- 1700: 'FLOAT'
- }
-
- return d[postgres_type] if postgres_type in d else 'STRING'
diff --git a/airflow/contrib/operators/pubsub_operator.py b/airflow/contrib/operators/pubsub_operator.py
deleted file mode 100644
index 2d55b19f85..0000000000
--- a/airflow/contrib/operators/pubsub_operator.py
+++ /dev/null
@@ -1,433 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.hooks.gcp_pubsub_hook import PubSubHook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class PubSubTopicCreateOperator(BaseOperator):
- """Create a PubSub topic.
-
- By default, if the topic already exists, this operator will
- not cause the DAG to fail. ::
-
- with DAG('successful DAG') as dag:
- (
- dag
- >> PubSubTopicCreateOperator(project='my-project',
- topic='my_new_topic')
- >> PubSubTopicCreateOperator(project='my-project',
- topic='my_new_topic')
- )
-
- The operator can be configured to fail if the topic already exists. ::
-
- with DAG('failing DAG') as dag:
- (
- dag
- >> PubSubTopicCreateOperator(project='my-project',
- topic='my_new_topic')
- >> PubSubTopicCreateOperator(project='my-project',
- topic='my_new_topic',
- fail_if_exists=True)
- )
-
- Both ``project`` and ``topic`` are templated so you can use
- variables in them.
- """
- template_fields = ['project', 'topic']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- project,
- topic,
- fail_if_exists=False,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param project: the GCP project ID where the topic will be created
- :type project: string
- :param topic: the topic to create. Do not include the
- full topic path. In other words, instead of
- ``projects/{project}/topics/{topic}``, provide only
- ``{topic}``. (templated)
- :type topic: string
- :param gcp_conn_id: The connection ID to use connecting to
- Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request
- must have domain-wide delegation enabled.
- :type delegate_to: string
- """
- super(PubSubTopicCreateOperator, self).__init__(*args, **kwargs)
-
- self.project = project
- self.topic = topic
- self.fail_if_exists = fail_if_exists
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- hook.create_topic(self.project, self.topic,
- fail_if_exists=self.fail_if_exists)
-
-
-class PubSubSubscriptionCreateOperator(BaseOperator):
- """Create a PubSub subscription.
-
- By default, the subscription will be created in ``topic_project``. If
- ``subscription_project`` is specified and the GCP credentials allow, the
- Subscription can be created in a different project from its topic.
-
- By default, if the subscription already exists, this operator will
- not cause the DAG to fail. However, the topic must exist in the project. ::
-
- with DAG('successful DAG') as dag:
- (
- dag
- >> PubSubSubscriptionCreateOperator(
- topic_project='my-project', topic='my-topic',
- subscription='my-subscription')
- >> PubSubSubscriptionCreateOperator(
- topic_project='my-project', topic='my-topic',
- subscription='my-subscription')
- )
-
- The operator can be configured to fail if the subscription already exists.
- ::
-
- with DAG('failing DAG') as dag:
- (
- dag
- >> PubSubSubscriptionCreateOperator(
- topic_project='my-project', topic='my-topic',
- subscription='my-subscription')
- >> PubSubSubscriptionCreateOperator(
- topic_project='my-project', topic='my-topic',
- subscription='my-subscription', fail_if_exists=True)
- )
-
- Finally, subscription is not required. If not passed, the operator will
- generated a universally unique identifier for the subscription's name. ::
-
- with DAG('DAG') as dag:
- (
- dag >> PubSubSubscriptionCreateOperator(
- topic_project='my-project', topic='my-topic')
- )
-
- ``topic_project``, ``topic``, ``subscription``, and
- ``subscription`` are templated so you can use variables in them.
- """
- template_fields = ['topic_project', 'topic', 'subscription',
- 'subscription_project']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- topic_project,
- topic,
- subscription=None,
- subscription_project=None,
- ack_deadline_secs=10,
- fail_if_exists=False,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param topic_project: the GCP project ID where the topic exists
- :type topic_project: string
- :param topic: the topic to create. Do not include the
- full topic path. In other words, instead of
- ``projects/{project}/topics/{topic}``, provide only
- ``{topic}``. (templated)
- :type topic: string
- :param subscription: the Pub/Sub subscription name. If empty, a random
- name will be generated using the uuid module
- :type subscription: string
- :param subscription_project: the GCP project ID where the subscription
- will be created. If empty, ``topic_project`` will be used.
- :type subscription_project: string
- :param ack_deadline_secs: Number of seconds that a subscriber has to
- acknowledge each message pulled from the subscription
- :type ack_deadline_secs: int
- :param gcp_conn_id: The connection ID to use connecting to
- Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request
- must have domain-wide delegation enabled.
- :type delegate_to: string
- """
- super(PubSubSubscriptionCreateOperator, self).__init__(*args, **kwargs)
-
- self.topic_project = topic_project
- self.topic = topic
- self.subscription = subscription
- self.subscription_project = subscription_project
- self.ack_deadline_secs = ack_deadline_secs
- self.fail_if_exists = fail_if_exists
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- return hook.create_subscription(
- self.topic_project, self.topic, self.subscription,
- self.subscription_project, self.ack_deadline_secs,
- self.fail_if_exists)
-
-
-class PubSubTopicDeleteOperator(BaseOperator):
- """Delete a PubSub topic.
-
- By default, if the topic does not exist, this operator will
- not cause the DAG to fail. ::
-
- with DAG('successful DAG') as dag:
- (
- dag
- >> PubSubTopicDeleteOperator(project='my-project',
- topic='non_existing_topic')
- )
-
- The operator can be configured to fail if the topic does not exist. ::
-
- with DAG('failing DAG') as dag:
- (
- dag
- >> PubSubTopicCreateOperator(project='my-project',
- topic='non_existing_topic',
- fail_if_not_exists=True)
- )
-
- Both ``project`` and ``topic`` are templated so you can use
- variables in them.
- """
- template_fields = ['project', 'topic']
- ui_color = '#cb4335'
-
- @apply_defaults
- def __init__(
- self,
- project,
- topic,
- fail_if_not_exists=False,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param project: the GCP project ID in which to work (templated)
- :type project: string
- :param topic: the topic to delete. Do not include the
- full topic path. In other words, instead of
- ``projects/{project}/topics/{topic}``, provide only
- ``{topic}``. (templated)
- :type topic: string
- :param fail_if_not_exists: If True and the topic does not exist, fail
- the task
- :type fail_if_not_exists: bool
- :param gcp_conn_id: The connection ID to use connecting to
- Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request
- must have domain-wide delegation enabled.
- :type delegate_to: string
- """
- super(PubSubTopicDeleteOperator, self).__init__(*args, **kwargs)
-
- self.project = project
- self.topic = topic
- self.fail_if_not_exists = fail_if_not_exists
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- hook.delete_topic(self.project, self.topic,
- fail_if_not_exists=self.fail_if_not_exists)
-
-
-class PubSubSubscriptionDeleteOperator(BaseOperator):
- """Delete a PubSub subscription.
-
- By default, if the subscription does not exist, this operator will
- not cause the DAG to fail. ::
-
- with DAG('successful DAG') as dag:
- (
- dag
- >> PubSubSubscriptionDeleteOperator(project='my-project',
- subscription='non-existing')
- )
-
- The operator can be configured to fail if the subscription already exists.
-
- ::
-
- with DAG('failing DAG') as dag:
- (
- dag
- >> PubSubSubscriptionDeleteOperator(
- project='my-project', subscription='non-existing',
- fail_if_not_exists=True)
- )
-
- ``project``, and ``subscription`` are templated so you can use
- variables in them.
- """
- template_fields = ['project', 'subscription']
- ui_color = '#cb4335'
-
- @apply_defaults
- def __init__(
- self,
- project,
- subscription,
- fail_if_not_exists=False,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param project: the GCP project ID in which to work (templated)
- :type project: string
- :param subscription: the subscription to delete. Do not include the
- full subscription path. In other words, instead of
- ``projects/{project}/subscription/{subscription}``, provide only
- ``{subscription}``. (templated)
- :type subscription: string
- :param fail_if_not_exists: If True and the subscription does not exist,
- fail the task
- :type fail_if_not_exists: bool
- :param gcp_conn_id: The connection ID to use connecting to
- Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request
- must have domain-wide delegation enabled.
- :type delegate_to: string
- """
- super(PubSubSubscriptionDeleteOperator, self).__init__(*args, **kwargs)
-
- self.project = project
- self.subscription = subscription
- self.fail_if_not_exists = fail_if_not_exists
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
-
- def execute(self, context):
- hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- hook.delete_subscription(self.project, self.subscription,
- fail_if_not_exists=self.fail_if_not_exists)
-
-
-class PubSubPublishOperator(BaseOperator):
- """Publish messages to a PubSub topic.
-
- Each Task publishes all provided messages to the same topic
- in a single GCP project. If the topic does not exist, this
- task will fail. ::
-
- from base64 import b64encode as b64e
-
- m1 = {'data': b64e('Hello, World!'),
- 'attributes': {'type': 'greeting'}
- }
- m2 = {'data': b64e('Knock, knock')}
- m3 = {'attributes': {'foo': ''}}
-
- t1 = PubSubPublishOperator(
- project='my-project',topic='my_topic',
- messages=[m1, m2, m3],
- create_topic=True,
- dag=dag)
-
- ``project`` , ``topic``, and ``messages`` are templated so you can use
- variables in them.
- """
- template_fields = ['project', 'topic', 'messages']
- ui_color = '#0273d4'
-
- @apply_defaults
- def __init__(
- self,
- project,
- topic,
- messages,
- gcp_conn_id='google_cloud_default',
- delegate_to=None,
- *args,
- **kwargs):
- """
- :param project: the GCP project ID in which to work (templated)
- :type project: string
- :param topic: the topic to which to publish. Do not include the
- full topic path. In other words, instead of
- ``projects/{project}/topics/{topic}``, provide only
- ``{topic}``. (templated)
- :type topic: string
- :param messages: a list of messages to be published to the
- topic. Each message is a dict with one or more of the
- following keys-value mappings:
- * 'data': a base64-encoded string
- * 'attributes': {'key1': 'value1', ...}
- Each message must contain at least a non-empty 'data' value
- or an attribute dict with at least one key (templated). See
- https://cloud.google.com/pubsub/docs/reference/rest/v1/PubsubMessage
- :type messages: list
- :param gcp_conn_id: The connection ID to use connecting to
- Google Cloud Platform.
- :type gcp_conn_id: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request
- must have domain-wide delegation enabled.
- :type delegate_to: string
- """
- super(PubSubPublishOperator, self).__init__(*args, **kwargs)
-
- self.gcp_conn_id = gcp_conn_id
- self.delegate_to = delegate_to
- self.project = project
- self.topic = topic
- self.messages = messages
-
- def execute(self, context):
- hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
- hook.publish(self.project, self.topic, self.messages)
diff --git a/airflow/contrib/operators/qubole_check_operator.py b/airflow/contrib/operators/qubole_check_operator.py
deleted file mode 100644
index 0e8d75e167..0000000000
--- a/airflow/contrib/operators/qubole_check_operator.py
+++ /dev/null
@@ -1,225 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.contrib.operators.qubole_operator import QuboleOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.contrib.hooks.qubole_check_hook import QuboleCheckHook
-from airflow.operators.check_operator import CheckOperator, ValueCheckOperator
-from airflow.exceptions import AirflowException
-
-
-class QuboleCheckOperator(CheckOperator, QuboleOperator):
- """
- Performs checks against Qubole Commands. ``QuboleCheckOperator`` expects
- a command that will be executed on QDS.
- By default, each value on first row of the result of this Qubole Commmand
- is evaluated using python ``bool`` casting. If any of the
- values return ``False``, the check is failed and errors out.
-
- Note that Python bool casting evals the following as ``False``:
-
- * ``False``
- * ``0``
- * Empty string (``""``)
- * Empty list (``[]``)
- * Empty dictionary or set (``{}``)
-
- Given a query like ``SELECT COUNT(*) FROM foo``, it will fail only if
- the count ``== 0``. You can craft much more complex query that could,
- for instance, check that the table has the same number of rows as
- the source table upstream, or that the count of today's partition is
- greater than yesterday's partition, or that a set of metrics are less
- than 3 standard deviation for the 7 day average.
-
- This operator can be used as a data quality check in your pipeline, and
- depending on where you put it in your DAG, you have the choice to
- stop the critical path, preventing from
- publishing dubious data, or on the side and receive email alerts
- without stopping the progress of the DAG.
-
- :param qubole_conn_id: Connection id which consists of qds auth_token
- :type qubole_conn_id: str
-
- kwargs:
-
- Arguments specific to Qubole command can be referred from QuboleOperator docs.
-
- :results_parser_callable: This is an optional parameter to
- extend the flexibility of parsing the results of Qubole
- command to the users. This is a python callable which
- can hold the logic to parse list of rows returned by Qubole command.
- By default, only the values on first row are used for performing checks.
- This callable should return a list of records on
- which the checks have to be performed.
-
- .. note:: All fields in common with template fields of
- QuboleOperator and CheckOperator are template-supported.
- """
-
- template_fields = QuboleOperator.template_fields + CheckOperator.template_fields
- template_ext = QuboleOperator.template_ext
- ui_fgcolor = '#000'
-
- @apply_defaults
- def __init__(self, qubole_conn_id="qubole_default", *args, **kwargs):
- sql = get_sql_from_qbol_cmd(kwargs)
- super(QuboleCheckOperator, self)\
- .__init__(qubole_conn_id=qubole_conn_id, sql=sql, *args, **kwargs)
- self.on_failure_callback = QuboleCheckHook.handle_failure_retry
- self.on_retry_callback = QuboleCheckHook.handle_failure_retry
-
- def execute(self, context=None):
- try:
- self.hook = self.get_hook(context=context)
- super(QuboleCheckOperator, self).execute(context=context)
- except AirflowException as e:
- handle_airflow_exception(e, self.get_hook())
-
- def get_db_hook(self):
- return self.get_hook()
-
- def get_hook(self, context=None):
- if hasattr(self, 'hook') and (self.hook is not None):
- return self.hook
- else:
- return QuboleCheckHook(context=context, *self.args, **self.kwargs)
-
- def __getattribute__(self, name):
- if name in QuboleCheckOperator.template_fields:
- if name in self.kwargs:
- return self.kwargs[name]
- else:
- return ''
- else:
- return object.__getattribute__(self, name)
-
- def __setattr__(self, name, value):
- if name in QuboleCheckOperator.template_fields:
- self.kwargs[name] = value
- else:
- object.__setattr__(self, name, value)
-
-
-class QuboleValueCheckOperator(ValueCheckOperator, QuboleOperator):
- """
- Performs a simple value check using Qubole command.
- By default, each value on the first row of this
- Qubole command is compared with a pre-defined value.
- The check fails and errors out if the output of the command
- is not within the permissible limit of expected value.
-
- :param qubole_conn_id: Connection id which consists of qds auth_token
- :type qubole_conn_id: str
-
- :param pass_value: Expected value of the query results.
- :type pass_value: str/int/float
-
- :param tolerance: Defines the permissible pass_value range, for example if
- tolerance is 2, the Qubole command output can be anything between
- -2*pass_value and 2*pass_value, without the operator erring out.
-
- :type tolerance: int/float
-
-
- kwargs:
-
- Arguments specific to Qubole command can be referred from QuboleOperator docs.
-
- :results_parser_callable: This is an optional parameter to
- extend the flexibility of parsing the results of Qubole
- command to the users. This is a python callable which
- can hold the logic to parse list of rows returned by Qubole command.
- By default, only the values on first row are used for performing checks.
- This callable should return a list of records on
- which the checks have to be performed.
-
-
- .. note:: All fields in common with template fields of
- QuboleOperator and ValueCheckOperator are template-supported.
- """
-
- template_fields = QuboleOperator.template_fields + ValueCheckOperator.template_fields
- template_ext = QuboleOperator.template_ext
- ui_fgcolor = '#000'
-
- @apply_defaults
- def __init__(self, pass_value, tolerance=None,
- qubole_conn_id="qubole_default", *args, **kwargs):
-
- sql = get_sql_from_qbol_cmd(kwargs)
- super(QuboleValueCheckOperator, self).__init__(
- qubole_conn_id=qubole_conn_id,
- sql=sql, pass_value=pass_value, tolerance=tolerance,
- *args, **kwargs)
-
- self.on_failure_callback = QuboleCheckHook.handle_failure_retry
- self.on_retry_callback = QuboleCheckHook.handle_failure_retry
-
- def execute(self, context=None):
- try:
- self.hook = self.get_hook(context=context)
- super(QuboleValueCheckOperator, self).execute(context=context)
- except AirflowException as e:
- handle_airflow_exception(e, self.get_hook())
-
- def get_db_hook(self):
- return self.get_hook()
-
- def get_hook(self, context=None):
- if hasattr(self, 'hook') and (self.hook is not None):
- return self.hook
- else:
- return QuboleCheckHook(context=context, *self.args, **self.kwargs)
-
- def __getattribute__(self, name):
- if name in QuboleValueCheckOperator.template_fields:
- if name in self.kwargs:
- return self.kwargs[name]
- else:
- return ''
- else:
- return object.__getattribute__(self, name)
-
- def __setattr__(self, name, value):
- if name in QuboleValueCheckOperator.template_fields:
- self.kwargs[name] = value
- else:
- object.__setattr__(self, name, value)
-
-
-def get_sql_from_qbol_cmd(params):
- sql = ''
- if 'query' in params:
- sql = params['query']
- elif 'sql' in params:
- sql = params['sql']
- return sql
-
-
-def handle_airflow_exception(airflow_exception, hook):
- cmd = hook.cmd
- if cmd is not None:
- if cmd.is_success:
- qubole_command_results = hook.get_query_results()
- qubole_command_id = cmd.id
- exception_message = '\nQubole Command Id: {qubole_command_id}' \
- '\nQubole Command Results:' \
- '\n{qubole_command_results}'.format(**locals())
- raise AirflowException(str(airflow_exception) + exception_message)
- raise AirflowException(airflow_exception.message)
diff --git a/airflow/contrib/operators/qubole_operator.py b/airflow/contrib/operators/qubole_operator.py
deleted file mode 100755
index 82ee293b93..0000000000
--- a/airflow/contrib/operators/qubole_operator.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-from airflow.contrib.hooks.qubole_hook import QuboleHook
-
-
-class QuboleOperator(BaseOperator):
- """
- Execute tasks (commands) on QDS (https://qubole.com).
-
- :param qubole_conn_id: Connection id which consists of qds auth_token
- :type qubole_conn_id: str
-
- kwargs:
- :command_type: type of command to be executed, e.g. hivecmd, shellcmd, hadoopcmd
- :tags: array of tags to be assigned with the command
- :cluster_label: cluster label on which the command will be executed
- :name: name to be given to command
- :notify: whether to send email on command completion or not (default is False)
-
- **Arguments specific to command types**
-
- hivecmd:
- :query: inline query statement
- :script_location: s3 location containing query statement
- :sample_size: size of sample in bytes on which to run query
- :macros: macro values which were used in query
- prestocmd:
- :query: inline query statement
- :script_location: s3 location containing query statement
- :macros: macro values which were used in query
- hadoopcmd:
- :sub_commnad: must be one these ["jar", "s3distcp", "streaming"] followed by
- 1 or more args
- shellcmd:
- :script: inline command with args
- :script_location: s3 location containing query statement
- :files: list of files in s3 bucket as file1,file2 format. These files will be
- copied into the working directory where the qubole command is being
- executed.
- :archives: list of archives in s3 bucket as archive1,archive2 format. These
- will be unarchived intothe working directory where the qubole command is
- being executed
- :parameters: any extra args which need to be passed to script (only when
- script_location is supplied)
- pigcmd:
- :script: inline query statement (latin_statements)
- :script_location: s3 location containing pig query
- :parameters: any extra args which need to be passed to script (only when
- script_location is supplied
- sparkcmd:
- :program: the complete Spark Program in Scala, SQL, Command, R, or Python
- :cmdline: spark-submit command line, all required information must be specify
- in cmdline itself.
- :sql: inline sql query
- :script_location: s3 location containing query statement
- :language: language of the program, Scala, SQL, Command, R, or Python
- :app_id: ID of an Spark job server app
- :arguments: spark-submit command line arguments
- :user_program_arguments: arguments that the user program takes in
- :macros: macro values which were used in query
- dbtapquerycmd:
- :db_tap_id: data store ID of the target database, in Qubole.
- :query: inline query statement
- :macros: macro values which were used in query
- dbexportcmd:
- :mode: 1 (simple), 2 (advance)
- :hive_table: Name of the hive table
- :partition_spec: partition specification for Hive table.
- :dbtap_id: data store ID of the target database, in Qubole.
- :db_table: name of the db table
- :db_update_mode: allowinsert or updateonly
- :db_update_keys: columns used to determine the uniqueness of rows
- :export_dir: HDFS/S3 location from which data will be exported.
- :fields_terminated_by: hex of the char used as column separator in the dataset
- dbimportcmd:
- :mode: 1 (simple), 2 (advance)
- :hive_table: Name of the hive table
- :dbtap_id: data store ID of the target database, in Qubole.
- :db_table: name of the db table
- :where_clause: where clause, if any
- :parallelism: number of parallel db connections to use for extracting data
- :extract_query: SQL query to extract data from db. $CONDITIONS must be part
- of the where clause.
- :boundary_query: Query to be used get range of row IDs to be extracted
- :split_column: Column used as row ID to split data into ranges (mode 2)
-
- .. note:: Following fields are template-supported : ``query``, ``script_location``,
- ``sub_command``, ``script``, ``files``, ``archives``, ``program``, ``cmdline``,
- ``sql``, ``where_clause``, ``extract_query``, ``boundary_query``, ``macros``,
- ``tags``, ``name``, ``parameters``, ``dbtap_id``, ``hive_table``, ``db_table``,
- ``split_column``, ``note_id``, ``db_update_keys``, ``export_dir``,
- ``partition_spec``, ``qubole_conn_id``, ``arguments``, ``user_program_arguments``.
- You can also use ``.txt`` files for template driven use cases.
-
- .. note:: In QuboleOperator there is a default handler for task failures and retries,
- which generally kills the command running at QDS for the corresponding task
- instance. You can override this behavior by providing your own failure and retry
- handler in task definition.
- """
-
- template_fields = ('query', 'script_location', 'sub_command', 'script', 'files',
- 'archives', 'program', 'cmdline', 'sql', 'where_clause', 'tags',
- 'extract_query', 'boundary_query', 'macros', 'name', 'parameters',
- 'dbtap_id', 'hive_table', 'db_table', 'split_column', 'note_id',
- 'db_update_keys', 'export_dir', 'partition_spec', 'qubole_conn_id',
- 'arguments', 'user_program_arguments', 'cluster_label')
-
- template_ext = ('.txt',)
- ui_color = '#3064A1'
- ui_fgcolor = '#fff'
-
- @apply_defaults
- def __init__(self, qubole_conn_id="qubole_default", *args, **kwargs):
- self.args = args
- self.kwargs = kwargs
- self.kwargs['qubole_conn_id'] = qubole_conn_id
- super(QuboleOperator, self).__init__(*args, **kwargs)
-
- if self.on_failure_callback is None:
- self.on_failure_callback = QuboleHook.handle_failure_retry
-
- if self.on_retry_callback is None:
- self.on_retry_callback = QuboleHook.handle_failure_retry
-
- def execute(self, context):
- return self.get_hook().execute(context)
-
- def on_kill(self, ti=None):
- self.get_hook().kill(ti)
-
- def get_results(self, ti=None, fp=None, inline=True, delim=None, fetch=True):
- return self.get_hook().get_results(ti, fp, inline, delim, fetch)
-
- def get_log(self, ti):
- return self.get_hook().get_log(ti)
-
- def get_jobs_id(self, ti):
- return self.get_hook().get_jobs_id(ti)
-
- def get_hook(self):
- # Reinitiating the hook, as some template fields might have changed
- return QuboleHook(*self.args, **self.kwargs)
-
- def __getattribute__(self, name):
- if name in QuboleOperator.template_fields:
- if name in self.kwargs:
- return self.kwargs[name]
- else:
- return ''
- else:
- return object.__getattribute__(self, name)
-
- def __setattr__(self, name, value):
- if name in QuboleOperator.template_fields:
- self.kwargs[name] = value
- else:
- object.__setattr__(self, name, value)
diff --git a/airflow/contrib/operators/s3_list_operator.py b/airflow/contrib/operators/s3_list_operator.py
deleted file mode 100644
index b85691b005..0000000000
--- a/airflow/contrib/operators/s3_list_operator.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 airflow.hooks.S3_hook import S3Hook
-from airflow.models import BaseOperator
-from airflow.utils.decorators import apply_defaults
-
-
-class S3ListOperator(BaseOperator):
- """
- List all objects from the bucket with the given string prefix in name.
-
- This operator returns a python list with the name of objects which can be
- used by `xcom` in the downstream task.
-
- :param bucket: The S3 bucket where to find the objects. (templated)
- :type bucket: string
- :param prefix: Prefix string to filters the objects whose name begin with
- such prefix. (templated)
- :type prefix: string
- :param delimiter: the delimiter marks key hierarchy. (templated)
- :type delimiter: string
- :param aws_conn_id: The connection ID to use when connecting to S3 storage.
- :type aws_conn_id: string
-
- **Example**:
- The following operator would list all the files
- (excluding subfolders) from the S3
- ``customers/2018/04/`` key in the ``data`` bucket. ::
-
- s3_file = S3ListOperator(
- task_id='list_3s_files',
- bucket='data',
- prefix='customers/2018/04/',
- delimiter='/',
- aws_conn_id='aws_customers_conn'
- )
- """
- template_fields = ('bucket', 'prefix', 'delimiter')
- ui_color = '#ffd700'
-
- @apply_defaults
- def __init__(self,
- bucket,
- prefix='',
- delimiter='',
- aws_conn_id='aws_default',
- *args,
- **kwargs):
- super(S3ListOperator, self).__init__(*args, **kwargs)
- self.bucket = bucket
- self.prefix = prefix
- self.delimiter = delimiter
- self.aws_conn_id = aws_conn_id
-
- def execute(self, context):
- hook = S3Hook(aws_conn_id=self.aws_conn_id)
-
- self.log.info(
- 'Getting the list of files from bucket: {0} in prefix: {1} (Delimiter {2})'.
- format(self.bucket, self.prefix, self.delimiter))
-
- return hook.list_keys(
- bucket_name=self.bucket,
- prefix=self.prefix,
- delimiter=self.delimiter)
diff --git a/airflow/contrib/operators/s3_to_gcs_operator.py b/airflow/contrib/operators/s3_to_gcs_operator.py
deleted file mode 100644
index 2898af1071..0000000000
--- a/airflow/contrib/operators/s3_to_gcs_operator.py
+++ /dev/null
@@ -1,190 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# 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 tempfile import NamedTemporaryFile
-
-from airflow.contrib.hooks.gcs_hook import (GoogleCloudStorageHook,
- _parse_gcs_url)
-from airflow.contrib.operators.s3_list_operator import S3ListOperator
-from airflow.exceptions import AirflowException
-from airflow.hooks.S3_hook import S3Hook
-from airflow.utils.decorators import apply_defaults
-
-
-class S3ToGoogleCloudStorageOperator(S3ListOperator):
- """
- Synchronizes an S3 key, possibly a prefix, with a Google Cloud Storage
- destination path.
-
- :param bucket: The S3 bucket where to find the objects. (templated)
- :type bucket: string
- :param prefix: Prefix string which filters objects whose name begin with
- such prefix. (templated)
- :type prefix: string
- :param delimiter: the delimiter marks key hierarchy. (templated)
- :type delimiter: string
- :param aws_conn_id: The source S3 connection
- :type aws_conn_id: string
- :param dest_gcs_conn_id: The destination connection ID to use
- when connecting to Google Cloud Storage.
- :type dest_gcs_conn_id: string
- :param dest_gcs: The destination Google Cloud Storage bucket and prefix
- where you want to store the files. (templated)
- :type dest_gcs: string
- :param delegate_to: The account to impersonate, if any.
- For this to work, the service account making the request must have
- domain-wide delegation enabled.
- :type delegate_to: string
- :param replace: Whether you want to replace existing destination files
- or not.
- :type replace: bool
-
-
- **Example**:
- .. code-block:: python
- s3_to_gcs_op = S3ToGoogleCloudStorageOperator(
- task_id='s3_to_gcs_example',
- bucket='my-s3-bucket',
- prefix='data/customers-201804',
- dest_gcs_conn_id='google_cloud_default',
- dest_gcs='gs://my.gcs.bucket/some/customers/',
- replace=False,
- dag=my-dag)
-
- Note that ``bucket``, ``prefix``, ``delimiter`` and ``dest_gcs`` are
- templated, so you can use variables in them if you wish.
- """
-
- template_fields = ('bucket', 'prefix', 'delimiter', 'dest_gcs')
- ui_color = '#e09411'
-
- @apply_defaults
- def __init__(self,
- bucket,
- prefix='',
- delimiter='',
- aws_conn_id='aws_default',
- dest_gcs_conn_id=None,
- dest_gcs=None,
- delegate_to=None,
- replace=False,
- *args,
- **kwargs):
-
- super(S3ToGoogleCloudStorageOperator, self).__init__(
- bucket=bucket,
- prefix=prefix,
- delimiter=delimiter,
- aws_conn_id=aws_conn_id,
- *args,
- **kwargs)
- self.dest_gcs_conn_id = dest_gcs_conn_id
- self.dest_gcs = dest_gcs
- self.delegate_to = delegate_to
- self.replace = replace
-
- if dest_gcs and not self._gcs_object_is_directory(self.dest_gcs):
- self.log.info(
- 'Destination Google Cloud Storage path is not a valid '
- '"directory", define a path that ends with a slash "/" or '
(This diff was longer than 20,000 lines, and has been truncated...)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
> Split hooks and operators out from core Airflow
> -----------------------------------------------
>
> Key: AIRFLOW-2732
> URL: https://issues.apache.org/jira/browse/AIRFLOW-2732
> Project: Apache Airflow
> Issue Type: Improvement
> Components: ci, hooks, operators, tests, travis
> Reporter: Taylor Edmiston
> Assignee: Taylor Edmiston
> Priority: Major
>
> The goal of this issue is to split out hooks and operators from the core incubator-airflow repo into a second repo to facilitate:
> # faster CI builds of core
> # more frequent releases of hooks & operators by decoupling them from core
> We have discussed this work with Max last month and it's also been mentioned on the Airflow Dev list in the thread "Apache Airflow 1.10.0b3".
--
This message was sent by Atlassian JIRA
(v7.6.3#76005)