Modify SQLite plugin to be lazy

turicas · Sep 11, 2017 · c22f6f7 · c22f6f7
1 parent c0d2df4
commit c22f6f7
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 7 deletions.
diff --git a/rows/plugins/sqlite.py b/rows/plugins/sqlite.py
@@ -21,6 +21,8 @@
 import sqlite3
 import string
 
+from itertools import chain
+
 import six
 
 import rows.fields as fields
@@ -31,6 +33,7 @@
 
 
 SQL_TABLE_NAMES = 'SELECT name FROM sqlite_master WHERE type="table"'
+# TODO: may use query args instead of string formatting
 SQL_CREATE_TABLE = 'CREATE TABLE IF NOT EXISTS "{table_name}" ({field_types})'
 SQL_SELECT_ALL = 'SELECT * FROM "{table_name}"'
 SQL_INSERT = 'INSERT INTO "{table_name}" ({field_names}) VALUES ({placeholders})'
@@ -117,13 +120,12 @@ def import_from_sqlite(filename_or_connection, table_name='table1', query=None,
     if query_args is None:
         query_args = tuple()
 
-    table_rows = list(cursor.execute(query, query_args)) # TODO: may be lazy
-    header = [six.text_type(info[0]) for info in cursor.description]
-    cursor.close()
-    # TODO: should close connection also?
+    cursor.execute(query, query_args)
+    data = chain([[six.text_type(info[0]) for info in cursor.description]],
+                 cursor)
 
     meta = {'imported_from': 'sqlite', 'filename': filename_or_connection, }
-    return create_table([header] + table_rows, meta=meta, *args, **kwargs)
+    return create_table(data, meta=meta, *args, **kwargs)
 
 
 def export_to_sqlite(table, filename_or_connection, table_name=None,
@@ -159,6 +161,7 @@ def export_to_sqlite(table, filename_or_connection, table_name=None,
             placeholders=', '.join('?' for _ in field_names))
     _convert_row = _python_to_sqlite(field_types)
     for batch in ipartition(prepared_table, batch_size):
+        # TODO: add a callback function
         cursor.executemany(insert_sql, map(_convert_row, batch))
 
     connection.commit()

diff --git a/tests/tests_plugin_sqlite.py b/tests/tests_plugin_sqlite.py
@@ -101,6 +101,24 @@ def test_export_to_sqlite_filename(self):
         table = rows.import_from_sqlite(temp.name)
         self.assert_table_equal(table, utils.table)
 
+    def test_import_from_sqlite_is_lazy(self): #, mocked_create_table):
+        connection = mock.MagicMock()
+        cursor = connection.cursor()
+        cursor.description = [('f1', None), ('f2', None), ('f3', None)]
+        gen = utils.LazyGenerator(max_number=1000)
+        igen = iter(gen)
+        next(igen)  # get header out -- does not happen on SQLite
+        cursor.__iter__.return_value = igen
+        cursor.fetchall = lambda: list(gen)
+
+        table = rows.import_from_sqlite(connection, lazy=True, samples=50)
+        self.assertIs(gen.last, 49)
+
+        for index, _ in enumerate(table):
+            if index == 99:
+                break
+        self.assertIs(gen.last, 99)
+
     def test_export_to_sqlite_connection(self):
         # TODO: may test file contents
         temp = tempfile.NamedTemporaryFile(delete=False, mode='wb')
@@ -124,9 +142,9 @@ def test_export_to_sqlite_create_unique_table_name(self):
         rows.export_to_sqlite(second_table, temp.name)  # table2
 
         result_first_table = rows.import_from_sqlite(temp.name,
-                                                     table_name='table1')
+                table_name='table1')
         result_second_table = rows.import_from_sqlite(temp.name,
-                                                      table_name='table2')
+                table_name='table2')
         self.assert_table_equal(result_first_table, first_table)
         self.assert_table_equal(result_second_table, second_table)