druid_hook — Airflow Documentation

# Source code for druid_hook

    # -*- 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 __future__ import print_function
    import logging
    import json
    import time

    from pydruid.client import PyDruid
    import requests

    from airflow.hooks.base_hook import BaseHook
    from airflow.exceptions import AirflowException

    LOAD_CHECK_INTERVAL = 5
    DEFAULT_TARGET_PARTITION_SIZE = 5000000

    class AirflowDruidLoadException(AirflowException):
        pass

    class DruidHook(BaseHook):
        '''
        Interact with druid.
        '''

        def __init__(
                self,
                druid_query_conn_id='druid_query_default',
                druid_ingest_conn_id='druid_ingest_default'):
            self.druid_query_conn_id = druid_query_conn_id
            self.druid_ingest_conn_id = druid_ingest_conn_id
            self.header = {'content-type': 'application/json'}

        def get_conn(self):
            """
            Returns a druid connection object for query
            """
            conn = self.get_connection(self.druid_query_conn_id)
            return PyDruid(
                "http://{conn.host}:{conn.port}".format(**locals()),
                conn.extra_dejson.get('endpoint', ''))

        @property
        def ingest_post_url(self):
            conn = self.get_connection(self.druid_ingest_conn_id)
            host = conn.host
            port = conn.port
            endpoint = conn.extra_dejson.get('endpoint', '')
            return "http://{host}:{port}/{endpoint}".format(**locals())

        def get_ingest_status_url(self, task_id):
            post_url = self.ingest_post_url
            return "{post_url}/{task_id}/status".format(**locals())

        def construct_ingest_query(
                self, datasource, static_path, ts_dim, columns, metric_spec,
                intervals, num_shards, target_partition_size,
                query_granularity="NONE", segment_granularity="DAY",
                hadoop_dependency_coordinates=None):
            """
            Builds an ingest query for an HDFS TSV load.

            :param datasource: target datasource in druid
            :param columns: list of all columns in the TSV, in the right order
            """

            # backward compatibilty for num_shards, but target_partition_size is the default setting
            # and overwrites the num_shards
            if target_partition_size == -1:
                if num_shards == -1:
                    target_partition_size = DEFAULT_TARGET_PARTITION_SIZE
            else:
                num_shards = -1

            metric_names = [
                m['fieldName'] for m in metric_spec if m['type'] != 'count']
            dimensions = [c for c in columns if c not in metric_names and c != ts_dim]
            ingest_query_dict = {
                "type": "index_hadoop",
                "spec": {
                    "dataSchema": {
                        "metricsSpec": metric_spec,
                        "granularitySpec": {
                            "queryGranularity": query_granularity,
                            "intervals": intervals,
                            "type": "uniform",
                            "segmentGranularity": segment_granularity,
                        },
                        "parser": {
                            "type": "string",
                            "parseSpec": {
                                "columns": columns,
                                "dimensionsSpec": {
                                    "dimensionExclusions": [],
                                    "dimensions": dimensions,  # list of names
                                    "spatialDimensions": []
                                },
                                "timestampSpec": {
                                    "column": ts_dim,
                                    "format": "auto"
                                },
                                "format": "tsv"
                            }
                        },
                        "dataSource": datasource
                    },
                    "tuningConfig": {
                        "type": "hadoop",
                        "jobProperties": {
                            "mapreduce.job.user.classpath.first": "false",
                            "mapreduce.map.output.compress": "false",
                            "mapreduce.output.fileoutputformat.compress": "false",
                        },
                        "partitionsSpec": {
                            "type": "hashed",
                            "targetPartitionSize": target_partition_size,
                            "numShards": num_shards,
                        },
                    },
                    "ioConfig": {
                        "inputSpec": {
                            "paths": static_path,
                            "type": "static"
                        },
                        "type": "hadoop"
                    }
                }
            }
            if hadoop_dependency_coordinates:
                ingest_query_dict[
                    'hadoopDependencyCoordinates'] = hadoop_dependency_coordinates

            return json.dumps(ingest_query_dict, indent=4)

        def send_ingest_query(
                self, datasource, static_path, ts_dim, columns, metric_spec,
                intervals, num_shards, target_partition_size, query_granularity, segment_granularity,
                hadoop_dependency_coordinates=None):
            query = self.construct_ingest_query(
                datasource, static_path, ts_dim, columns,
                metric_spec, intervals, num_shards, target_partition_size,
                query_granularity, segment_granularity, hadoop_dependency_coordinates)
            r = requests.post(
                self.ingest_post_url, headers=self.header, data=query)
            logging.info(self.ingest_post_url)
            logging.info(query)
            logging.info(r.text)
            d = json.loads(r.text)
            if "task" not in d:
                raise AirflowDruidLoadException(
                    "[Error]: Ingesting data to druid failed.")
            return d["task"]

        def load_from_hdfs(
                self, datasource, static_path,  ts_dim, columns,
                intervals, num_shards, target_partition_size, query_granularity, segment_granularity,
                metric_spec=None, hadoop_dependency_coordinates=None):
            """
            load data to druid from hdfs

            :param ts_dim: The column name to use as a timestamp
            :param metric_spec: A list of dictionaries
            """
            task_id = self.send_ingest_query(
                datasource, static_path, ts_dim, columns, metric_spec,
                intervals, num_shards, target_partition_size, query_granularity, segment_granularity,
                hadoop_dependency_coordinates)
            status_url = self.get_ingest_status_url(task_id)
            while True:
                r = requests.get(status_url)
                d = json.loads(r.text)
                if d['status']['status'] == 'FAILED':
                    logging.error(d)
                    raise AirflowDruidLoadException(
                        "[Error]: Ingesting data to druid failed.")
                elif d['status']['status'] == 'SUCCESS':
                    break
                time.sleep(LOAD_CHECK_INTERVAL)