Python cassandra.cluster.Cluster() Examples

def test_cassandra_connect(instrument, elasticapm_client, cassandra_cluster):
    elasticapm_client.begin_transaction("transaction.test")
    sess = cassandra_cluster.connect()
    elasticapm_client.end_transaction("test")

    transactions = elasticapm_client.events[TRANSACTION]
    span = elasticapm_client.spans_for_transaction(transactions[0])[0]

    assert span["type"] == "db"
    assert span["subtype"] == "cassandra"
    assert span["action"] == "connect"
    assert span["duration"] > 0
    assert span["name"] == "Cluster.connect"
    assert span["context"]["destination"] == {
        "address": socket.gethostbyname(os.environ.get("CASSANDRA_HOST", "localhost")),
        "port": 9042,
        "service": {"name": "cassandra", "resource": "cassandra", "type": "db"},
    } 

def get_session(cls, seeds, **kwargs):
            _ = kwargs
            if cls.__session is None:
                # Allow dependency injection
                session = kwargs.get("session")
                if session is None:
                    cluster = c_cluster.Cluster(seeds)
                    session = cluster.connect()
                keyspace = kwargs["keyspace"]
                replication = kwargs["replication"]
                if kwargs.get("drop_keyspace", False):
                    session.execute(cls.QUERY_DROP_KEYSPACE.format(keyspace))
                session.execute(cls.QUERY_CREATE_KEYSPACE.format(
                    keyspace=keyspace,
                    replication=str(replication),
                ))
                session.set_keyspace(keyspace)
                cls.__session = session
            return cls.__session 

def __init__(self, cluster=Cluster(['127.0.0.1'], connect_timeout=30), keyspace='cstar_perf', email_notifications=False):
        """Instantiate DB model object for interacting with the C* backend.

        cluster - Python driver object for accessing Cassandra
        keyspace - the keyspace to use
        email_notifications - if True, perform email notifications for some actions. Defaults to False.
        """
        log.info("Initializing Model...")
        self.cluster = cluster if type(cluster) == Cluster else Cluster(cluster)
        self.keyspace = keyspace
        self.email_notifications = email_notifications
        self.__maybe_create_schema()
        self.__shared_session = self.get_session()
        ## Prepare statements:
        self.__prepared_statements = {}
        for name, stmt in Model.statements.items():
            # log.debug("Preparing statement: {stmt}".format(stmt=stmt))
            self.__prepared_statements[name] = self.get_session().prepare(stmt)

        ### ZeroMQ publisher for announcing jobs as they come in:
        zmq_context = zmq.Context()
        self.zmq_socket = zmq_context.socket(zmq.PUSH)
        self.zmq_socket.connect("tcp://127.0.0.1:5556")

        log.info("Model initialized") 

def __init__(self):
        self.MORPHL_SERVER_IP_ADDRESS = getenv('MORPHL_SERVER_IP_ADDRESS')
        self.MORPHL_CASSANDRA_USERNAME = getenv('MORPHL_CASSANDRA_USERNAME')
        self.MORPHL_CASSANDRA_PASSWORD = getenv('MORPHL_CASSANDRA_PASSWORD')
        self.MORPHL_CASSANDRA_KEYSPACE = getenv('MORPHL_CASSANDRA_KEYSPACE')

        self.CASS_REQ_TIMEOUT = 3600.0

        self.auth_provider = PlainTextAuthProvider(
            username=self.MORPHL_CASSANDRA_USERNAME,
            password=self.MORPHL_CASSANDRA_PASSWORD)
        self.cluster = Cluster(
            [self.MORPHL_SERVER_IP_ADDRESS], auth_provider=self.auth_provider)

        self.session = self.cluster.connect(self.MORPHL_CASSANDRA_KEYSPACE)
        self.session.row_factory = dict_factory
        self.session.default_fetch_size = 100

        self.prepare_statements() 

def context():
    global session
    if not session:
        cluster = Cluster()
        session = cluster.connect()

        # create the keyspace
        ks = 'test_cdm'
        if 'test_cdm' not in cluster.metadata.keyspaces:
            q = "create KEYSPACE {} WITH \
                    replication = {{'class': \
                    'SimpleStrategy', \
                    'replication_factor': 1}}".format(ks)
            session.execute(q)
        if not os.path.exists("./cache"):
            os.mkdir("./cache")
        session.set_keyspace(ks)

    return Context(root=os.getcwd(),
                   dataset="test",
                   session=session,
                   cache_dir="./cache",
                   ) 

def cloud_generator(test_split=0.2):
    """
    Generator for cloud use.
    test_split: Percentage of data to reserve for testing.
    Return X, y (input image and label)
    """
    # retrieve raw data
    cluster = Cluster(['10.44.27.2'])
    session = cluster.connect('healthcare_keyspace_v2')
    f_names = _get_image_filenames(test_split=test_split)
    while True:
        for x in f_names.iloc:
            query = "SELECT image_blob, is_label FROM processed_data"
            query = query + " WHERE image_filename='" + x[0] + "'" + "ALLOW FILTERING"
            rows = session.execute(query)
            X, y = _return_input_and_label(rows)
            yield X, y 

def __init__(self):
        self.MORPHL_SERVER_IP_ADDRESS = getenv('MORPHL_SERVER_IP_ADDRESS')
        self.MORPHL_CASSANDRA_USERNAME = getenv('MORPHL_CASSANDRA_USERNAME')
        self.MORPHL_CASSANDRA_PASSWORD = getenv('MORPHL_CASSANDRA_PASSWORD')
        self.MORPHL_CASSANDRA_KEYSPACE = getenv('MORPHL_CASSANDRA_KEYSPACE')

        self.QUERY = 'SELECT * FROM ga_chp_bq_predictions WHERE client_id = ? LIMIT 1'
        self.CASS_REQ_TIMEOUT = 3600.0

        self.auth_provider = PlainTextAuthProvider(
            username=self.MORPHL_CASSANDRA_USERNAME,
            password=self.MORPHL_CASSANDRA_PASSWORD)
        self.cluster = Cluster(
            [self.MORPHL_SERVER_IP_ADDRESS], auth_provider=self.auth_provider)

        self.session = self.cluster.connect(self.MORPHL_CASSANDRA_KEYSPACE)
        self.session.row_factory = dict_factory
        self.session.default_fetch_size = 100

        self.prepare_statements() 

def __init__(self):
        self.MORPHL_SERVER_IP_ADDRESS = getenv('MORPHL_SERVER_IP_ADDRESS')
        self.MORPHL_CASSANDRA_USERNAME = getenv('MORPHL_CASSANDRA_USERNAME')
        self.MORPHL_CASSANDRA_PASSWORD = getenv('MORPHL_CASSANDRA_PASSWORD')
        self.MORPHL_CASSANDRA_KEYSPACE = getenv('MORPHL_CASSANDRA_KEYSPACE')

        self.prep_stmt = {}

        template_for_prediction = 'INSERT INTO ga_chp_bq_predictions (client_id,prediction) VALUES (?,?)'
        template_for_predictions_by_date = 'INSERT INTO ga_chp_bq_predictions_by_prediction_date (prediction_date, client_id, prediction) VALUES (?,?,?)'
        template_for_predictions_statistics = 'UPDATE ga_chp_bq_predictions_statistics SET loyal=loyal+?, neutral=neutral+?, churning=churning+?, lost=lost+? WHERE prediction_date=?'


        self.CASS_REQ_TIMEOUT = 3600.0

        self.auth_provider = PlainTextAuthProvider(
            username=self.MORPHL_CASSANDRA_USERNAME,
            password=self.MORPHL_CASSANDRA_PASSWORD)
        self.cluster = Cluster(
            [self.MORPHL_SERVER_IP_ADDRESS], auth_provider=self.auth_provider)
        self.session = self.cluster.connect(self.MORPHL_CASSANDRA_KEYSPACE)

        self.prep_stmt['prediction'] = self.session.prepare(template_for_prediction)
        self.prep_stmt['predictions_by_date'] = self.session.prepare(template_for_predictions_by_date)
        self.prep_stmt['predictions_statistics'] = self.session.prepare(template_for_predictions_statistics) 

def __enter__(self):
        domsconfig = ConfigParser.RawConfigParser()
        domsconfig.readfp(pkg_resources.resource_stream(__name__, "domsconfig.ini"), filename='domsconfig.ini')

        cassHost = domsconfig.get("cassandra", "host")
        cassKeyspace = domsconfig.get("cassandra", "keyspace")
        cassDatacenter = domsconfig.get("cassandra", "local_datacenter")
        cassVersion = int(domsconfig.get("cassandra", "protocol_version"))

        dc_policy = DCAwareRoundRobinPolicy(cassDatacenter)
        token_policy = TokenAwarePolicy(dc_policy)

        self._cluster = Cluster([host for host in cassHost.split(',')], load_balancing_policy=token_policy,
                                protocol_version=cassVersion)

        self._session = self._cluster.connect(cassKeyspace)
        return self 

def init(args):
    global solr_connection
    solr_connection = SolrConnection(args.solr)
    global solr_collection
    solr_collection = solr_connection[args.collection]
    global SOLR_UNIQUE_KEY
    SOLR_UNIQUE_KEY = args.solrIdField

    dc_policy = RoundRobinPolicy()
    token_policy = TokenAwarePolicy(dc_policy)

    global cassandra_cluster
    cassandra_cluster = Cluster(contact_points=args.cassandra, port=args.cassandraPort,
                                protocol_version=int(args.cassandraProtocolVersion),
                                load_balancing_policy=token_policy)
    global cassandra_session
    cassandra_session = cassandra_cluster.connect(keyspace=args.cassandraKeyspace)

    global cassandra_table
    cassandra_table = args.cassandraTable 

def get_db(args):
    log = logging.getLogger("cassandra")
    patch_tables(args)
    try:
        cas_host, cas_port = args.cassandra.split(":")
    except ValueError:
        cas_host = args.cassandra
        cas_port = "9042"

    def get_cluster():
        return Cluster((cas_host,), port=int(cas_port),
                       load_balancing_policy=RoundRobinPolicy())
    cluster = get_cluster()
    log.info("Connecting to %s", args.cassandra)
    try:
        session = cluster.connect(args.keyspace)
    except NoHostAvailable:
        log.warning("Keyspace %s does not exist", args.keyspace)
        cluster = get_cluster()
        session = cluster.connect()
    return session 

def __cassandra_connect__(cassandra_instance):
        """
        Connect to the Cassandra DB - either a local one or the Zooniverse aws one. If unable to connect, re-try up to 10 times and then raise an error.

        Raises
        ------
        cassandra.cluster.NoHostAvailable
            If we are not able to connect to the Cassandra DB after 10 tries.
        """
        for i in range(10):
            try:
                if cassandra_instance == "local":
                    print("connecting to local Cassandra instance")
                    cluster = Cluster()
                else:
                    print("connecting to Cassandra: " + cassandra_instance)
                    cluster = Cluster([cassandra_instance])

                try:
                    cassandra_session = cluster.connect("zooniverse")
                except cassandra.InvalidRequest:
                    cassandra_session = cluster.connect()
                    cassandra_session.execute("CREATE KEYSPACE zooniverse WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 2 }")
                    cassandra_session = cluster.connect('zooniverse')

                return cassandra_session
            except cassandra.cluster.NoHostAvailable as err:
                if i == 9:
                    raise err
                warning(err)

        raise Exception('Could not connect to cassandra due to unknown error') 

def cassandra_cluster():
    cluster = Cluster([os.environ.get("CASSANDRA_HOST", "localhost")])
    yield cluster
    del cluster 

def __cassandra_connect__(self):
        """
        connect to the AWS instance of Cassandra - try 10 times and raise an error
        :return:
        """
        for i in range(10):
            try:
                self.cluster = Cluster()
                self.cassandra_session = self.cluster.connect('zooniverse')
                return
            except cassandra.cluster.NoHostAvailable:
                pass

        assert False 

def initialize(self, db_ip, db_port, **args):
        ips, default_port = _parse_hosts(args['config'].remote_db_hosts)
        lb_policy = self._get_loadbalancing_policy(
            self.config.load_balancing)
        consistency = self._get_consistency_level(
            self.config.consistency_level)

        self.client = cluster.Cluster(ips, port=default_port,
                                      load_balancing_policy=lb_policy)
        self.session = self.client.connect(ROOT_KS)
        self.session.default_consistency_level = consistency
        self.session.row_factory = query.dict_factory 

def __init__(self, contact_points=_DEFAULT_CONTACT_POINTS, port=9042):
        from cassandra.cluster import Cluster
        self._cluster = Cluster(contact_points, port)
        self._sessions = {} 

def connect(self):
        log.debug("Connecting to CASSANDRA NODES {}".format(current_app.config['CASSANDRA_NODES']))
        if self.cluster is None:
            if isinstance(current_app.config['CASSANDRA_NODES'], (list, tuple)):
                self.cluster = Cluster(current_app.config['CASSANDRA_NODES'])
            elif isinstance(current_app.config['CASSANDRA_NODES'], (str, unicode)):
                self.cluster = Cluster([current_app.config['CASSANDRA_NODES']])
            else:
                raise TypeError("CASSANDRA_NODES must be defined as a list, tuple, string, or unicode object.")

        online_cluster = self.cluster.connect()
        return online_cluster 

def run_server():
    # Initialize database before gunicorn workers startup:
    config = load_app_config()
    cassandra_hosts = [h.strip() for h in config.get('server','cassandra_hosts').split(",")]
    from cstar_perf.frontend.server.model import Model
    from cassandra.cluster import Cluster

    auth_provider = auth_provider_if_configured(config)
    cluster = Cluster(contact_points=cassandra_hosts, auth_provider=auth_provider, connect_timeout=30)

    keyspace = config.get('server', 'cassandra_keyspace') if config.has_option('server', 'cassandra_keyspace') else 'cstar_perf'
    db = Model(cluster=cluster, keyspace=keyspace)
    del db

    app_path = os.path.realpath(os.path.join(os.path.dirname(
        os.path.realpath(__file__)), os.path.pardir, "server"))
    os.chdir(app_path)

    log.info('Waiting for cstar_perf_notifications startup...')
    # this will block until cstar_perf_notifications is up and running
    console_publish('dummy_cluster', {'job_id': 'startup_check', 'msg': 'checking for notification server'})

    # TODO when refactoring how the app is started, do not listen on all interfaces
    proc = subprocess.Popen(shlex.split("gunicorn -k flask_sockets.worker --bind=0.0.0.0:8000"
                                        " -t 300 --log-file=- --workers=10 app:app"))

    # Capture SIGTERM events to shutdown child gunicorn processes..
    def on_terminate(sig, frame):
        print("Killing child processes...")
        proc.terminate()
    signal.signal(signal.SIGTERM, on_terminate)

    proc.communicate() 

def __init__(self):
        cluster = Cluster(["panoptes-cassandra.zooniverse.org"])

        try:
            self.cassandra_session = cluster.connect("active_weather")
        except InvalidRequest as e:
            print(e)
            self.cassandra_session = cluster.connect()
            self.cassandra_session.execute("CREATE KEYSPACE active_weather WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 2 }")
            self.cassandra_session = cluster.connect('active_weather')

        # self.__reset__() 

def init(self, config):
        log = logging.getLogger(__name__)
        log.info("*** STARTING DOMS INITIALIZATION ***")

        domsconfig = ConfigParser.RawConfigParser()
        domsconfig.readfp(pkg_resources.resource_stream(__name__, "domsconfig.ini"), filename='domsconfig.ini')

        cassHost = domsconfig.get("cassandra", "host")
        cassPort = domsconfig.get("cassandra", "port")
        cassKeyspace = domsconfig.get("cassandra", "keyspace")
        cassDatacenter = domsconfig.get("cassandra", "local_datacenter")
        cassVersion = int(domsconfig.get("cassandra", "protocol_version"))

        log.info("Cassandra Host(s): %s" % (cassHost))
        log.info("Cassandra Keyspace: %s" % (cassKeyspace))
        log.info("Cassandra Datacenter: %s" % (cassDatacenter))
        log.info("Cassandra Protocol Version: %s" % (cassVersion))

        dc_policy = DCAwareRoundRobinPolicy(cassDatacenter)
        token_policy = TokenAwarePolicy(dc_policy)

        with Cluster([host for host in cassHost.split(',')], port=cassPort, load_balancing_policy=token_policy,
                     protocol_version=cassVersion) as cluster:
            session = cluster.connect()

            self.createKeyspace(session, cassKeyspace)
            self.createTables(session) 

def _connect(self, cluster, session, contact_points, port):
        lb_policy = c_policies.TokenAwarePolicy(c_policies.DCAwareRoundRobinPolicy())
        # See https://datastax-oss.atlassian.net/browse/PYTHON-643
        lb_policy.shuffle_replicas = True

        if not cluster:
            cluster = c_cluster.Cluster(
                contact_points,
                port,
                compression=self.__compression,
                auth_provider=self.auth_provider,
                # Metrics are disabled because too expensive to compute.
                metrics_enabled=False,
                load_balancing_policy=lb_policy,
            )

            # Limits in flight requests
            cluster.connection_class = getConnectionClass()

        if session:
            session.shutdown()

        session = cluster.connect()
        session.row_factory = c_query.tuple_factory  # Saves 2% CPU
        if self.__timeout:
            session.default_timeout = self.__timeout
        return cluster, session 

def setUpClass(cls):
        """Create the test Cassandra Cluster as cls.cassandra."""
        cls.cassandra = None
        if CASSANDRA_HOSTPORT:
            host, cls.port = CASSANDRA_HOSTPORT.split(":")
            cls.contact_points = [host]
        else:
            cls.setUpCassandra()

        # Make it easy to do raw queries to Cassandra.
        cls.cluster = c_cluster.Cluster(cls.contact_points, cls.port)
        cls.session = cls.cluster.connect()
        cls._reset_keyspace(cls.session, cls.KEYSPACE)
        cls._reset_keyspace(cls.session, cls.KEYSPACE + "_metadata") 

def tearDownClass(cls):
        """Stop the test Cassandra Cluster."""
        cls.cluster.shutdown()
        cls.session = None
        cls.cluster = None
        if cls.cassandra:
            cls.cassandra.stop()
            cls.cassandra = None 

def check_status():
        try:
            cluster = Cluster(
                CONF.cassandra.contact_points
            )
            session = cluster.connect(CONF.cassandra.keyspace)
            session.shutdown()
        except Exception as ex:
            LOG.exception(str(ex))
            return False, str(ex)
        return True, 'OK' 

def __init__(self):
        self.DAY_OF_DATA_CAPTURE = getenv('DAY_OF_DATA_CAPTURE')
        self.MORPHL_SERVER_IP_ADDRESS = getenv('MORPHL_SERVER_IP_ADDRESS')
        self.MORPHL_CASSANDRA_USERNAME = getenv('MORPHL_CASSANDRA_USERNAME')
        self.MORPHL_CASSANDRA_PASSWORD = getenv('MORPHL_CASSANDRA_PASSWORD')
        self.MORPHL_CASSANDRA_KEYSPACE = getenv('MORPHL_CASSANDRA_KEYSPACE')
        self.CASS_REQ_TIMEOUT = 3600.0

        self.auth_provider = PlainTextAuthProvider(
            username=self.MORPHL_CASSANDRA_USERNAME, password=self.MORPHL_CASSANDRA_PASSWORD)
        self.cluster = Cluster(
            contact_points=[self.MORPHL_SERVER_IP_ADDRESS], auth_provider=self.auth_provider)
        self.session = self.cluster.connect(self.MORPHL_CASSANDRA_KEYSPACE)

        self.prepare_statements() 

def __init__(self):
        self.MORPHL_SERVER_IP_ADDRESS = getenv('MORPHL_SERVER_IP_ADDRESS')
        self.MORPHL_CASSANDRA_USERNAME = getenv('MORPHL_CASSANDRA_USERNAME')
        self.MORPHL_CASSANDRA_PASSWORD = getenv('MORPHL_CASSANDRA_PASSWORD')
        self.MORPHL_CASSANDRA_KEYSPACE = getenv('MORPHL_CASSANDRA_KEYSPACE')

        self.prep_stmt = {}

        template_for_prediction = 'INSERT INTO ga_chp_predictions (client_id,prediction) VALUES (?,?)'
        template_for_predictions_by_date = 'INSERT INTO ga_chp_predictions_by_prediction_date (prediction_date, client_id, prediction) VALUES (?,?,?)'
        template_for_predictions_statistics = 'UPDATE ga_chp_predictions_statistics SET loyal=loyal+?, neutral=neutral+?, churning=churning+?, lost=lost+? WHERE prediction_date=?'

        self.CASS_REQ_TIMEOUT = 3600.0

        self.auth_provider = PlainTextAuthProvider(
            username=self.MORPHL_CASSANDRA_USERNAME,
            password=self.MORPHL_CASSANDRA_PASSWORD)
        self.cluster = Cluster(
            [self.MORPHL_SERVER_IP_ADDRESS], auth_provider=self.auth_provider)
        self.session = self.cluster.connect(self.MORPHL_CASSANDRA_KEYSPACE)

        self.prep_stmt['prediction'] = self.session.prepare(
            template_for_prediction)
        self.prep_stmt['predictions_by_date'] = self.session.prepare(
            template_for_predictions_by_date)
        self.prep_stmt['predictions_statistics'] = self.session.prepare(
            template_for_predictions_statistics) 

def init(self, conf):
        conf = conf.with_fallback(CassandraExtractor.DEFAULT_CONFIG)
        self._cluster = '{}'.format(conf.get_string(CassandraExtractor.CLUSTER_KEY))
        self._filter = conf.get(CassandraExtractor.FILTER_FUNCTION_KEY)
        ips = conf.get_list(CassandraExtractor.IPS_KEY)
        kwargs = conf.get(CassandraExtractor.KWARGS_KEY)
        self._client = Cluster(ips, **kwargs)
        self._client.connect()
        self._extract_iter = None  # type: Union[None, Iterator] 

def test_main_thread_loop(loop, session):
    asyncio.set_event_loop(loop)
    cluster = Cluster()
    session = cluster.connect()

    aiosession(session)

    assert loop is session._asyncio_loop 

def cluster():
    cluster = Cluster()

    yield cluster

    cluster.shutdown() 

def shutdown_cluster(self) -> None:
        """
        Closes all sessions and connections associated with this Cluster.
        """
        if not self.cluster.is_shutdown:
            self.cluster.shutdown()