From cfb3f1235848d000ba8609bf84e634bf56ac8291 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 9 Jan 2022 20:39:58 -0800
Subject: [PATCH] Only buffer input if --sniff, closes #364

---
 sqlite_utils/cli.py      | 13 ++++++++++---
 tests/test_cli_insert.py | 41 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index b235068..05a7bcc 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -757,13 +757,20 @@ def insert_upsert_implementation(
     if pk and len(pk) == 1:
         pk = pk[0]
     encoding = encoding or "utf-8-sig"
-    buffered = io.BufferedReader(file, buffer_size=4096)
-    decoded = io.TextIOWrapper(buffered, encoding=encoding)
+
+    # The --sniff option needs us to buffer the file to peek ahead
+    sniff_buffer = None
+    if sniff:
+        sniff_buffer = io.BufferedReader(file, buffer_size=4096)
+        decoded = io.TextIOWrapper(sniff_buffer, encoding=encoding)
+    else:
+        decoded = io.TextIOWrapper(file, encoding=encoding)
+
     tracker = None
     if csv or tsv:
         if sniff:
             # Read first 2048 bytes and use that to detect
-            first_bytes = buffered.peek(2048)
+            first_bytes = sniff_buffer.peek(2048)
             dialect = csv_std.Sniffer().sniff(first_bytes.decode(encoding, "ignore"))
         else:
             dialect = "excel-tab" if tsv else "excel"
diff --git a/tests/test_cli_insert.py b/tests/test_cli_insert.py
index c40dd85..675e9eb 100644
--- a/tests/test_cli_insert.py
+++ b/tests/test_cli_insert.py
@@ -2,6 +2,9 @@ from sqlite_utils import cli, Database
 from click.testing import CliRunner
 import json
 import pytest
+import subprocess
+import sys
+import time
 
 
 def test_insert_simple(tmpdir):
@@ -453,3 +456,41 @@ def test_insert_convert_row_modifying_in_place(db_path):
     db = Database(db_path)
     rows = list(db.query("select name, is_chicken from rows"))
     assert rows == [{"name": "Azi", "is_chicken": 1}]
+
+
+def test_insert_streaming_batch_size_1(db_path):
+    # https://github.com/simonw/sqlite-utils/issues/364
+    # Streaming with --batch-size 1 should commit on each record
+    # Can't use CliRunner().invoke() here bacuse we need to
+    # run assertions in between writing to process stdin
+    # First, create the DB with WAL mode enabled
+    CliRunner().invoke(cli.cli, ["create-database", db_path, "--enable-wal"])
+    # Now start streaming rows to insert --nl
+    proc = subprocess.Popen(
+        [
+            sys.executable,
+            "-m",
+            "sqlite_utils",
+            "insert",
+            db_path,
+            "rows",
+            "-",
+            "--nl",
+            "--batch-size",
+            "1",
+        ],
+        stdin=subprocess.PIPE,
+        stdout=sys.stdout,
+    )
+    proc.stdin.write(b'{"name": "Azi"}\n')
+    proc.stdin.flush()
+    # Without this delay the data wasn't yet visible
+    time.sleep(0.2)
+    assert list(Database(db_path)["rows"].rows) == [{"name": "Azi"}]
+    proc.stdin.write(b'{"name": "Suna"}\n')
+    proc.stdin.flush()
+    time.sleep(0.2)
+    assert list(Database(db_path)["rows"].rows) == [{"name": "Azi"}, {"name": "Suna"}]
+    proc.stdin.close()
+    proc.wait()
+    assert proc.returncode == 0