From 833ddd9411ac6f89ef1ddabc23d8e4d4a2ac1610 Mon Sep 17 00:00:00 2001 From: yzqzss Date: Sun, 28 Jul 2024 12:35:48 +0800 Subject: [PATCH] --id-name --- stwp_mongodb_helper/grass_generator.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/stwp_mongodb_helper/grass_generator.py b/stwp_mongodb_helper/grass_generator.py index dc97474..784af72 100644 --- a/stwp_mongodb_helper/grass_generator.py +++ b/stwp_mongodb_helper/grass_generator.py @@ -21,6 +21,7 @@ def arg_parser(): parser.add_argument("--chunk-size", help="Chunk size", type=int, required=True) parser.add_argument("--end", help="End ID", type=int, required=True) parser.add_argument("--delay", help="Delay in seconds", type=float, required=False, default=1.0) + parser.add_argument("--id-name", help="ID name", default="id", choices=["id", "feed_id"]) return parser.parse_args() def qos(delay: float): @@ -34,6 +35,7 @@ def main(): c_queue_name = args.co chunk_size: int = args.chunk_size delay: float = args.delay + ID_NAME = args.id_name client = pymongo.MongoClient(os.environ['MONGODB_URI']) dbs = client.list_database_names() @@ -55,7 +57,7 @@ def main(): print(todos_now, "TODO documents in queue") if todos_now > chunk_size * 5: continue - max_id_now = find_max_id(coll, "id") or 0 + max_id_now = find_max_id(coll, ID_NAME) or 0 if tqd is None: tqd = tqdm.tqdm(total=args.end, initial=max_id_now) tqd.n = max_id_now @@ -68,11 +70,11 @@ def main(): docs = [] for i in range(1,chunk_size+1): docs.append({ - "id": max_id_now + i, + ID_NAME: max_id_now + i, "status": "TODO" }) print("Inserting", chunk_size, "documents") - print([doc["id"] for doc in docs]) + print([doc[ID_NAME] for doc in docs]) coll.insert_many(docs)