Python/SQLAlchemy: Fix multirow/batched exercise (#1189)

Large inserts without using batch are source of OOMemory, so CrateDB
introduced a setting (configurable) to have a protective limit.

In this spirit, large inserts using the "multirow" strategy no longer
work, so accompany the test harness to only use them with "batched".

In this case, large is ~50_000 records, while ~25_000 records of
the shape at hand still work with traditional "multirow" inserts.

Author: amotl

by-language/python-sqlalchemy/insert_efficient.py

@@ -53,9 +53,6 @@
 
 import sqlalchemy as sa
 
-# INSERT_RECORDS = 1275
-INSERT_RECORDS = 50_000
-# INSERT_RECORDS = 2_750_000
 
 BATCHED_PAGE_SIZE = 20_000
 
@@ -88,7 +85,7 @@ def insert_batched(engine, table, records):
         conn.execute(insertable, parameters=records)
 
 
-def run_example(dburi: str, variant: str):
+def run_example(dburi: str, variant: str, record_count: int):
     metadata = sa.MetaData()
     table = sa.Table(
         "testdrive",
@@ -98,7 +95,7 @@ def run_example(dburi: str, variant: str):
     )
 
     # Create 275 test records.
-    records = [{"id": i, "name": f"foo_{i}"} for i in range(1, INSERT_RECORDS + 1)]
+    records = [{"id": i, "name": f"foo_{i}"} for i in range(1, record_count + 1)]
 
     # Run multi-row insert, with a specified batch-/page-size.
     engine = sa.create_engine(dburi, insertmanyvalues_page_size=BATCHED_PAGE_SIZE, echo=True)
@@ -119,7 +116,7 @@ def run_example(dburi: str, variant: str):
         print("Number of records:", result.scalar_one())
 
 
-def run_database(database: str, variant: str):
+def run_database(database: str, variant: str, record_count: int):
     if database == "sqlite":
         dburi = "sqlite:///:memory:"
     elif database == "postgresql":
@@ -129,10 +126,17 @@ def run_database(database: str, variant: str):
     else:
         raise ValueError(f"Unknown database: {database}")
 
-    run_example(dburi, variant)
+    run_example(dburi, variant, record_count)
 
 
 if __name__ == "__main__":
+    """
+    Invocation examples:
+
+    python insert_efficient.py cratedb multirow 25000
+    python insert_efficient.py cratedb batched 50000
+    """
     database = sys.argv[1]
     variant = sys.argv[2]
-    run_database(database, variant)
+    record_count = int(sys.argv[3])
+    run_database(database, variant, record_count)

by-language/python-sqlalchemy/test.py

@@ -8,17 +8,19 @@ def run(command: str):
 
 
 def test_insert_efficient_multirow():
-    cmd = "time python insert_efficient.py cratedb multirow"
+    insert_records = 25_000
+    cmd = f"time python insert_efficient.py cratedb multirow {insert_records}"
     run(cmd)
 
 
 def test_insert_efficient_batched():
-    cmd = "time python insert_efficient.py cratedb batched"
+    insert_records = 50_000
+    cmd = f"time python insert_efficient.py cratedb batched {insert_records}"
     run(cmd)
 
 
 def test_insert_efficient_unknown(capfd):
-    cmd = "time python insert_efficient.py cratedb unknown"
+    cmd = "time python insert_efficient.py cratedb unknown 1000"
     with pytest.raises(subprocess.CalledProcessError):
         run(cmd)
     out, err = capfd.readouterr()