sqlx_pg_seeder/
lib.rs

1use std::env::current_dir;
2use std::error::Error;
3use std::fs::{self};
4use std::io::Read;
5
6use serde_json::{self, Value};
7
8use sqlx::{Pool, Postgres};
9
10use chrono::{DateTime, FixedOffset};
11use sqlx::types::chrono;
12use sqlx::types::Uuid;
13
14use config::Config;
15
16struct Seeder {
17  file_names: Vec<String>,
18  table_names: Vec<String>,
19}
20
21trait SeederFn {
22  fn new() -> Self;
23}
24
25impl SeederFn for Seeder {
26  fn new() -> Self {
27    Seeder {
28      file_names: Vec::new(),
29      table_names: Vec::new(),
30    }
31  }
32}
33
34trait MyDateTimeEncode {
35  fn my_encode(&self) -> String;
36}
37
38impl MyDateTimeEncode for DateTime<FixedOffset> {
39  fn my_encode(&self) -> String {
40    self.to_rfc3339()
41  }
42}
43
44#[derive(Debug)]
45struct SeederConfig {
46  task_folder_path: String,
47  success_folder_path: String,
48  jsonb_name: String,
49  array_string_name: String,
50  created_at_name: String,
51  updated_at_name: String,
52}
53
54trait SeederConfigFn {
55  fn new() -> Self;
56}
57
58impl SeederConfigFn for SeederConfig {
59  fn new() -> Self {
60    SeederConfig {
61      task_folder_path: String::new(),
62      success_folder_path: String::new(),
63      jsonb_name: String::new(),
64      array_string_name: String::new(),
65      created_at_name: String::new(),
66      updated_at_name: String::new(),
67    }
68  }
69}
70
71fn read_config() -> Result<SeederConfig, config::ConfigError> {
72  let mut new_seeder_config = SeederConfig::new();
73  let settings = match Config::builder()
74    .add_source(config::File::with_name("pg-seeder"))
75    .build()
76  {
77    Ok(config) => config,
78    Err(err) => {
79      eprintln!("Error: Failed to load configuration: {}", err);
80      Config::default()
81    }
82  };
83
84  match settings.get::<String>("seeders.task_folder") {
85    Ok(value) => new_seeder_config.task_folder_path = value,
86    Err(_) => new_seeder_config.task_folder_path = "src/seeders/task".to_string(),
87  };
88  match settings.get::<String>("seeders.success_folder") {
89    Ok(value) => new_seeder_config.success_folder_path = value,
90    Err(_) => new_seeder_config.success_folder_path = "src/seeders/success".to_string(),
91  };
92  match settings.get::<String>("seeders.created_at_name") {
93    Ok(value) => new_seeder_config.created_at_name = value,
94    Err(_) => new_seeder_config.created_at_name = "created_at".to_string(),
95  };
96  match settings.get::<String>("seeders.updated_at_name") {
97    Ok(value) => new_seeder_config.updated_at_name = value,
98    Err(_) => new_seeder_config.updated_at_name = "updated_at".to_string(),
99  };
100  match settings.get::<String>("seeders.jsonb_name") {
101    Ok(value) => new_seeder_config.jsonb_name = value,
102    Err(_) => new_seeder_config.jsonb_name = "size".to_string(),
103  };
104  match settings.get::<String>("seeders.array_string_name") {
105    Ok(value) => new_seeder_config.array_string_name = value,
106    Err(_) => new_seeder_config.array_string_name = "thumbnail_src".to_string(),
107  };
108
109  Ok(new_seeder_config)
110}
111
112#[warn(unused_variables)]
113pub async fn seeder(pool: &Pool<Postgres>) -> Result<(), Box<dyn Error>> {
114  let seed_config = read_config().unwrap();
115  // println!("read config {:?}", seed_config.success_folder_path);
116
117  let mut new_seeder = Seeder::new();
118
119  let seeder_folder = current_dir()
120    .and_then(|a| Ok(a.join(seed_config.task_folder_path)))
121    .expect("No seed_folder exists");
122
123  let success_folder = current_dir()
124    .and_then(|a| Ok(a.join(seed_config.success_folder_path)))
125    .expect("No seed_folder exists");
126
127  if let Ok(entries) = fs::read_dir(&seeder_folder) {
128    for entry in entries {
129      if let Ok(entry) = entry {
130        if let Ok(file_name) = entry.file_name().into_string() {
131          new_seeder.file_names.push(file_name.to_string());
132
133          if let Some(ext) = file_name.split(".").last() {
134            if ext == "json" {
135              if let Some(first_part) = file_name.split(".").next() {
136                new_seeder.table_names.push(first_part.to_string());
137
138                let json_data = read_json_file();
139
140                for json_value in json_data {
141                  if let Some(field_value) = json_value.get(first_part) {
142                    let arr_field_value = field_value.as_array().unwrap();
143
144                    for each in arr_field_value {
145                      let mut field_names: Vec<&str> = Vec::new();
146                      let mut field_values: Vec<String> = Vec::new();
147
148                      if let Some(json_obj) = each.as_object() {
149                        for (key, value) in json_obj {
150                          field_names.push(key);
151
152                          field_values.push(value.to_string());
153                        }
154                      }
155                      // println!("Field Names: {:?}", &field_names);
156                      // println!("Field Values: {:?}", &field_values);
157
158                      let mut placeholders = String::new();
159
160                      for (idx, n) in (1..=field_values.len()).enumerate() {
161                        let field_name = &field_names[idx];
162
163                        let placeholder = match field_name {
164                          _ if field_name == &seed_config.jsonb_name => {
165                            format!("${}::JSONB", n)
166                          }
167
168                          field if field == &seed_config.array_string_name => {
169                            format!("${}::TEXT[]", n)
170                          }
171                          _ => format!("${}", n),
172                        };
173
174                        placeholders.push_str(&placeholder);
175
176                        if idx < field_values.len() - 1 {
177                          placeholders.push_str(", ");
178                        }
179                      }
180
181                      let query = format!(
182                        "insert into {} ({}) values ({})",
183                        &first_part,
184                        &field_names.join(", "),
185                        placeholders
186                      );
187
188                      // println!("postgres query : {:?}", &query);
189
190                      let mut query = sqlx::query(&query);
191
192                      for (index, value) in field_values.iter().enumerate() {
193                        match each.get(field_names[index]) {
194                          Some(json_value) => match json_value {
195                            Value::Bool(bool_value) => {
196                              query = query.bind(bool_value);
197                            }
198                            Value::Number(int_value) => {
199                              if let Some(i64_value) = int_value.as_i64() {
200                                query = query.bind(i64_value);
201                              } else if let Some(f64_value) = int_value.as_f64() {
202                                query = query.bind(f64_value);
203                              } else {
204                                println!("Number Error")
205                              }
206                            }
207
208                            Value::String(uuid_string) => {
209                              // println!("index {:?}", field_names[index]);
210                              match Uuid::parse_str(uuid_string) {
211                                Ok(uuid_value) => {
212                                  query = query.bind(uuid_value);
213                                }
214                                Err(_) => match field_names[index]
215                                  == seed_config.created_at_name
216                                  || field_names[index] == seed_config.updated_at_name
217                                {
218                                  true => {
219                                    if let Ok(timestamp) =
220                                      chrono::DateTime::parse_from_rfc3339(uuid_string)
221                                    {
222                                      query = query.bind(timestamp);
223                                      // println!(
224                                      //   "string: created_at!!! true, filed_name : {:?}",
225                                      //   field_names[index]
226                                      // )
227                                    }
228                                  }
229                                  false => {
230                                    query = query.bind(uuid_string);
231                                    // println!(
232                                    //   "string: created_at!!! false, filed_name : {:?}",
233                                    //   field_names[index]
234                                    // )
235                                  }
236                                },
237                              }
238                            }
239                            Value::Array(array_value) => {
240                              // println!("array {:?}", field_names[index]);
241
242                              query = query.bind(array_value);
243                            }
244                            Value::Object(obj_value) => {
245                              // println!("JSONB!!!! {:?}", field_names[index]);
246                              let json_string = serde_json::to_string(obj_value)
247                                .expect("Failed to serialize JSON object to string");
248
249                              // Bind the JSON string to the SQL query
250                              query = query.bind(json_string);
251                            }
252                            _ => {
253                              query = query.bind(value);
254                            }
255                          },
256                          None => {
257                            println!("Seeder Error!")
258                          }
259                        }
260                      }
261                      // 쿼리 실행
262                      query.execute(pool).await.unwrap();
263                    }
264                  }
265                }
266              }
267            }
268          }
269          println!("✅ Seed completed for the {:?}.json", file_name);
270        }
271
272        let new_path = success_folder.join(entry.file_name());
273        // println!("success folder : {:?}", new_path);
274        if let Err(err) = fs::rename(entry.path(), new_path) {
275          println!("Failed to move file: {}", err);
276        };
277      } else {
278        println!("Failed to read the directory.");
279      }
280    }
281  }
282
283  // println!(
284  //   "file_name: {:?} parts : {:?}",
285  //   new_seeder.file_names, new_seeder.table_names
286  // );
287
288  print!("✅ Seeder Work Success! ✅ \n");
289  Ok(())
290}
291
292pub fn read_json_file() -> Vec<Value> {
293  let dir_path = current_dir()
294    .expect("Can't retreive directory")
295    .join("src/seeders/task");
296
297  let mut json_values: Vec<Value> = Vec::new();
298
299  if let Ok(entries) = fs::read_dir(dir_path) {
300    for entry in entries {
301      if let Ok(entry) = entry {
302        let file_path = entry.path();
303
304        if let Some(ext) = file_path.extension() {
305          if ext == "json" {
306            let mut file = fs::File::open(&file_path).expect("Cannot open the file");
307
308            let mut contents = String::new();
309            file
310              .read_to_string(&mut contents)
311              .expect("There was an issue reading the file");
312
313            let json_value: Value =
314              serde_json::from_str(&contents).expect("Failed to parse JSON");
315
316            json_values.push(json_value);
317          }
318        }
319      }
320    }
321  }
322
323  json_values
324}