diff --git a/.travis.yml b/.travis.yml index 41b71b1..1bd3729 100644 --- a/.travis.yml +++ b/.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 diff --git a/cassandra_migrate/cli.py b/cassandra_migrate/cli.py index 634b45a..2d33d80 100644 --- a/cassandra_migrate/cli.py +++ b/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) diff --git a/cassandra_migrate/migrator.py b/cassandra_migrate/migrator.py index 2db6d9e..0e2c45d 100644 --- a/cassandra_migrate/migrator.py +++ b/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, diff --git a/tox.ini b/tox.ini index c5f0a1c..584d1f5 100644 --- a/tox.ini +++ b/tox.ini @@ -1,7 +1,6 @@ [tox] envlist = py27 - py34 py35 py36 pypy3