Source code for airflow.providers.apache.cassandra.hooks.cassandra

#
# 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 hook to integrate with Apache Cassandra."""

from typing import Any, Dict, Union

from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster, Session
from cassandra.policies import (
    DCAwareRoundRobinPolicy,
    RoundRobinPolicy,
    TokenAwarePolicy,
    WhiteListRoundRobinPolicy,
)

from airflow.hooks.base import BaseHook
from airflow.utils.log.logging_mixin import LoggingMixin

[docs]Policy = Union[DCAwareRoundRobinPolicy, RoundRobinPolicy, TokenAwarePolicy, WhiteListRoundRobinPolicy]
[docs]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. """
[docs] conn_name_attr = 'cassandra_conn_id'
[docs] default_conn_name = 'cassandra_default'
[docs] conn_type = 'cassandra'
[docs] hook_name = 'Cassandra'
def __init__(self, cassandra_conn_id: str = default_conn_name): super().__init__() 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 protocol_version = conn.extra_dejson.get('protocol_version', None) if protocol_version: conn_config['protocol_version'] = protocol_version self.cluster = Cluster(**conn_config) self.keyspace = conn.schema self.session = None
[docs] def get_conn(self) -> Session: """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
[docs] def get_cluster(self) -> Cluster: """Returns Cassandra cluster.""" return self.cluster
[docs] def shutdown_cluster(self) -> None: """Closes all sessions and connections associated with this Cluster.""" if not self.cluster.is_shutdown: self.cluster.shutdown()
@staticmethod
[docs] def get_lb_policy(policy_name: str, policy_args: Dict[str, Any]) -> Policy: """ Creates load balancing policy. :param policy_name: Name of the policy to use. :param policy_args: Parameters for the policy. """ 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) # Fallback to default RoundRobinPolicy return RoundRobinPolicy()
[docs] def table_exists(self, table: str) -> bool: """ Checks if a table exists in Cassandra :param table: Target Cassandra table. Use dot notation to target a specific keyspace. """ 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
[docs] def record_exists(self, table: str, keys: Dict[str, str]) -> bool: """ Checks if a record exists in Cassandra :param table: Target Cassandra table. Use dot notation to target a specific keyspace. :param keys: The keys and their values to check the existence. """ keyspace = self.keyspace if '.' in table: keyspace, table = table.split('.', 1) ks_str = " AND ".join(f"{key}=%({key})s" for key in keys.keys()) query = f"SELECT * FROM {keyspace}.{table} WHERE {ks_str}" try: result = self.get_conn().execute(query, keys) return result.one() is not None except Exception: return False

Was this entry helpful?