def setUp(self): """ Setup sessions and pause node1 """ self.cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile( load_balancing_policy=HostFilterPolicy( RoundRobinPolicy(), lambda host: host.address == "127.0.0.1" ) ) } ) self.session = self.cluster.connect(wait_for_all_pools=True) self.control_connection_host_number = 1 self.node_to_stop = get_node(self.control_connection_host_number) ddl = ''' CREATE TABLE test3rf.timeout ( k int PRIMARY KEY, v int )''' self.session.execute(ddl) self.node_to_stop.pause()
def use_cluster_with_graph(num_nodes): """ This is a work around to account for the fact that spark nodes will conflict over master assignment when started all at once. """ # Create the cluster but don't start it. use_singledc(start=False, workloads=['graph', 'spark']) # Start first node. get_node(1).start(wait_for_binary_proto=True) # Wait binary protocol port to open wait_for_node_socket(get_node(1), 120) # Wait for spark master to start up spark_master_http = ("localhost", 7080) common.check_socket_listening(spark_master_http, timeout=60) tmp_cluster = Cluster(protocol_version=PROTOCOL_VERSION) # Start up remaining nodes. try: session = tmp_cluster.connect() statement = "ALTER KEYSPACE dse_leases WITH REPLICATION = {'class': 'NetworkTopologyStrategy', 'dc1': '%d'}" % ( num_nodes) session.execute(statement) finally: tmp_cluster.shutdown() for i in range(1, num_nodes + 1): if i is not 1: node = get_node(i) node.start(wait_for_binary_proto=True) wait_for_node_socket(node, 120) # Wait for workers to show up as Alive on master wait_for_spark_workers(3, 120)
def main(): cluster = Cluster() session = cluster.connect() start_time = datetime.datetime.fromtimestamp(1564524000) end_time = datetime.datetime.fromtimestamp(1564610400) results = [] current_time = start_time while current_time <= end_time: rows = session.execute( """ SELECT user_id, product_id, year, month, day, hour, interaction_time, type FROM prodcat.user_interaction_product_history WHERE year = %s AND month = %s AND day = %s AND hour = %s """, (current_time.year, current_time.month, current_time.day, current_time.hour) ) results += rows.current_rows current_time += datetime.timedelta(hours=1) interactions = parse_interaction_data(results, start_time) interaction_details = calc_interaction_details(interactions) model = initialize_model(interaction_details) for interaction in interactions[:50]: predictions = evaluate_model(interaction, interaction_details, model, 18) print("Predicted Product: %s, Actual: %s", predictions[0], interaction.product_sequence_id) print("Predicted State Change: %s, Actual: %s", predictions[1], interaction.state)
def use_cluster_with_graph(num_nodes): """ This is a work around to account for the fact that spark nodes will conflict over master assignment when started all at once. """ # Create the cluster but don't start it. use_singledc(start=False, workloads=['graph', 'spark']) # Start first node. get_node(1).start(wait_for_binary_proto=True) # Wait binary protocol port to open wait_for_node_socket(get_node(1), 120) # Wait for spark master to start up spark_master_http = ("localhost", 7080) common.check_socket_listening(spark_master_http, timeout=60) tmp_cluster = Cluster(protocol_version=PROTOCOL_VERSION) # Start up remaining nodes. try: session = tmp_cluster.connect() statement = "ALTER KEYSPACE dse_leases WITH REPLICATION = {'class': 'NetworkTopologyStrategy', 'dc1': '%d'}" % (num_nodes) session.execute(statement) finally: tmp_cluster.shutdown() for i in range(1, num_nodes+1): if i is not 1: node = get_node(i) node.start(wait_for_binary_proto=True) wait_for_node_socket(node, 120) # Wait for workers to show up as Alive on master wait_for_spark_workers(3, 120)
def _cluster_session_with_lbp(self, lbp): # create a cluster with no delay on events cluster = Cluster(protocol_version=PROTOCOL_VERSION, topology_event_refresh_window=0, status_event_refresh_window=0, execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(load_balancing_policy=lbp)}) session = cluster.connect() return cluster, session
def test_numpy_results_paged(self): """ Test Numpy-based parser that returns a NumPy array """ # arrays = { 'a': arr1, 'b': arr2, ... } cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=tuple_factory)}) session = cluster.connect(keyspace="testspace") session.client_protocol_handler = NumpyProtocolHandler session.default_fetch_size = 2 expected_pages = (self.N_ITEMS + session.default_fetch_size - 1) // session.default_fetch_size self.assertLess(session.default_fetch_size, self.N_ITEMS) results = session.execute("SELECT * FROM test_table") self.assertTrue(results.has_more_pages) for count, page in enumerate(results, 1): self.assertIsInstance(page, dict) for colname, arr in page.items(): if count <= expected_pages: self.assertGreater(len(arr), 0, "page count: %d" % (count,)) self.assertLessEqual(len(arr), session.default_fetch_size) else: # we get one extra item out of this iteration because of the way NumpyParser returns results # The last page is returned as a dict with zero-length arrays self.assertEqual(len(arr), 0) self.assertEqual(self._verify_numpy_page(page), len(arr)) self.assertEqual(count, expected_pages + 1) # see note about extra 'page' above cluster.shutdown()
def test_protocol_divergence_v5_fail_by_continuous_paging(self): """ Test to validate that V5 and DSE_V1 diverge. ContinuousPagingOptions is not supported by V5 @since DSE 2.0b3 GRAPH 1.0b1 @jira_ticket PYTHON-694 @expected_result NoHostAvailable will be risen when the continuous_paging_options parameter is set @test_category connection """ cluster = Cluster(protocol_version=ProtocolVersion.V5, allow_beta_protocol_version=True) session = cluster.connect() max_pages = 4 max_pages_per_second = 3 continuous_paging_options = ContinuousPagingOptions( max_pages=max_pages, max_pages_per_second=max_pages_per_second) future = self._send_query_message( session, timeout=cluster._default_timeout, consistency_level=ConsistencyLevel.ONE, continuous_paging_options=continuous_paging_options) # This should raise NoHostAvailable because continuous paging is not supported under ProtocolVersion.DSE_V1 with self.assertRaises(NoHostAvailable) as context: future.result() self.assertIn( "Continuous paging may only be used with protocol version ProtocolVersion.DSE_V1 or higher", str(context.exception)) cluster.shutdown()
def test_white_list(self): use_singledc() keyspace = 'test_white_list' cluster = Cluster(('127.0.0.2',), protocol_version=PROTOCOL_VERSION, topology_event_refresh_window=0, status_event_refresh_window=0, execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(load_balancing_policy=WhiteListRoundRobinPolicy((IP_FORMAT % 2,)))}) session = cluster.connect() self._wait_for_nodes_up([1, 2, 3]) create_schema(cluster, session, keyspace) self._insert(session, keyspace) self._query(session, keyspace) self.coordinator_stats.assert_query_count_equals(self, 1, 0) self.coordinator_stats.assert_query_count_equals(self, 2, 12) self.coordinator_stats.assert_query_count_equals(self, 3, 0) # white list policy should not allow reconnecting to ignored hosts force_stop(3) self._wait_for_nodes_down([3]) self.assertFalse(cluster.metadata._hosts[IP_FORMAT % 3].is_currently_reconnecting()) self.coordinator_stats.reset_counts() force_stop(2) self._wait_for_nodes_down([2]) try: self._query(session, keyspace) self.fail() except NoHostAvailable: pass finally: cluster.shutdown()
def _setup_for_proxy(self, grant=True): os.environ['KRB5_CONFIG'] = self.krb_conf self.refresh_kerberos_tickets(self.cassandra_keytab, "*****@*****.**", self.krb_conf) auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"], principal='*****@*****.**') cluster = Cluster(auth_provider=auth_provider) session = cluster.connect() session.execute("CREATE ROLE IF NOT EXISTS '{0}' WITH LOGIN = TRUE;".format('*****@*****.**')) session.execute("CREATE ROLE IF NOT EXISTS '{0}' WITH LOGIN = TRUE;".format('*****@*****.**')) session.execute("GRANT EXECUTE ON ALL AUTHENTICATION SCHEMES to '*****@*****.**'") session.execute("CREATE ROLE IF NOT EXISTS '{0}' WITH LOGIN = TRUE;".format('*****@*****.**')) session.execute("GRANT EXECUTE ON ALL AUTHENTICATION SCHEMES to '*****@*****.**'") # Create a keyspace and allow only charlie to query it. session.execute( "CREATE KEYSPACE testkrbproxy WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}") session.execute("CREATE TABLE testkrbproxy.testproxy (id int PRIMARY KEY, value text)") session.execute("GRANT ALL PERMISSIONS ON KEYSPACE testkrbproxy to '{0}'".format('*****@*****.**')) if grant: session.execute("GRANT PROXY.LOGIN ON ROLE '{0}' to '{1}'".format('*****@*****.**', '*****@*****.**')) cluster.shutdown()
def test_pool_management(self): # Ensure that in_flight and request_ids quiesce after cluster operations cluster = Cluster( protocol_version=PROTOCOL_VERSION, idle_heartbeat_interval=0 ) # no idle heartbeat here, pool management is tested in test_idle_heartbeat session = cluster.connect() session2 = cluster.connect() # prepare p = session.prepare("SELECT * FROM system.local WHERE key=?") self.assertTrue(session.execute(p, ('local', ))) # simple self.assertTrue( session.execute("SELECT * FROM system.local WHERE key='local'")) # set keyspace session.set_keyspace('system') session.set_keyspace('system_traces') # use keyspace session.execute('USE system') session.execute('USE system_traces') # refresh schema cluster.refresh_schema_metadata() cluster.refresh_schema_metadata(max_schema_agreement_wait=0) assert_quiescent_pool_state(self, cluster) cluster.shutdown()
def test_rfthree_tokenaware_none_down(self): keyspace = 'test_rfthree_tokenaware_none_down' cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile( TokenAwarePolicy(RoundRobinPolicy())) }) session = cluster.connect(wait_for_all_pools=True) wait_for_up(cluster, 1) wait_for_up(cluster, 2) create_schema(cluster, session, keyspace, replication_factor=3) self._insert(session, keyspace, count=1) self._query(session, keyspace, count=1) self.coordinator_stats.assert_query_count_equals(self, 1, 0) self.coordinator_stats.assert_query_count_equals(self, 2, 1) self.coordinator_stats.assert_query_count_equals(self, 3, 0) self.coordinator_stats.reset_counts() self._assert_writes_succeed(session, keyspace, SINGLE_DC_CONSISTENCY_LEVELS) self._assert_reads_succeed(session, keyspace, SINGLE_DC_CONSISTENCY_LEVELS - set([ConsistencyLevel.ANY]), expected_reader=2) cluster.shutdown()
def validate_ssl_options(ssl_options): # find absolute path to client CA_CERTS tries = 0 while True: if tries > 5: raise RuntimeError( "Failed to connect to SSL cluster after 5 attempts") try: cluster = Cluster(protocol_version=PROTOCOL_VERSION, ssl_options=ssl_options) session = cluster.connect(wait_for_all_pools=True) break except Exception: ex_type, ex, tb = sys.exc_info() log.warn("{0}: {1} Backtrace: {2}".format( ex_type.__name__, ex, traceback.extract_tb(tb))) del tb tries += 1 # attempt a few simple commands. insert_keyspace = """CREATE KEYSPACE ssltest WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '3'} """ statement = SimpleStatement(insert_keyspace) statement.consistency_level = 1 session.execute(statement) drop_keyspace = "DROP KEYSPACE ssltest" statement = SimpleStatement(drop_keyspace) statement.consistency_level = ConsistencyLevel.ANY session.execute(statement) cluster.shutdown()
def load_schema(self, use_schema=False): """加载schema文件 """ if use_schema: # 创建默认的执行配置,指向特定的graph ep = GraphExecutionProfile(graph_options=GraphOptions( graph_name=self.graph_name)) cluster = Cluster( execution_profiles={EXEC_PROFILE_GRAPH_DEFAULT: ep}) session = cluster.connect() # 创建graph session.execute_graph( "system.graph(name).ifNotExists().create()", {'name': self.graph_name}, execution_profile=EXEC_PROFILE_GRAPH_SYSTEM_DEFAULT) # 批量执行gremlin创建schema的命令 with open(self.schema_path) as f: pat = re.compile("[/ \n]") # 正则表达式 for line in f: if not re.match(pat, line): # 开头不是斜杠、空格、换行的 print("正在加载 {}".format(line.strip())) session.execute_graph(line.strip()) else: print("schema未加载,请确保graph中存在schema")
class ConnectionTimeoutTest(unittest.TestCase): def setUp(self): self.cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles= {EXEC_PROFILE_DEFAULT: ExecutionProfile(load_balancing_policy= WhiteListRoundRobinPolicy([DSE_IP]))}) self.session = self.cluster.connect() def tearDown(self): self.cluster.shutdown() @patch('dse.connection.Connection.max_in_flight', 2) def test_in_flight_timeout(self): """ Test to ensure that connection id fetching will block when max_id is reached/ In previous versions of the driver this test will cause a NoHostAvailable exception to be thrown, when the max_id is restricted @since 3.3 @jira_ticket PYTHON-514 @expected_result When many requests are run on a single node connection acquisition should block until connection is available or the request times out. @test_category connection timeout """ futures = [] query = '''SELECT * FROM system.local''' for _ in range(100): futures.append(self.session.execute_async(query)) for future in futures: future.result()
def main(): cluster = Cluster() session = cluster.connect() start_time = datetime.datetime.fromtimestamp(1564524000) end_time = datetime.datetime.fromtimestamp(1564610400) results = [] current_time = start_time while current_time <= end_time: rows = session.execute( """ SELECT user_id, product_id, year, month, day, hour, interaction_time, type FROM prodcat.user_interaction_product_history WHERE year = %s AND month = %s AND day = %s AND hour = %s """, (current_time.year, current_time.month, current_time.day, current_time.hour)) results += rows.current_rows current_time += datetime.timedelta(hours=1) interactions = parse_interaction_data(results, start_time) interaction_details = calc_interaction_details(interactions) model = initialize_model(interaction_details) train(interactions, interaction_details, model, 150)
def test_token_aware_is_used_by_default(self): """ Test for default load balancing policy test_token_aware_is_used_by_default tests that the default loadbalancing policy is policies.TokenAwarePolicy. It creates a simple Cluster and verifies that the default loadbalancing policy is TokenAwarePolicy if the murmur3 C extension is found. Otherwise, the default loadbalancing policy is DCAwareRoundRobinPolicy. @since 2.6.0 @jira_ticket PYTHON-160 @expected_result TokenAwarePolicy should be the default loadbalancing policy. @test_category load_balancing:token_aware """ cluster = Cluster(protocol_version=PROTOCOL_VERSION) if murmur3 is not None: self.assertTrue( isinstance( cluster.profile_manager.default.load_balancing_policy, TokenAwarePolicy)) else: self.assertTrue( isinstance( cluster.profile_manager.default.load_balancing_policy, DCAwareRoundRobinPolicy)) cluster.shutdown()
def _protocol_divergence_fail_by_flag_uses_int(self, version, uses_int_query_flag, int_flag=True, beta=False): cluster = Cluster(protocol_version=version, allow_beta_protocol_version=beta) session = cluster.connect() query_one = SimpleStatement( "INSERT INTO test3rf.test (k, v) VALUES (1, 1)") query_two = SimpleStatement( "INSERT INTO test3rf.test (k, v) VALUES (2, 2)") execute_with_long_wait_retry(session, query_one) execute_with_long_wait_retry(session, query_two) with mock.patch('dse.protocol.ProtocolVersion.uses_int_query_flags', new=mock.Mock(return_value=int_flag)): future = self._send_query_message( session, cluster._default_timeout, consistency_level=ConsistencyLevel.ONE, fetch_size=1) response = future.result() # This means the flag are not handled as they are meant by the server if uses_int=False self.assertEqual(response.has_more_pages, uses_int_query_flag) execute_with_long_wait_retry(session, SimpleStatement("TRUNCATE test3rf.test")) cluster.shutdown()
class BasicGraphUnitTestCase(BasicKeyspaceUnitTestCase): """ This is basic graph unit test case that provides various utility methods that can be leveraged for testcase setup and tear down """ @property def graph_name(self): return self._testMethodName.lower() def session_setup(self): self.cluster = Cluster(protocol_version=PROTOCOL_VERSION) self.session = self.cluster.connect() self.ks_name = self._testMethodName.lower() self.cass_version, self.cql_version = get_server_versions() def setUp(self): self.session_setup() self.reset_graph() profiles = self.cluster.profile_manager.profiles profiles[EXEC_PROFILE_GRAPH_DEFAULT].graph_options.graph_name = self.graph_name profiles[EXEC_PROFILE_GRAPH_ANALYTICS_DEFAULT].graph_options.graph_name = self.graph_name self.clear_schema() def tearDown(self): self.cluster.shutdown() def clear_schema(self): self.session.execute_graph('schema.clear()') def reset_graph(self): reset_graph(self.session, self.graph_name) def wait_for_graph_inserted(self): wait_for_graph_inserted(self.session, self.graph_name)
def _test_downgrading_cl(self, keyspace, rf, accepted): cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile( TokenAwarePolicy(RoundRobinPolicy()), DowngradingConsistencyRetryPolicy()) }) session = cluster.connect(wait_for_all_pools=True) create_schema(cluster, session, keyspace, replication_factor=rf) self._insert(session, keyspace, 1) self._query(session, keyspace, 1) self.coordinator_stats.assert_query_count_equals(self, 1, 0) self.coordinator_stats.assert_query_count_equals(self, 2, 1) self.coordinator_stats.assert_query_count_equals(self, 3, 0) try: force_stop(2) wait_for_down(cluster, 2) self._assert_writes_succeed(session, keyspace, accepted) self._assert_reads_succeed(session, keyspace, accepted - set([ConsistencyLevel.ANY])) self._assert_writes_fail(session, keyspace, SINGLE_DC_CONSISTENCY_LEVELS - accepted) self._assert_reads_fail(session, keyspace, SINGLE_DC_CONSISTENCY_LEVELS - accepted) finally: start(2) wait_for_up(cluster, 2) cluster.shutdown()
def setUp(self): if PROTOCOL_VERSION < 4: raise unittest.SkipTest( "Native protocol 4,0+ is required for custom payloads, currently using %r" % (PROTOCOL_VERSION, )) self.cluster = Cluster(protocol_version=PROTOCOL_VERSION) self.session = self.cluster.connect()
def test_connect_to_already_shutdown_cluster(self): """ Ensure you cannot connect to a cluster that's been shutdown """ cluster = Cluster(protocol_version=PROTOCOL_VERSION) cluster.shutdown() self.assertRaises(Exception, cluster.connect)
def test_for_schema_disagreements_same_keyspace(self): """ Tests for any schema disagreements using the same keyspace multiple times """ cluster = Cluster(protocol_version=PROTOCOL_VERSION) session = cluster.connect(wait_for_all_pools=True) for i in range(30): try: execute_until_pass( session, "CREATE KEYSPACE test WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}" ) except AlreadyExists: execute_until_pass(session, "DROP KEYSPACE test") execute_until_pass( session, "CREATE KEYSPACE test WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}" ) execute_until_pass( session, "CREATE TABLE test.cf (key int PRIMARY KEY, value int)") for j in range(100): execute_until_pass( session, "INSERT INTO test.cf (key, value) VALUES ({0}, {0})". format(j)) execute_until_pass(session, "DROP KEYSPACE test") cluster.shutdown()
def test_prepare_on_all_hosts(self): """ Test to validate prepare_on_all_hosts flag is honored. Use a special ForcedHostSwitchPolicy to ensure prepared queries are cycled over nodes that should not have them prepared. Check the logs to insure they are being re-prepared on those nodes @since 3.4.0 @jira_ticket PYTHON-556 @expected_result queries will have to re-prepared on hosts that aren't the control connection """ white_list = ForcedHostSwitchPolicy() clus = Cluster( execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile(load_balancing_policy=white_list)}, protocol_version=PROTOCOL_VERSION, prepare_on_all_hosts=False, reprepare_on_up=False) self.addCleanup(clus.shutdown) session = clus.connect(wait_for_all_pools=True) mock_handler = MockLoggingHandler() logger = logging.getLogger(cluster.__name__) logger.addHandler(mock_handler) select_statement = session.prepare("SELECT * FROM system.local") session.execute(select_statement) session.execute(select_statement) session.execute(select_statement) self.assertEqual(2, mock_handler.get_message_count('debug', "Re-preparing"))
def setup(self): """Setup the connection""" global cluster, session if 'username' in self.cluster_options or 'password' in self.cluster_options: raise CQLEngineException("Username & Password are now handled by using the native driver's auth_provider") if self.lazy_connect: return self.cluster = Cluster(self.hosts, **self.cluster_options) try: self.session = self.cluster.connect() log.debug(format_log_context("connection initialized with internally created session", connection=self.name)) except NoHostAvailable: if self.retry_connect: log.warning(format_log_context("connect failed, setting up for re-attempt on first use", connection=self.name)) self.lazy_connect = True raise if DEFAULT_CONNECTION in _connections and _connections[DEFAULT_CONNECTION] == self: cluster = _connections[DEFAULT_CONNECTION].cluster session = _connections[DEFAULT_CONNECTION].session self.setup_session()
def __init__(self): self.secure_connect_bundle = '/home/ubuntu/workspace/creds.zip' self.path_to_creds = '' self.cluster = Cluster( cloud={'secure_connect_bundle': self.secure_connect_bundle}, auth_provider=PlainTextAuthProvider('KVUser', 'KVPassword')) self.session = self.cluster.connect()
def test_duplicate_metrics_per_cluster(self): """ Test to validate that cluster metrics names can't overlap. @since 3.6.0 @jira_ticket PYTHON-561 @expected_result metric names should not be allowed to be same. @test_category metrics """ cluster2 = Cluster( metrics_enabled=True, protocol_version=PROTOCOL_VERSION, execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile(retry_policy=FallthroughRetryPolicy()) }) cluster3 = Cluster( metrics_enabled=True, protocol_version=PROTOCOL_VERSION, execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile(retry_policy=FallthroughRetryPolicy()) }) # Ensure duplicate metric names are not allowed cluster2.metrics.set_stats_name("appcluster") cluster2.metrics.set_stats_name("appcluster") with self.assertRaises(ValueError): cluster3.metrics.set_stats_name("appcluster") cluster3.metrics.set_stats_name("devops") session2 = cluster2.connect(self.ks_name, wait_for_all_pools=True) session3 = cluster3.connect(self.ks_name, wait_for_all_pools=True) # Basic validation that naming metrics doesn't impact their segration or accuracy for i in range(10): query = SimpleStatement("SELECT * FROM {0}.{0}".format( self.ks_name), consistency_level=ConsistencyLevel.ALL) session2.execute(query) for i in range(5): query = SimpleStatement("SELECT * FROM {0}.{0}".format( self.ks_name), consistency_level=ConsistencyLevel.ALL) session3.execute(query) self.assertEqual( cluster2.metrics.get_stats()['request_timer']['count'], 10) self.assertEqual( cluster3.metrics.get_stats()['request_timer']['count'], 5) # Check scales to ensure they are appropriately named self.assertTrue("appcluster" in scales._Stats.stats.keys()) self.assertTrue("devops" in scales._Stats.stats.keys()) cluster2.shutdown() cluster3.shutdown()
def setUp(self): self.c = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=dict_factory) }) self.session1 = self.c.connect(keyspace=self.keyspace1) self.session2 = self.c.connect(keyspace=self.keyspace2)
def _results_from_row_factory(self, row_factory): cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={ EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=row_factory) }) with cluster: return cluster.connect().execute(self.select)
def setup_cluster(): if args.host is None: nodes = ['localhost'] else: nodes = [args.host] if args.port is None: port = 9042 else: port = int(args.port) if args.connect_timeout is None: connect_timeout = 5 else: connect_timeout = int(args.connect_timeout) if args.ssl is not None and args.certfile is not None: ssl_opts = { 'ca_certs': args.certfile, 'ssl_version': PROTOCOL_TLSv1, 'keyfile': args.userkey, 'certfile': args.usercert } else: ssl_opts = {} cluster = None if args.protocol_version is not None: auth = None if args.username is not None and args.password is not None: if args.protocol_version == 1: auth = get_credentials elif args.protocol_version > 1: auth = PlainTextAuthProvider(username=args.username, password=args.password) cluster = Cluster(control_connection_timeout=connect_timeout, connect_timeout=connect_timeout, contact_points=nodes, port=port, protocol_version=args.protocol_version, auth_provider=auth, ssl_options=ssl_opts) else: cluster = Cluster(control_connection_timeout=connect_timeout, connect_timeout=connect_timeout, contact_points=nodes, port=port, ssl_options=ssl_opts) session = cluster.connect() session.default_timeout = TIMEOUT session.default_fetch_size = FETCH_SIZE session.row_factory = dse.query.ordered_dict_factory return session
def test_connect_no_auth_provider(self): cluster = Cluster(protocol_version=PROTOCOL_VERSION) try: self.assertRaisesRegexp(NoHostAvailable, '.*AuthenticationFailed.*', cluster.connect) assert_quiescent_pool_state(self, cluster) finally: cluster.shutdown()
def setUp(self): self.cluster = Cluster(protocol_version=PROTOCOL_VERSION) self.session = self.cluster.connect() query = """ INSERT INTO test3rf.test (k, v) VALUES (?, ?) """ self.simple_statement = SimpleStatement(query, routing_key='ss_rk', keyspace='keyspace_name') self.prepared = self.session.prepare(query)
def setUpClass(self): """ This will setup the necessary infrastructure to run unified authentication tests. """ if DSE_VERSION and DSE_VERSION < '5.1': return self.cluster = None ccm_cluster = get_cluster() # Stop cluster if running and configure it with the correct options ccm_cluster.stop() if isinstance(ccm_cluster, DseCluster): # Setup dse options in dse.yaml config_options = { 'authentication_options': { 'enabled': 'true', 'default_scheme': 'internal', 'scheme_permissions': 'true', 'transitional_mode': 'normal' }, 'authorization_options': { 'enabled': 'true' } } # Setup dse authenticator in cassandra.yaml ccm_cluster.set_configuration_options({ 'authenticator': 'com.datastax.bdp.cassandra.auth.DseAuthenticator', 'authorizer': 'com.datastax.bdp.cassandra.auth.DseAuthorizer' }) ccm_cluster.set_dse_configuration_options(config_options) ccm_cluster.start(wait_for_binary_proto=True, wait_other_notice=True) else: log.error("Cluster is not dse cluster test will fail") # Create users and test keyspace self.user_role = 'user1' self.server_role = 'server' self.root_cluster = Cluster( auth_provider=DSEPlainTextAuthProvider('cassandra', 'cassandra')) self.root_session = self.root_cluster.connect() self.root_session.execute("CREATE USER {0} WITH PASSWORD '{1}'".format( self.server_role, self.server_role)) self.root_session.execute("CREATE USER {0} WITH PASSWORD '{1}'".format( self.user_role, self.user_role)) self.root_session.execute( "CREATE KEYSPACE testproxy WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}" ) self.root_session.execute( "CREATE TABLE testproxy.testproxy (id int PRIMARY KEY, value text)" ) self.root_session.execute( "GRANT ALL PERMISSIONS ON KEYSPACE testproxy to {0}".format( self.user_role))
def connect_and_query(self, auth_provider): """ Runs a simple system query with the auth_provided specified. """ os.environ['KRB5_CONFIG'] = self.krb_conf self.cluster = Cluster(auth_provider=auth_provider) self.session = self.cluster.connect() query = "SELECT * FROM system.local" statement = SimpleStatement(query) rs = self.session.execute(statement) return rs
def test_graph_profile(self): """ Test verifying various aspects of graph config properties. @since 1.0.0 @jira_ticket PYTHON-570 @test_category dse graph """ generate_classic(self.session) # Create variou execution policies exec_dif_factory = GraphExecutionProfile(row_factory=single_object_row_factory) exec_dif_factory.graph_options.graph_name = self.graph_name exec_dif_lbp = GraphExecutionProfile(load_balancing_policy=WhiteListRoundRobinPolicy(['127.0.0.1'])) exec_dif_lbp.graph_options.graph_name = self.graph_name exec_bad_lbp = GraphExecutionProfile(load_balancing_policy=WhiteListRoundRobinPolicy(['127.0.0.2'])) exec_dif_lbp.graph_options.graph_name = self.graph_name exec_short_timeout = GraphExecutionProfile(request_timeout=1, load_balancing_policy=WhiteListRoundRobinPolicy(['127.0.0.1'])) exec_short_timeout.graph_options.graph_name = self.graph_name # Add a single exection policy on cluster creation local_cluster = Cluster(protocol_version=PROTOCOL_VERSION, execution_profiles={"exec_dif_factory": exec_dif_factory}) local_session = local_cluster.connect() rs1 = self.session.execute_graph('g.V()') rs2 = local_session.execute_graph('g.V()', execution_profile='exec_dif_factory') # Verify default and non default policy works self.assertFalse(isinstance(rs2[0], Vertex)) self.assertTrue(isinstance(rs1[0], Vertex)) # Add other policies validate that lbp are honored local_cluster.add_execution_profile("exec_dif_ldp", exec_dif_lbp) local_session.execute_graph('g.V()', execution_profile="exec_dif_ldp") local_cluster.add_execution_profile("exec_bad_lbp", exec_bad_lbp) with self.assertRaises(NoHostAvailable): local_session.execute_graph('g.V()', execution_profile="exec_bad_lbp") # Try with missing EP with self.assertRaises(ValueError): local_session.execute_graph('g.V()', execution_profile='bad_exec_profile') # Validate that timeout is honored local_cluster.add_execution_profile("exec_short_timeout", exec_short_timeout) with self.assertRaises(OperationTimedOut): local_session.execute_graph('java.util.concurrent.TimeUnit.MILLISECONDS.sleep(2000L);', execution_profile='exec_short_timeout')
class BasicDseAuthTest(unittest.TestCase): @classmethod def setUpClass(self): """ This will setup the necessary infrastructure to run our authentication tests. It requres the ADS_HOME environment variable and our custom embedded apache directory server jar in order to run. """ clear_kerberos_tickets() # Setup variables for various keytab and other files self.conf_file_dir = ADS_HOME+"conf/" self.krb_conf = self.conf_file_dir+"krb5.conf" self.dse_keytab = self.conf_file_dir+"dse.keytab" self.dseuser_keytab = self.conf_file_dir+"dseuser.keytab" self.cassandra_keytab = self.conf_file_dir+"cassandra.keytab" actual_jar = ADS_HOME+"embedded-ads.jar" # Create configuration directories if they don't already exists if not os.path.exists(self.conf_file_dir): os.makedirs(self.conf_file_dir) log.warning("Starting adserver") # Start the ADS, this will create the keytab con configuration files listed above self.proc = subprocess.Popen(['java', '-jar', actual_jar, '-k', '--confdir', self.conf_file_dir], shell=False) time.sleep(10) # TODO poll for server to come up log.warning("Starting adserver started") ccm_cluster = get_cluster() log.warning("fetching tickets") # Stop cluster if running and configure it with the correct options ccm_cluster.stop() if isinstance(ccm_cluster, DseCluster): # Setup kerberos options in dse.yaml config_options = {'kerberos_options': {'keytab': self.dse_keytab, 'service_principal': 'dse/[email protected]', 'qop': 'auth'}, 'authentication_options': {'enabled': 'true', 'default_scheme': 'kerberos', 'scheme_permissions': 'true', 'allow_digest_with_kerberos': 'true', 'plain_text_without_ssl': 'warn', 'transitional_mode': 'disabled'} } krb5java = "-Djava.security.krb5.conf=" + self.krb_conf # Setup dse authenticator in cassandra.yaml ccm_cluster.set_configuration_options({'authenticator': 'com.datastax.bdp.cassandra.auth.DseAuthenticator'}) ccm_cluster.set_dse_configuration_options(config_options) ccm_cluster.start(wait_for_binary_proto=True, wait_other_notice=True, jvm_args=[krb5java]) else: log.error("Cluster is not dse cluster test will fail") @classmethod def tearDownClass(self): """ Terminates running ADS (Apache directory server). """ self.proc.terminate() def tearDown(self): """ This will clear any existing kerberos tickets by using kdestroy """ clear_kerberos_tickets() self.cluster.shutdown() def refresh_kerberos_tickets(self, keytab_file, user_name, krb_conf): """ Fetches a new ticket for using the keytab file and username provided. """ self.ads_pid = subprocess.call(['kinit', '-t', keytab_file, user_name], env={'KRB5_CONFIG': krb_conf}, shell=False) def connect_and_query(self, auth_provider): """ Runs a simple system query with the auth_provided specified. """ os.environ['KRB5_CONFIG'] = self.krb_conf self.cluster = Cluster(auth_provider=auth_provider) self.session = self.cluster.connect() query = "SELECT * FROM system.local" statement = SimpleStatement(query) rs = self.session.execute(statement) return rs def test_should_not_authenticate_with_bad_user_ticket(self): """ This tests will attempt to authenticate with a user that has a valid ticket, but is not a valid dse user. @since 1.0.0 @jira_ticket PYTHON-457 @test_category dse auth @expected_result NoHostAvailable exception should be thrown """ self.refresh_kerberos_tickets(self.dseuser_keytab, "*****@*****.**", self.krb_conf) auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"]) self.assertRaises(NoHostAvailable, self.connect_and_query, auth_provider) def test_should_not_athenticate_without_ticket(self): """ This tests will attempt to authenticate with a user that is valid but has no ticket @since 1.0.0 @jira_ticket PYTHON-457 @test_category dse auth @expected_result NoHostAvailable exception should be thrown """ auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"]) self.assertRaises(NoHostAvailable, self.connect_and_query, auth_provider) def test_connect_with_kerberos(self): """ This tests will attempt to authenticate with a user that is valid and has a ticket @since 1.0.0 @jira_ticket PYTHON-457 @test_category dse auth @expected_result Client should be able to connect and run a basic query """ self.refresh_kerberos_tickets(self.cassandra_keytab, "*****@*****.**", self.krb_conf) auth_provider = DSEGSSAPIAuthProvider() rs = self.connect_and_query(auth_provider) self.assertIsNotNone(rs) connections = [c for holders in self.cluster.get_connection_holders() for c in holders.get_connections()] # Check to make sure our server_authenticator class is being set appropriate for connection in connections: self.assertTrue('DseAuthenticator' in connection.authenticator.server_authenticator_class) def test_connect_with_kerberos_and_graph(self): """ This tests will attempt to authenticate with a user and execute a graph query @since 1.0.0 @jira_ticket PYTHON-457 @test_category dse auth @expected_result Client should be able to connect and run a basic graph query with authentication """ self.refresh_kerberos_tickets(self.cassandra_keytab, "*****@*****.**", self.krb_conf) auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"]) rs = self.connect_and_query(auth_provider) self.assertIsNotNone(rs) reset_graph(self.session, self._testMethodName.lower()) profiles = self.cluster.profile_manager.profiles profiles[EXEC_PROFILE_GRAPH_DEFAULT].graph_options.graph_name = self._testMethodName.lower() generate_classic(self.session) rs = self.session.execute_graph('g.V()') self.assertIsNotNone(rs) def test_connect_with_kerberos_host_not_resolved(self): """ This tests will attempt to authenticate with IP, this will fail. @since 1.0.0 @jira_ticket PYTHON-566 @test_category dse auth @expected_result Client should error when ip is used """ self.refresh_kerberos_tickets(self.cassandra_keytab, "*****@*****.**", self.krb_conf) auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"], resolve_host_name=False) self.assertRaises(NoHostAvailable, self.connect_and_query, auth_provider) def test_connect_with_explicit_principal(self): """ This tests will attempt to authenticate using valid and invalid user principals @since 1.0.0 @jira_ticket PYTHON-574 @test_category dse auth @expected_result Client principals should be used by the underlying mechanism """ # Connect with valid principal self.refresh_kerberos_tickets(self.cassandra_keytab, "*****@*****.**", self.krb_conf) auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"], principal="*****@*****.**") rs = self.connect_and_query(auth_provider) connections = [c for holders in self.cluster.get_connection_holders() for c in holders.get_connections()] # Check to make sure our server_authenticator class is being set appropriate for connection in connections: self.assertTrue('DseAuthenticator' in connection.authenticator.server_authenticator_class) # Use invalid principal auth_provider = DSEGSSAPIAuthProvider(service='dse', qops=["auth"], principal="*****@*****.**") self.assertRaises(NoHostAvailable, self.connect_and_query, auth_provider)
def session_setup(self): self.cluster = Cluster(protocol_version=PROTOCOL_VERSION) self.session = self.cluster.connect() self.ks_name = self._testMethodName.lower() self.cass_version, self.cql_version = get_server_versions()
def main(): arguments = docopt(__doc__) host = [arguments['--host']] if arguments["--host"] else ["localhost"] session = Cluster(host).connect() graph = arguments['<keyspace>'] if graph and graph not in session.cluster.metadata.keyspaces: print "Graph {}{}{} not found".format(Fore.RED, graph, Style.RESET_ALL) sys.exit(1) session.default_graph_options.graph_name = graph print Fore.GREEN + "Connected to {}/{}".format(host[0], graph) + Style.RESET_ALL accum = None eof = None print "Gremlin REPL, use heredocs for multiline ex:<<EOF, help for help\n" while True: graph = session.default_graph_options.graph_name prompt = "gremlin [{}/{}]> ".format(host[0], graph) if eof is None else "gremlin (cont)> " input = raw_input(prompt) output_time = False start_time = time.time() if input.startswith("<<"): # heredoc print "Multiline mode activated" eof = input[2:] accum = [] continue if eof and input == eof: eof = None input = "\n".join(accum) print input elif eof: accum.append(input) continue if input == "quit" or input == "exit": break if input == "%schema": continue if input == "help": print_help() continue if input == "docs": webbrowser.open("http://docs.datastax.com/en/datastax_enterprise/5.0/datastax_enterprise/graph/graphTOC.html#graphTOC__missing-elem-id--graphTOC") continue total_time = None try: try: parsed = parse_line(input) result = parsed.execute(session) print_result_set(result) continue except ParseError as e: pass except Exception as e: print e continue stmt = SimpleGraphStatement(input) if input.startswith("a"): print Fore.GREEN + "Spark Graph Traversal Enabled, this may take a while..." + Style.RESET_ALL stmt.options.graph_source = "a" stmt.options.graph_alias = "a" start = time.time() result = session.execute_graph(stmt) total_time = time.time() - start except Exception as e: print e continue if isinstance(result, ResultSet): print_result_set(result) if total_time: print Fore.RED + "Query Time: {}s".format(round(total_time, 3)) + Style.RESET_ALL else: try: print "Unknown result", type(result), result except Exception as e: print e