| srikanth | 116e6e8 | 2014-08-19 07:22:37 -0700 | [diff] [blame] | 1 | # |
| 2 | # Copyright (c) 2013 Big Switch Networks, Inc. |
| 3 | # |
| 4 | # Licensed under the Eclipse Public License, Version 1.0 (the |
| 5 | # "License"); you may not use this file except in compliance with the |
| 6 | # License. You may obtain a copy of the License at |
| 7 | # |
| 8 | # http://www.eclipse.org/legal/epl-v10.html |
| 9 | # |
| 10 | # Unless required by applicable law or agreed to in writing, software |
| 11 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or |
| 13 | # implied. See the License for the specific language governing |
| 14 | # permissions and limitations under the License. |
| 15 | # |
| 16 | |
| 17 | from thrift.transport import TTransport |
| 18 | from thrift.transport import TSocket |
| 19 | from thrift.protocol import TBinaryProtocol |
| 20 | from cassandra import Cassandra |
| 21 | from cassandra.ttypes import * |
| 22 | |
| 23 | # FIXME: This class is derived from the django_cassandra backend. |
| 24 | # Should refactor this code better. |
| 25 | |
| 26 | class CassandraConnection(object): |
| 27 | def __init__(self, host, port, keyspace, user, password): |
| 28 | self.host = host |
| 29 | self.port = port |
| 30 | self.keyspace = keyspace |
| 31 | self.user = user |
| 32 | self.password = password |
| 33 | self.transport = None |
| 34 | self.client = None |
| 35 | self.keyspace_set = False |
| 36 | self.logged_in = False |
| 37 | |
| 38 | def set_keyspace(self): |
| 39 | if not self.keyspace_set: |
| 40 | self.client.set_keyspace(self.keyspace) |
| 41 | self.keyspace_set = True |
| 42 | |
| 43 | def login(self): |
| 44 | # TODO: This user/password auth code hasn't been tested |
| 45 | if not self.logged_in: |
| 46 | if self.user: |
| 47 | credentials = {'username': self.user, 'password': self.password} |
| 48 | self.client.login(AuthenticationRequest(credentials)) |
| 49 | self.logged_in = True |
| 50 | |
| 51 | def connect(self, set_keyspace=False, login=False): |
| 52 | if self.transport == None: |
| 53 | # Create the client connection to the Cassandra daemon |
| 54 | socket = TSocket.TSocket(self.host, int(self.port)) |
| 55 | transport = TTransport.TFramedTransport(TTransport.TBufferedTransport(socket)) |
| 56 | protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport) |
| 57 | transport.open() |
| 58 | self.transport = transport |
| 59 | self.client = Cassandra.Client(protocol) |
| 60 | |
| 61 | if login: |
| 62 | self.login() |
| 63 | |
| 64 | if set_keyspace: |
| 65 | self.set_keyspace() |
| 66 | |
| 67 | def disconnect(self): |
| 68 | if self.transport != None: |
| 69 | self.transport.close() |
| 70 | self.transport = None |
| 71 | self.client = None |
| 72 | self.keyspace_set = False |
| 73 | self.logged_in = False |
| 74 | |
| 75 | def is_connected(self): |
| 76 | return self.transport != None |
| 77 | |
| 78 | def reconnect(self): |
| 79 | self.disconnect() |
| 80 | self.connect(True, True) |
| 81 | |
| 82 | def get_client(self): |
| 83 | if self.client == None: |
| 84 | self.connect(True,True) |
| 85 | return self.client |