Connect to a cluster using the DC aware policy

dc_aware_policy.py

#!/usr/bin/env python

# Copyright 2014 DataStax

# 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.

# Using the DataStax python driver: https://github.com/datastax/python-driver

import cassandra
from cassandra.cluster import Cluster
from cassandra.policies import *

# this is who to contact initially, once we're in we'll auto discover the rest of the hosts
SEED_HOSTS = ['127.0.0.1']

# the name of our local datacenter
DATACENTER = 'US_EAST'

class CustomRetryPolicy(RetryPolicy):

    def on_write_timeout(self, query, consistency, write_type,
                         required_responses, received_responses, retry_num):

        # retry at most 5 times regardless of query type
        if retry_num >= 5:
            return (self.RETHROW, None)

        return (self.RETRY, consistency)


def connect(seeds, keyspace, datacenter=None, default_retry_policy=CustomRetryPolicy(), port=9042):

    load_balancing_policy = None
    if datacenter:
        # If you are using multiple datacenters it's important to use
        # the DCAwareRoundRobinPolicy. If not then the client will
        # make cross DC connections. This defaults to round robin
        # which means round robin across all nodes irrespective of
        # data center.
        load_balancing_policy = DCAwareRoundRobinPolicy(local_dc=datacenter)


    cluster = Cluster(contact_points=seeds,
                      port=port,
                      default_retry_policy=default_retry_policy,
                      load_balancing_policy=load_balancing_policy)

    return cluster.connect(keyspace)


session = connect(SEED_HOSTS, "system", datacenter=DATACENTER)
rows = session.execute("SELECT * FROM schema_keyspaces;")
for i in rows:
    print repr(i)