11#![ cfg( feature = "seed" ) ]
22
3+ use crate :: db:: connection:: DbPool ;
4+ use crate :: handlers:: user:: store_heartbeats_in_db;
35use crate :: models:: heartbeat:: { NewHeartbeat , SourceType } ;
46use crate :: models:: user:: User ;
5- use crate :: schema:: heartbeats;
67use chrono:: Utc ;
7- use diesel:: prelude:: * ;
88use ipnetwork:: { IpNetwork , Ipv4Network } ;
9- use rand:: Rng ;
109use rand:: prelude:: IndexedRandom ;
10+ use rand:: Rng ;
1111use std:: net:: Ipv4Addr ;
1212use tracing:: { info, warn} ;
1313
@@ -29,41 +29,38 @@ struct HeartbeatParams<'a> {
2929 user_agent : & ' a str ,
3030}
3131
32- pub fn seed_database ( conn : & mut PgConnection ) -> Result < ( ) , Box < dyn std:: error:: Error > > {
32+ pub async fn seed_database ( pool : & DbPool ) -> Result < ( ) , Box < dyn std:: error:: Error > > {
3333 info ! ( "🔄 Starting database seeding..." ) ;
3434
35- let Ok ( user) = create_dummy_user ( conn) else {
36- warn ! ( "⚠️ Dummy user already exists, skipping seeding." ) ;
37- return Ok ( ( ) ) ;
35+ let user = {
36+ let mut conn = pool. get ( ) ?;
37+
38+ if let Some ( existing_user) = User :: find_by_github_id ( & mut conn, -1 ) ? {
39+ warn ! ( "⚠️ Dummy user already exists (id: {}), skipping seeding." , existing_user. id) ;
40+ return Ok ( ( ) ) ;
41+ }
42+
43+ User :: create_or_update (
44+ & mut conn,
45+ -1 ,
46+ "Test User" ,
47+ "https://avatars.githubusercontent.com/u/999999" ,
48+ ) ?
3849 } ;
50+
3951 info ! (
4052 "✅ Created dummy user: {} (API Key: {})" ,
4153 user. name, user. api_key
4254 ) ;
4355
44- generate_random_heartbeats ( conn , user. id , TOTAL_HEARTBEATS ) ?;
56+ generate_random_heartbeats ( pool , user. id , TOTAL_HEARTBEATS ) . await ?;
4557
4658 info ! ( "✅ Database seeding completed successfully!" ) ;
4759 Ok ( ( ) )
4860}
49-
50- fn create_dummy_user ( conn : & mut PgConnection ) -> Result < User , Box < dyn std:: error:: Error > > {
51- if User :: find_by_github_id ( conn, -1 ) ?. is_some ( ) {
52- return Err ( "Dummy user already exists" . into ( ) ) ;
53- }
54-
55- let user = User :: create_or_update (
56- conn,
57- -1 ,
58- "Test User" ,
59- "https://avatars.githubusercontent.com/u/999999" ,
60- ) ?;
61-
62- Ok ( user)
63- }
64-
65- fn generate_random_heartbeats (
66- conn : & mut PgConnection ,
61+
62+ async fn generate_random_heartbeats (
63+ pool : & DbPool ,
6764 user_id : i32 ,
6865 count : usize ,
6966) -> Result < ( ) , Box < dyn std:: error:: Error > > {
@@ -81,18 +78,25 @@ fn generate_random_heartbeats(
8178 user_agent : USER_AGENT ,
8279 } ;
8380
84- let mut heartbeats = Vec :: with_capacity ( count) ;
81+ let mut batch = Vec :: with_capacity ( BATCH_SIZE ) ;
82+
8583 for _ in 0 ..count {
8684 let heartbeat = generate_random_heartbeat ( & mut rng, user_id, & params) ;
87- heartbeats. push ( heartbeat) ;
85+ batch. push ( heartbeat) ;
86+
87+ if batch. len ( ) == BATCH_SIZE {
88+ let current_batch: Vec < _ > = batch. drain ( ..) . collect ( ) ;
89+ store_heartbeats_in_db ( pool, current_batch) . await ?;
90+ }
8891 }
8992
90- for batch in heartbeats. chunks ( BATCH_SIZE ) {
91- diesel:: insert_into ( heartbeats:: table)
92- . values ( batch)
93- . execute ( conn) ?;
93+ if !batch. is_empty ( ) {
94+ let current_batch: Vec < _ > = batch. drain ( ..) . collect ( ) ;
95+ store_heartbeats_in_db ( pool, current_batch) . await ?;
9496 }
9597
98+ info ! ( "✅ Inserted {} heartbeats into the database" , count) ;
99+
96100 Ok ( ( ) )
97101}
98102
@@ -137,7 +141,7 @@ fn generate_random_heartbeat<R: Rng>(
137141 is_write : Some ( rng. random_bool ( 0.5 ) ) ,
138142 editor : Some ( "vscode" . to_string ( ) ) ,
139143 operating_system : Some ( "linux" . to_string ( ) ) ,
140- machine : "test-machine" . to_string ( ) ,
144+ machine : Some ( "test-machine" . to_string ( ) ) ,
141145 user_agent : params. user_agent . to_string ( ) ,
142146 lines : Some ( rng. random_range ( 1 ..=1000 ) ) ,
143147 project_root_count : None ,
@@ -147,5 +151,6 @@ fn generate_random_heartbeat<R: Rng>(
147151 lineno : Some ( rng. random_range ( 1 ..=100 ) ) ,
148152 cursorpos : Some ( rng. random_range ( 0 ..500 ) ) ,
149153 source_type : SourceType :: SEEDING . to_string ( ) . into ( ) ,
154+ project_id : None ,
150155 }
151156}
0 commit comments