allow passing load-balancing policy to the cassandra driver

.travis.yml

@@ -8,7 +8,6 @@ matrix:
 
   include:
     - python: "2.7"
-    - python: "3.4"
     - python: "3.5"
     - python: "3.6"
     - python: pypy3.5-5.8.0

cassandra_migrate/cli.py

@@ -46,6 +46,23 @@ def main():
                         help='Path to configuration file')
     parser.add_argument('-m', '--profile', default='dev',
                         help='Name of keyspace profile to use')
+    parser.add_argument('-l', '--load-balancing-policy',
+                        choices=('round-robin', 'dc-aware-round-robin'),
+                        default='round-robin',
+                        help="""
+                        Cassandra connection load balancing policy. if
+                        dc-aware-round-robin is passed, you must set
+                        --local-dc. """)
+    parser.add_argument('-L', '--local-dc',
+                        help='Local datacenter name')
+    parser.add_argument('--used-hosts-per-remote-dc',
+                        type=int,
+                        default=0,
+                        help="""
+                        If --load-balancing-policy is dc-aware-round-robin,
+                        this controls how many hosts to use from non-local
+                        datacenters.  Only set this if you know what it
+                        means.""")
     parser.add_argument('-s', '--ssl-cert', default=None,
                         help="""
                         File path of .pem or .crt containing certificate of the
@@ -134,12 +151,23 @@ def main():
 
         print(os.path.basename(new_path))
     else:
-        with Migrator(config=config, profile=opts.profile,
-                      hosts=opts.hosts.split(','), port=opts.port,
-                      user=opts.user, password=opts.password,
-                      host_cert_path=opts.ssl_cert,
-                      client_key_path=opts.ssl_client_private_key,
-                      client_cert_path=opts.ssl_client_cert) as migrator:
+        if (
+            opts.load_balancing_policy == 'dc-aware-round-robin' and
+            not opts.local_dc
+        ):
+            parser.error('must pass --local-dc when setting --load-balancing-'
+                         'policy to dc-aware-round-robin')
+        with Migrator(
+            config=config, profile=opts.profile,
+            hosts=opts.hosts.split(','), port=opts.port,
+            user=opts.user, password=opts.password,
+            host_cert_path=opts.ssl_cert,
+            client_key_path=opts.ssl_client_private_key,
+            client_cert_path=opts.ssl_client_cert,
+            load_balancing_policy=opts.load_balancing_policy,
+            local_dc=opts.local_dc,
+            used_hosts_per_remote_dc=opts.used_hosts_per_remote_dc,
+        ) as migrator:
             cmd_method = getattr(migrator, opts.action)
             if not callable(cmd_method):
                 print('Error: invalid command', file=sys.stderr)

cassandra_migrate/migrator.py

@@ -21,6 +21,7 @@
 from cassandra.auth import PlainTextAuthProvider
 from cassandra_migrate import (Migration, FailedMigration, InconsistentState,
                                UnknownMigration, ConcurrentMigration)
+from cassandra.policies import (RoundRobinPolicy, DCAwareRoundRobinPolicy)
 from cassandra_migrate.cql import CqlSplitter
 
 
@@ -110,13 +111,18 @@ class Migrator(object):
     - user, password: authentication options. May be None to not use it.
     - hosts: comma-separated list of contact points
     - port: connection port
+    - load_balancing_policy: "round-robin" or "dc-aware-round-robin"
+    - local_dc: name of local datacenter
+    - used_hosts_per_remote_dc: number of remote hosts to use
     """
 
     logger = logging.getLogger("Migrator")
 
     def __init__(self, config, profile='dev', hosts=['127.0.0.1'], port=9042,
                  user=None, password=None, host_cert_path=None,
-                 client_key_path=None, client_cert_path=None):
+                 client_key_path=None, client_cert_path=None,
+                 load_balancing_policy='round-robin',
+                 local_dc=None, used_hosts_per_remote_dc=0):
         self.config = config
 
         try:
@@ -137,10 +143,19 @@ def __init__(self, config, profile='dev', hosts=['127.0.0.1'], port=9042,
         else:
             ssl_options = None
 
+        if load_balancing_policy == 'round-robin':
+            cluster_load_balancing_policy = RoundRobinPolicy()
+        elif load_balancing_policy == 'dc-aware-round-robin':
+            cluster_load_balancing_policy = DCAwareRoundRobinPolicy(
+                local_dc=local_dc,
+                use_hosts_per_remote_dc=used_hosts_per_remote_dc,
+            )
+
         self.cluster = Cluster(
             contact_points=hosts,
             port=port,
             auth_provider=auth_provider,
+            load_balancing_policy=cluster_load_balancing_policy,
             max_schema_agreement_wait=300,
             control_connection_timeout=10,
             connect_timeout=30,

tox.ini

@@ -1,7 +1,6 @@
 [tox]
 envlist =
     py27
-    py34
     py35
     py36
     pypy3