fix_ingest_csv.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. import re
  2. import os
  3. file_path = "ingest_csv.py"
  4. with open(file_path, "r", encoding="utf-8", errors="replace") as f:
  5. content = f.read()
  6. # Replace encoding issues in header
  7. content = content.replace("Lange FranA ois", "Francois Lange")
  8. content = content.replace("Lange FranAois", "Francois Lange")
  9. content = content.replace("Lange François", "Francois Lange")
  10. # 1. Add init_ingestion_status_table function definition right before ingest_file
  11. old_def = "def ingest_file(filename, engine):"
  12. new_def = """def init_ingestion_status_table(engine):
  13. with engine.begin() as conn:
  14. conn.execute(text(\"\"\"
  15. CREATE TABLE IF NOT EXISTS ingestion_status (
  16. id INT AUTO_INCREMENT PRIMARY KEY,
  17. filename VARCHAR(255),
  18. start_time DATETIME,
  19. end_time DATETIME,
  20. rows_loaded INT,
  21. status VARCHAR(50),
  22. error_message TEXT
  23. )
  24. \"\"\"))
  25. def ingest_file(filename, engine):"""
  26. content = content.replace(old_def, new_def)
  27. # 2. Add table initialization and insert running status in ingest_file
  28. old_start = """ print(f"\\ndYs? Found {filename}! Starting grouped vertical partition ingestion...")
  29. chunk_size = 10000
  30. total_processed = 0"""
  31. new_start = """ print(f"\\ndYs? Found {filename}! Starting grouped vertical partition ingestion...")
  32. # Initialize ingestion status tracking
  33. init_ingestion_status_table(engine)
  34. import datetime
  35. start_time = datetime.datetime.now()
  36. with engine.begin() as conn:
  37. conn.execute(text(\"\"\"
  38. INSERT INTO ingestion_status (filename, start_time, rows_loaded, status)
  39. VALUES (:filename, :start_time, 0, 'RUNNING')
  40. \"\"\"), {"filename": filename, "start_time": start_time})
  41. last_id_res = conn.execute(text("SELECT LAST_INSERT_ID()"))
  42. ingest_id = last_id_res.scalar()
  43. chunk_size = 10000
  44. total_processed = 0"""
  45. content = content.replace(old_start, new_start)
  46. # 3. Add milestone updates inside chunk loop
  47. old_milestone = """ total_processed += len(df)
  48. print(f" Successfully appended {total_processed} rows into grouped tables...", end="\\r")
  49. if total_processed % 50000 == 0:
  50. notifier.send_alert(f"Ingestion Milestone: {total_processed} rows processed")"""
  51. new_milestone = """ total_processed += len(df)
  52. print(f" Successfully appended {total_processed} rows into grouped tables...", end="\\r")
  53. # Update rows loaded in database
  54. with engine.begin() as conn:
  55. conn.execute(text(\"\"\"
  56. UPDATE ingestion_status
  57. SET rows_loaded = :rows
  58. WHERE id = :id
  59. \"\"\"), {"rows": total_processed, "id": ingest_id})
  60. if total_processed % 50000 == 0:
  61. notifier.send_alert(f"Ingestion Milestone: {total_processed} rows processed")"""
  62. content = content.replace(old_milestone, new_milestone)
  63. # 4. Add success and failure callbacks at the end
  64. old_end = """ except BaseException as e:
  65. notifier.send_alert(f"Ingestion Exception: {str(e)}")
  66. print(f"\\n [Warning] Chunk skipped due to error: {e}")
  67. notifier.send_alert(f"Ingestion Finished: {filename}")
  68. print(f"\\no. Finished importing {filename}.")
  69. return True"""
  70. new_end = """ except BaseException as e:
  71. end_time = datetime.datetime.now()
  72. with engine.begin() as conn:
  73. conn.execute(text(\"\"\"
  74. UPDATE ingestion_status
  75. SET end_time = :end_time, status = 'FAILED', error_message = :err
  76. WHERE id = :id
  77. \"\"\"), {"end_time": end_time, "err": str(e), "id": ingest_id})
  78. notifier.send_alert(f"Ingestion Exception: {str(e)}")
  79. print(f"\\n [Warning] Chunk skipped due to error: {e}")
  80. end_time = datetime.datetime.now()
  81. with engine.begin() as conn:
  82. conn.execute(text(\"\"\"
  83. UPDATE ingestion_status
  84. SET end_time = :end_time, status = 'COMPLETED'
  85. WHERE id = :id
  86. \"\"\"), {"end_time": end_time, "id": ingest_id})
  87. notifier.send_alert(f"Ingestion Finished: {filename}")
  88. print(f"\\n. Finished importing {filename}.")
  89. return True"""
  90. content = content.replace(old_end, new_end)
  91. with open(file_path, "w", encoding="utf-8") as f:
  92. f.write(content)
  93. print("ingest_csv.py successfully updated.")