You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by fo...@apache.org on 2018/06/12 20:39:33 UTC
incubator-airflow git commit: [AIRFLOW-2585] Fix several bugs in
CassandraHook and CassandraToGCSOperator
Repository: incubator-airflow
Updated Branches:
refs/heads/master 5bb547a85 -> 48fccefab
[AIRFLOW-2585] Fix several bugs in CassandraHook and CassandraToGCSOperator
Closes #3483 from jgao54/fix-cassandra-hook-bug
Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/48fccefa
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/48fccefa
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/48fccefa
Branch: refs/heads/master
Commit: 48fccefaba35f16dfff10b7629f64844bb850553
Parents: 5bb547a
Author: Joy Gao <Jo...@apache.org>
Authored: Tue Jun 12 22:39:26 2018 +0200
Committer: Fokko Driesprong <fo...@godatadriven.com>
Committed: Tue Jun 12 22:39:26 2018 +0200
----------------------------------------------------------------------
airflow/contrib/hooks/cassandra_hook.py | 100 ++++++++++++++++---
airflow/contrib/operators/cassandra_to_gcs.py | 20 ++--
tests/contrib/hooks/test_cassandra_hook.py | 70 ++++++++++++-
.../operators/test_cassandra_to_gcs_operator.py | 13 +--
4 files changed, 172 insertions(+), 31 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/48fccefa/airflow/contrib/hooks/cassandra_hook.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/hooks/cassandra_hook.py b/airflow/contrib/hooks/cassandra_hook.py
index 90046a8..6fce98b 100644
--- a/airflow/contrib/hooks/cassandra_hook.py
+++ b/airflow/contrib/hooks/cassandra_hook.py
@@ -19,8 +19,7 @@
from cassandra.cluster import Cluster
from cassandra.policies import (RoundRobinPolicy, DCAwareRoundRobinPolicy,
- TokenAwarePolicy, HostFilterPolicy,
- WhiteListRoundRobinPolicy)
+ TokenAwarePolicy, WhiteListRoundRobinPolicy)
from cassandra.auth import PlainTextAuthProvider
from airflow.hooks.base_hook import BaseHook
@@ -31,12 +30,45 @@ 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. Load_alancing_policy, ssl_options, cql_version can be specified
- in the extra field of the connection.
-
- For details of the Cluster config, see cassandra.cluster for more details.
+ 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)
@@ -52,7 +84,9 @@ class CassandraHook(BaseHook, LoggingMixin):
conn_config['auth_provider'] = PlainTextAuthProvider(
username=conn.login, password=conn.password)
- lb_policy = self.get_policy(conn.extra_dejson.get('load_balancing_policy', None))
+ 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
@@ -66,23 +100,59 @@ class CassandraHook(BaseHook, LoggingMixin):
self.cluster = Cluster(**conn_config)
self.keyspace = conn.schema
+ self.session = None
def get_conn(self):
"""
- Returns a cassandra connection object
+ Returns a cassandra Session object
"""
+ if self.session:
+ return self.session
return self.cluster.connect(self.keyspace)
def get_cluster(self):
return self.cluster
- @classmethod
- def get_policy(cls, policy_name):
+ 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,
- 'TokenAwarePolicy': TokenAwarePolicy,
- 'HostFilterPolicy': HostFilterPolicy,
'WhiteListRoundRobinPolicy': WhiteListRoundRobinPolicy,
+ 'TokenAwarePolicy': TokenAwarePolicy,
}
- return policies.get(policy_name)
+
+ 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)
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/48fccefa/airflow/contrib/operators/cassandra_to_gcs.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/operators/cassandra_to_gcs.py b/airflow/contrib/operators/cassandra_to_gcs.py
index b4e216d..211444b 100644
--- a/airflow/contrib/operators/cassandra_to_gcs.py
+++ b/airflow/contrib/operators/cassandra_to_gcs.py
@@ -98,11 +98,13 @@ class CassandraToGoogleCloudStorageOperator(BaseOperator):
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': 'STRING',
+ 'UUIDType': 'BYTES',
'BooleanType': 'BOOL',
'ByteType': 'INTEGER',
'AsciiType': 'STRING',
@@ -142,12 +144,15 @@ class CassandraToGoogleCloudStorageOperator(BaseOperator):
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.
"""
- hook = CassandraHook(cassandra_conn_id=self.cassandra_conn_id)
- session = hook.get_conn()
+ self.hook = CassandraHook(cassandra_conn_id=self.cassandra_conn_id)
+ session = self.hook.get_conn()
cursor = session.execute(self.cql)
return cursor
@@ -221,11 +226,10 @@ class CassandraToGoogleCloudStorageOperator(BaseOperator):
elif isinstance(value, (text_type, int, float, bool, dict)):
return value
elif isinstance(value, binary_type):
- encoded_value = b64encode(value)
- if PY3:
- encoded_value = encoded_value.decode('ascii')
- return encoded_value
- elif isinstance(value, (datetime, Date, UUID)):
+ 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)
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/48fccefa/tests/contrib/hooks/test_cassandra_hook.py
----------------------------------------------------------------------
diff --git a/tests/contrib/hooks/test_cassandra_hook.py b/tests/contrib/hooks/test_cassandra_hook.py
index 42afd9e..fd9e93c 100644
--- a/tests/contrib/hooks/test_cassandra_hook.py
+++ b/tests/contrib/hooks/test_cassandra_hook.py
@@ -19,11 +19,14 @@
import unittest
import mock
+from mock import patch
from airflow import configuration
from airflow.contrib.hooks.cassandra_hook import CassandraHook
from cassandra.cluster import Cluster
-from cassandra.policies import TokenAwarePolicy
+from cassandra.policies import (
+ TokenAwarePolicy, RoundRobinPolicy, DCAwareRoundRobinPolicy, WhiteListRoundRobinPolicy
+)
from airflow import models
from airflow.utils import db
@@ -35,7 +38,7 @@ class CassandraHookTest(unittest.TestCase):
models.Connection(
conn_id='cassandra_test', conn_type='cassandra',
host='host-1,host-2', port='9042', schema='test_keyspace',
- extra='{"load_balancing_policy":"TokenAwarePolicy"'))
+ extra='{"load_balancing_policy":"TokenAwarePolicy"}'))
def test_get_conn(self):
with mock.patch.object(Cluster, "connect") as mock_connect, \
@@ -51,6 +54,69 @@ class CassandraHookTest(unittest.TestCase):
self.assertEqual(cluster.port, 9042)
self.assertTrue(isinstance(cluster.load_balancing_policy, TokenAwarePolicy))
+ def test_get_lb_policy(self):
+ # test LB policies with no args
+ self._assert_get_lb_policy('RoundRobinPolicy', {}, RoundRobinPolicy)
+ self._assert_get_lb_policy('DCAwareRoundRobinPolicy', {}, DCAwareRoundRobinPolicy)
+ self._assert_get_lb_policy('TokenAwarePolicy', {}, TokenAwarePolicy,
+ expected_child_policy_type=RoundRobinPolicy)
+
+ # test DCAwareRoundRobinPolicy with args
+ self._assert_get_lb_policy('DCAwareRoundRobinPolicy',
+ {'local_dc': 'foo', 'used_hosts_per_remote_dc': '3'},
+ DCAwareRoundRobinPolicy)
+
+ # test WhiteListRoundRobinPolicy with args
+ fake_addr_info = [['family', 'sockettype', 'proto', 'canonname', ('2606:2800:220:1:248:1893:25c8:1946', 80, 0, 0)]] # noqa
+ with patch('socket.getaddrinfo', return_value=fake_addr_info):
+ self._assert_get_lb_policy('WhiteListRoundRobinPolicy',
+ {'hosts': ['host1', 'host2']},
+ WhiteListRoundRobinPolicy)
+
+ # test TokenAwarePolicy with args
+ with patch('socket.getaddrinfo', return_value=fake_addr_info):
+ self._assert_get_lb_policy('TokenAwarePolicy',
+ {'child_load_balancing_policy': 'WhiteListRoundRobinPolicy', # noqa
+ 'child_load_balancing_policy_args': {'hosts': ['host-1', 'host-2']}}, # noqa
+ TokenAwarePolicy,
+ expected_child_policy_type=WhiteListRoundRobinPolicy) # noqa
+
+ # test invalid policy name should default to RoundRobinPolicy
+ self._assert_get_lb_policy('DoesNotExistPolicy', {}, RoundRobinPolicy)
+
+ # test invalid child policy name should default child policy to RoundRobinPolicy
+ self._assert_get_lb_policy('TokenAwarePolicy', {}, TokenAwarePolicy,
+ expected_child_policy_type=RoundRobinPolicy)
+ self._assert_get_lb_policy('TokenAwarePolicy',
+ {'child_load_balancing_policy': 'DoesNotExistPolicy'},
+ TokenAwarePolicy,
+ expected_child_policy_type=RoundRobinPolicy)
+
+ # test host not specified for WhiteListRoundRobinPolicy should throw exception
+ self._assert_get_lb_policy('WhiteListRoundRobinPolicy',
+ {},
+ WhiteListRoundRobinPolicy,
+ should_throw=True)
+ self._assert_get_lb_policy('TokenAwarePolicy',
+ {'child_load_balancing_policy': 'WhiteListRoundRobinPolicy'}, # noqa
+ TokenAwarePolicy,
+ expected_child_policy_type=RoundRobinPolicy,
+ should_throw=True)
+
+ def _assert_get_lb_policy(self, policy_name, policy_args, expected_policy_type,
+ expected_child_policy_type=None,
+ should_throw=False):
+ thrown = False
+ try:
+ policy = CassandraHook.get_lb_policy(policy_name, policy_args)
+ self.assertTrue(isinstance(policy, expected_policy_type))
+ if expected_child_policy_type:
+ self.assertTrue(isinstance(policy._child_policy,
+ expected_child_policy_type))
+ except Exception:
+ thrown = True
+ self.assertEqual(should_throw, thrown)
+
if __name__ == '__main__':
unittest.main()
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/48fccefa/tests/contrib/operators/test_cassandra_to_gcs_operator.py
----------------------------------------------------------------------
diff --git a/tests/contrib/operators/test_cassandra_to_gcs_operator.py b/tests/contrib/operators/test_cassandra_to_gcs_operator.py
index add115f..42f8195 100644
--- a/tests/contrib/operators/test_cassandra_to_gcs_operator.py
+++ b/tests/contrib/operators/test_cassandra_to_gcs_operator.py
@@ -61,18 +61,19 @@ class CassandraToGCSTest(unittest.TestCase):
self.assertEquals(op.convert_value('date', date), str(date_str))
import uuid
+ from base64 import b64encode
test_uuid = uuid.uuid4()
- self.assertEquals(op.convert_value('uuid', test_uuid), str(test_uuid))
-
- from decimal import Decimal
- d = Decimal(1.0)
- self.assertEquals(op.convert_value('decimal', d), float(d))
+ encoded_uuid = b64encode(test_uuid.bytes).decode('ascii')
+ self.assertEquals(op.convert_value('uuid', test_uuid), encoded_uuid)
- from base64 import b64encode
b = b'abc'
encoded_b = b64encode(b).decode('ascii')
self.assertEquals(op.convert_value('binary', b), encoded_b)
+ from decimal import Decimal
+ d = Decimal(1.0)
+ self.assertEquals(op.convert_value('decimal', d), float(d))
+
from cassandra.util import Time
time = Time(0)
self.assertEquals(op.convert_value('time', time), '00:00:00')