From 7c8f7ab600b8f014c0263192100325b028c8ecea Mon Sep 17 00:00:00 2001
From: Bo-Chun Louis Chen <louistw@uab.edu>
Date: Wed, 27 Apr 2022 17:04:15 -0500
Subject: [PATCH] Add insert many method
---
init_user_state.py | 20 ++++++++++++++++++++
1 file changed, 20 insertions(+)
diff --git a/init_user_state.py b/init_user_state.py
index a700082..d18724e 100644
--- a/init_user_state.py
+++ b/init_user_state.py
@@ -15,6 +15,9 @@ args = parser.parse_args()
default_state = "ok"
+# Chunk size for insert into db
+size = 1000
+
db = dataset.connect(f"sqlite:///prod_rmq_agents/{rcfg.db_path}/user_reg.db")
table = db["user_state"]
@@ -28,6 +31,23 @@ users = subprocess.run(
).stdout.split()
# Update user_state table
+# Insert many
+if len(users) > 50:
+ start = 0
+ today = datetime.now()
+ while start < len(users):
+ end = start + size if start + size < len(users) else len(users)
+ data = [
+ dict(username=user, state=default_state, date=today)
+ for user in users[start:end]
+ ]
+ if args.dry_run:
+ print(f"Table insert many from {start} to {end - 1}")
+ else:
+ table.insert_many(data, chunk_size=size)
+ start = end
+
+# Insert one by one
for user in users:
if args.dry_run:
print(f"Table insert user: {user}, state: {default_state}")
--
GitLab