From 50589f8523751191559e8d812c2ee0889da06e50 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 27 Jan 2019 22:26:45 -0800
Subject: [PATCH] CLI insert tool now uses generators, closes #7

Also cleaned up the logic so we commit rows in batches too.
---
 sqlite_utils/cli.py |  8 ++++----
 sqlite_utils/db.py  | 21 +++++++++++----------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index c6117ad..17d7667 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -72,17 +72,17 @@ def optimize(path, no_vacuum):
 @click.argument("json_file", type=click.File(), required=True)
 @click.option("--pk", help="Column to use as the primary key, e.g. id")
 @click.option("--nl", is_flag=True, help="Expect newline-delimited JSON")
-def insert(path, table, json_file, pk, nl):
+@click.option("--batch-size", type=int, default=100, help="Commit every X records")
+def insert(path, table, json_file, pk, nl, batch_size):
     "Insert records from JSON file into the table, create table if it is missing"
     db = sqlite_utils.Database(path)
     if nl:
-        # TODO: Use a generator once #7 is solved
-        docs = [json_std.loads(line) for line in json_file]
+        docs = (json_std.loads(line) for line in json_file)
     else:
         docs = json_std.load(json_file)
         if isinstance(docs, dict):
             docs = [docs]
-    db[table].insert_all(docs, pk=pk)
+    db[table].insert_all(docs, pk=pk, batch_size=batch_size)
 
 
 @cli.command()
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index e3a042c..1b806bf 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -185,13 +185,14 @@ class Table:
 
     def create(self, columns, pk=None, foreign_keys=None, column_order=None):
         columns = {name: value for (name, value) in columns.items()}
-        self.db.create_table(
-            self.name,
-            columns,
-            pk=pk,
-            foreign_keys=foreign_keys,
-            column_order=column_order,
-        )
+        with self.db.conn:
+            self.db.create_table(
+                self.name,
+                columns,
+                pk=pk,
+                foreign_keys=foreign_keys,
+                column_order=column_order,
+            )
         self.exists = True
         return self
 
@@ -384,9 +385,9 @@ class Table:
                 values.extend(
                     jsonify_if_needed(record.get(key, None)) for key in all_columns
                 )
-            result = self.db.conn.execute(sql, values)
-            self.db.conn.commit()
-            self.last_id = result.lastrowid
+            with self.db.conn:
+                result = self.db.conn.execute(sql, values)
+                self.last_id = result.lastrowid
         return self
 
     def upsert(self, record, pk=None, foreign_keys=None, column_order=None):