Merge branch 'cassandra-2.1' into trunk

CHANGES.txt

@@ -17,6 +17,7 @@
 
 
 2.1.1
+ * (cqlsh) COPY TO/FROM improvements (CASSANDRA-7405)
  * Support list index operations with conditions (CASSANDRA-7499)
  * Add max live/tombstoned cells to nodetool cfstats output (CASSANDRA-7731)
  * Validate IPv6 wildcard addresses properly (CASSANDRA-7680)

bin/cqlsh

@@ -118,7 +118,7 @@ cqlshlibdir = os.path.join(CASSANDRA_PATH, 'pylib')
 if os.path.isdir(cqlshlibdir):
     sys.path.insert(0, cqlshlibdir)
 
-from cqlshlib import cqlhandling, cql3handling, pylexotron, sslhandling
+from cqlshlib import cqlhandling, cql3handling, pylexotron, sslhandling, async_insert
 from cqlshlib.displaying import (RED, BLUE, CYAN, ANSI_RESET, COLUMN_NAME_COLORS,
                                  FormattedValue, colorme)
 from cqlshlib.formatting import format_by_type, formatter_for, format_value_utype
@@ -1348,27 +1348,33 @@ class Shell(cmd.Cmd):
             if header:
                 linesource.next()
             table_meta = self.get_table_meta(ks, cf)
-            rownum = -1
             reader = csv.reader(linesource, **dialect_options)
-            for rownum, row in enumerate(reader):
+            from functools import partial
-                if len(row) != len(columns):
+            rownum, error = \
-                    self.printerr("Record #%d (line %d) has the wrong number of fields "
+                async_insert.insert_concurrent(self.session, enumerate(reader, start=1),
-                                  "(%d instead of %d)."
+                                               partial(
-                                  % (rownum, reader.line_num, len(row), len(columns)))
+                                                   self.create_insert_statement,
-                    return rownum
+                                                   columns, nullval,
-                if not self.do_import_row(columns, nullval, table_meta, row):
+                                                   table_meta))
-                    self.printerr("Aborting import at record #%d (line %d). "
+            if error:
-                                  "Previously-inserted values still present."
+                self.printerr(str(error[0]))
-                                  % (rownum, reader.line_num))
+                self.printerr("Aborting import at record #%d. "
-                    return rownum
+                              "Previously-inserted values still present."
+                               % error[1])
         finally:
             if do_close:
                 linesource.close()
             elif self.tty:
                 print
-        return rownum + 1
+        return rownum-1
+
+    def create_insert_statement(self, columns, nullval, table_meta, row):
+
+        if len(row) != len(columns):
+            raise ValueError(
+                "Record has the wrong number of fields (%d instead of %d)."
+                % (len(row), len(columns)))
 
-    def do_import_row(self, columns, nullval, table_meta, row):
         rowmap = {}
         primary_key_columns = [col.name for col in table_meta.primary_key]
         for name, value in zip(columns, row):
@@ -1390,9 +1396,6 @@ class Shell(cmd.Cmd):
                 return False
             else:
                 rowmap[name] = 'null'
-        return self.do_import_insert(table_meta, rowmap)
-
-    def do_import_insert(self, table_meta, rowmap):
         # would be nice to be able to use a prepared query here, but in order
         # to use that interface, we'd need to have all the input as native
         # values already, reading them from text just like the various
@@ -1406,7 +1409,8 @@ class Shell(cmd.Cmd):
         )
         if self.debug:
             print 'Import using CQL: %s' % query
-        return self.perform_simple_statement(SimpleStatement(query))
+        return SimpleStatement(query)
+
 
     def perform_csv_export(self, ks, cf, columns, fname, opts):
         dialect_options = self.csv_dialect_defaults.copy()
@@ -1460,13 +1464,7 @@ class Shell(cmd.Cmd):
         if columns is None:
             columns = self.get_column_names(ks, cf)
         columnlist = ', '.join(protect_names(columns))
-        # this limit is pretty awful. would be better to use row-key-paging, so
+        query = 'SELECT %s FROM %s.%s' % (columnlist, protect_name(ks), protect_name(cf))
-        # that the dump could be pretty easily aborted if necessary, but that
-        # can be kind of tricky with cql3. Punt for now, until the real cursor
-        # API is added in CASSANDRA-4415.
-        # https://datastax-oss.atlassian.net/browse/PYTHON-16
-        query = 'SELECT %s FROM %s.%s LIMIT 99999999' \
-                % (columnlist, protect_name(ks), protect_name(cf))
         return self.session.execute(query)
 
     def do_show(self, parsed):

pylib/cqlshlib/async_insert.py

@@ -0,0 +1,110 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from itertools import count
+from threading import Event, Condition
+import sys
+
+
+class _CountDownLatch(object):
+    def __init__(self, counter=1):
+        self._count = counter
+        self._lock = Condition()
+
+    def count_down(self):
+        with self._lock:
+            self._count -= 1
+            if self._count <= 0:
+                self._lock.notifyAll()
+
+    def await(self):
+        with self._lock:
+            while self._count > 0:
+                self._lock.wait()
+
+
+class _ChainedWriter(object):
+
+    CONCURRENCY = 100
+
+    def __init__(self, session, enumerated_reader, statement_func):
+        self._sentinel = object()
+        self._session = session
+        self._cancellation_event = Event()
+        self._first_error = None
+        self._num_finished = count(start=1)
+        self._task_counter = _CountDownLatch(self.CONCURRENCY)
+        self._enumerated_reader = enumerated_reader
+        self._statement_func = statement_func
+
+    def insert(self):
+        if not self._enumerated_reader:
+            return 0, None
+
+        for i in xrange(self.CONCURRENCY):
+            self._execute_next(self._sentinel, 0)
+
+        self._task_counter.await()
+        return next(self._num_finished), self._first_error
+
+    def _abort(self, error, failed_record):
+        if not self._first_error:
+            self._first_error = error, failed_record
+        self._task_counter.count_down()
+        self._cancellation_event.set()
+
+    def _handle_error(self, error, failed_record):
+        self._abort(error, failed_record)
+
+    def _execute_next(self, result, last_completed_record):
+        if self._cancellation_event.is_set():
+            self._task_counter.count_down()
+            return
+
+        if result is not self._sentinel:
+            finished = next(self._num_finished)
+            if not finished % 1000:
+                sys.stdout.write('Imported %s rows\r' % finished)
+                sys.stdout.flush()
+
+        try:
+            (current_record, row) = next(self._enumerated_reader)
+        except StopIteration:
+            self._task_counter.count_down()
+            return
+        except Exception as exc:
+            self._abort(exc, last_completed_record)
+            return
+
+        if self._cancellation_event.is_set():
+            self._task_counter.count_down()
+            return
+
+        try:
+            statement = self._statement_func(row)
+            future = self._session.execute_async(statement)
+            future.add_callbacks(callback=self._execute_next,
+                                 callback_args=(current_record,),
+                                 errback=self._handle_error,
+                                 errback_args=(current_record,))
+        except Exception as exc:
+            self._abort(exc, current_record)
+            return
+
+
+def insert_concurrent(session, enumerated_reader, statement_func):
+    return _ChainedWriter(session, enumerated_reader, statement_func).insert()
+