-
Notifications
You must be signed in to change notification settings - Fork 63
/
util.rs
1583 lines (1436 loc) · 55.1 KB
/
util.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#[cfg(any(feature = "feature_capable", feature = "lite"))]
use std::borrow::Cow;
use std::{
cmp::min,
env, fs,
fs::File,
io::{BufReader, BufWriter, Read, Write},
path::{Path, PathBuf},
str,
sync::OnceLock,
time::SystemTime,
};
use docopt::Docopt;
#[cfg(any(feature = "feature_capable", feature = "lite"))]
use indicatif::{HumanCount, ProgressBar, ProgressDrawTarget, ProgressStyle};
use log::log_enabled;
use reqwest::Client;
use serde::de::DeserializeOwned;
#[cfg(any(feature = "feature_capable", feature = "lite"))]
use serde::de::{Deserialize, Deserializer, Error};
use sysinfo::System;
use crate::{
config,
config::{Config, Delimiter, DEFAULT_WTR_BUFFER_CAPACITY},
CliError, CliResult, CURRENT_COMMAND,
};
#[macro_export]
macro_rules! regex_oncelock {
($re:literal $(,)?) => {{
static RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
RE.get_or_init(|| regex::Regex::new($re).unwrap())
}};
}
// leave at least 20% of the available memory free
const DEFAULT_FREEMEMORY_HEADROOM_PCT: u8 = 20;
static ROW_COUNT: OnceLock<Option<u64>> = OnceLock::new();
pub type ByteString = Vec<u8>;
#[inline]
pub fn num_cpus() -> usize {
num_cpus::get()
}
const CARGO_BIN_NAME: &str = env!("CARGO_BIN_NAME");
const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
const TARGET: &str = match option_env!("TARGET") {
Some(target) => target,
None => "Unknown_target",
};
const QSV_KIND: &str = match option_env!("QSV_KIND") {
Some(kind) => kind,
None => "installed",
};
fn default_user_agent() -> String {
let unknown_command = "Unknown".to_string();
let current_command = CURRENT_COMMAND.get().unwrap_or(&unknown_command);
format!("{CARGO_BIN_NAME}/{CARGO_PKG_VERSION} ({TARGET}; {current_command}; {QSV_KIND}; https://github.com/jqnatividad/qsv)")
}
pub fn max_jobs() -> usize {
let num_cpus = num_cpus();
let max_jobs = match env::var("QSV_MAX_JOBS") {
Ok(val) => val.parse::<usize>().unwrap_or(1_usize),
Err(_) => num_cpus,
};
if (1..=num_cpus).contains(&max_jobs) {
max_jobs
} else {
num_cpus
}
}
/// Given a desired number of cores to use
/// returns number of cores to actually use.
/// If desired is None, zero, or greater than available cores,
/// returns max_jobs, which is equal to number of available cores
/// If desired is Some and less than available cores,
/// returns desired number of cores
pub fn njobs(flag_jobs: Option<usize>) -> usize {
let max_jobs = max_jobs();
let jobs_to_use = flag_jobs.map_or(max_jobs, |jobs| {
if jobs == 0 || jobs > max_jobs {
max_jobs
} else {
jobs
}
});
env::set_var("RAYON_NUM_THREADS", jobs_to_use.to_string());
log::info!("Using {jobs_to_use} jobs...");
jobs_to_use
}
pub fn timeout_secs(timeout: u16) -> Result<u64, String> {
let timeout = match env::var("QSV_TIMEOUT") {
Ok(val) => val.parse::<u16>().unwrap_or(30_u16),
Err(_) => timeout,
};
if timeout > 3600 {
return fail_format!("Timeout cannot be more than 3,600 seconds (1 hour): {timeout}");
} else if timeout == 0 {
return fail!("Timeout cannot be zero.");
}
log::info!("TIMEOUT: {timeout}");
Ok(timeout as u64)
}
/// sets custom user agent
/// if user agent is not set, then use the default user agent
/// it supports four special LITERALs: $QSV_BIN_NAME, $QSV_VERSION, $QSV_TARGET, $QSV_KIND
/// and $QSV_COMMAND which will be replaced with the actual values during runtime
pub fn set_user_agent(user_agent: Option<String>) -> CliResult<String> {
use reqwest::header::HeaderValue;
let ua = match user_agent {
Some(ua_arg) => ua_arg,
None => env::var("QSV_USER_AGENT").unwrap_or_else(|_| default_user_agent()),
};
let unknown_command = "Unknown".to_string();
let current_command = CURRENT_COMMAND.get().unwrap_or(&unknown_command);
// look for special literals - $QSV_VERSION and $QSV_TARGET and replace them
let ua = ua
.replace("$QSV_BIN_NAME", CARGO_BIN_NAME)
.replace("$QSV_VERSION", CARGO_PKG_VERSION)
.replace("$QSV_TARGET", TARGET)
.replace("$QSV_KIND", QSV_KIND)
.replace("$QSV_COMMAND", current_command);
match HeaderValue::from_str(ua.as_str()) {
Ok(_) => (),
Err(e) => return fail_incorrectusage_clierror!("Invalid user-agent value: {e}"),
};
log::info!("set user agent: {ua}");
Ok(ua)
}
pub fn version() -> String {
let mut enabled_features = String::new();
#[cfg(all(feature = "apply", not(feature = "lite")))]
enabled_features.push_str("apply;");
#[cfg(all(feature = "fetch", not(feature = "lite")))]
enabled_features.push_str("fetch;");
#[cfg(all(feature = "foreach", not(feature = "lite")))]
enabled_features.push_str("foreach;");
#[cfg(all(feature = "geocode", not(feature = "lite")))]
enabled_features.push_str("geocode;");
#[cfg(all(feature = "luau", not(feature = "lite")))]
{
let luau = mlua::Lua::new();
match luau.load("return _VERSION").eval() {
Ok(version_info) => {
if let mlua::Value::String(luaustring_val) = version_info {
let string_val = luaustring_val.to_str().unwrap_or("Luau - invalid version");
if string_val == "Luau" {
enabled_features.push_str("Luau - version not specified;");
} else {
enabled_features.push_str(&format!("{string_val};"));
}
} else {
enabled_features.push_str("Luau - ?;");
}
},
Err(e) => enabled_features.push_str(&format!("Luau - cannot retrieve version: {e};")),
};
}
#[cfg(all(feature = "python", not(feature = "lite")))]
{
enabled_features.push_str("python-");
pyo3::Python::with_gil(|py| {
enabled_features.push_str(py.version());
enabled_features.push(';');
});
}
#[cfg(all(feature = "to", not(feature = "lite")))]
enabled_features.push_str("to;");
#[cfg(all(feature = "polars", not(feature = "lite")))]
enabled_features.push_str(format!("polars-{};", polars::VERSION).as_str());
#[cfg(feature = "self_update")]
enabled_features.push_str("self_update");
enabled_features.push('-');
// get max_file_size & memory info. max_file_size is based on QSV_FREEMEMORY_HEADROOM_PCT
// setting and is only enforced when qsv is running in "non-streaming" mode (i.e. needs to
// load the entire file into memory).
let mut sys = System::new();
sys.refresh_memory();
let avail_mem = sys.available_memory();
let total_mem = sys.total_memory();
let free_swap = sys.free_swap();
let max_file_size = mem_file_check(Path::new(""), true, false).unwrap_or(0) as u64;
#[cfg(feature = "mimalloc")]
let malloc_kind = "mimalloc";
#[cfg(feature = "jemallocator")]
let malloc_kind = "jemalloc";
#[cfg(not(any(feature = "mimalloc", feature = "jemallocator")))]
let malloc_kind = "standard";
let (qsvtype, maj, min, pat, pre, rustversion) = (
option_env!("CARGO_BIN_NAME"),
option_env!("CARGO_PKG_VERSION_MAJOR"),
option_env!("CARGO_PKG_VERSION_MINOR"),
option_env!("CARGO_PKG_VERSION_PATCH"),
option_env!("CARGO_PKG_VERSION_PRE"),
option_env!("CARGO_PKG_RUST_VERSION"),
);
if let (Some(qsvtype), Some(maj), Some(min), Some(pat), Some(pre), Some(rustversion)) =
(qsvtype, maj, min, pat, pre, rustversion)
{
if pre.is_empty() {
format!(
"{qsvtype} {maj}.{min}.{pat}-{malloc_kind}-{enabled_features}{maxjobs}-{numcpus};\
{max_file_size}-{free_swap}-{avail_mem}-{total_mem} ({TARGET} compiled with Rust \
{rustversion}) {QSV_KIND}",
maxjobs = max_jobs(),
numcpus = num_cpus(),
max_file_size = indicatif::HumanBytes(max_file_size),
free_swap = indicatif::HumanBytes(free_swap),
avail_mem = indicatif::HumanBytes(avail_mem),
total_mem = indicatif::HumanBytes(total_mem),
)
} else {
format!(
"{qsvtype} {maj}.{min}.\
{pat}-{pre}-{malloc_kind}-{enabled_features}{maxjobs}-{numcpus};\
{max_file_size}-{free_swap}-{avail_mem}-{total_mem} ({TARGET} compiled with Rust \
{rustversion}) {QSV_KIND}",
maxjobs = max_jobs(),
numcpus = num_cpus(),
max_file_size = indicatif::HumanBytes(max_file_size),
free_swap = indicatif::HumanBytes(free_swap),
avail_mem = indicatif::HumanBytes(avail_mem),
total_mem = indicatif::HumanBytes(total_mem),
)
}
} else {
String::new()
}
}
const OTHER_ENV_VARS: &[&str] = &["all_proxy", "no_proxy", "http_proxy", "https_proxy"];
pub fn show_env_vars() {
let mut env_var_set = false;
for (n, v) in env::vars_os() {
let env_var = n.into_string().unwrap();
#[cfg(feature = "mimalloc")]
if env_var.starts_with("QSV_")
|| env_var.starts_with("MIMALLOC_")
|| OTHER_ENV_VARS.contains(&env_var.to_ascii_lowercase().as_str())
{
env_var_set = true;
woutinfo!("{env_var}: {v:?}");
}
#[cfg(feature = "jemallocator")]
if env_var.starts_with("QSV_")
|| env_var.starts_with("JEMALLOC_")
|| env_var.starts_with("MALLOC_CONF")
|| OTHER_ENV_VARS.contains(&env_var.to_ascii_lowercase().as_str())
{
env_var_set = true;
woutinfo!("{env_var}: {v:?}");
}
#[cfg(not(any(feature = "mimalloc", feature = "jemallocator")))]
if env_var.starts_with("QSV_")
|| OTHER_ENV_VARS.contains(&env_var.to_ascii_lowercase().as_str())
{
env_var_set = true;
woutinfo!("{env_var}: {v:?}");
}
}
if !env_var_set {
woutinfo!("No qsv-relevant environment variables set.");
}
}
#[inline]
pub fn count_rows(conf: &Config) -> Result<u64, CliError> {
if let Some(idx) = conf.indexed().unwrap_or(None) {
Ok(idx.count())
} else {
// index does not exist or is stale,
// count records by iterating through records
// Do this only once per invocation and cache the result in ROW_COUNT,
// so we don't have to re-count rows every time we need to know the
// rowcount for CSVs that don't have an index.
let count_opt = ROW_COUNT.get_or_init(|| {
if let Ok(mut rdr) = conf.reader() {
let mut count = 0_u64;
let mut _record = csv::ByteRecord::new();
#[allow(clippy::used_underscore_binding)]
while rdr.read_byte_record(&mut _record).unwrap_or_default() {
count += 1;
}
Some(count)
} else {
None
}
});
match *count_opt {
Some(count) => Ok(count),
None => Err(CliError::Other("Unable to get row count".to_string())),
}
}
}
#[cfg(any(feature = "feature_capable", feature = "lite"))]
pub fn prep_progress(progress: &ProgressBar, record_count: u64) {
progress.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}] [{wide_bar} {percent}%{msg}] ({per_sec} - {eta})")
.unwrap(),
);
progress.set_message(format!(" of {} records", HumanCount(record_count)));
// draw progress bar for the first time using specified style
progress.set_length(record_count);
log::info!("Progress started... {record_count} records");
}
#[cfg(any(feature = "feature_capable", feature = "lite"))]
pub fn finish_progress(progress: &ProgressBar) {
progress.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}] [{wide_bar} {percent}%{msg}] ({per_sec})")
.unwrap(),
);
if progress.length().unwrap_or_default() == progress.position() {
progress.finish();
log::info!("Progress done... {}", progress.message());
} else {
progress.abandon();
log::info!("Progress abandoned... {}", progress.message());
}
}
#[cfg(all(any(feature = "fetch", feature = "geocode"), not(feature = "lite")))]
macro_rules! update_cache_info {
($progress:expr, $cache_instance:expr) => {
use cached::Cached;
use indicatif::HumanCount;
match $cache_instance.lock() {
Ok(cache) => {
let size = cache.cache_size();
if size > 0 {
let hits = cache.cache_hits().unwrap_or_default();
let misses = cache.cache_misses().unwrap_or(1);
let hit_ratio = (hits as f64 / (hits + misses) as f64) * 100.0;
let capacity = cache.cache_capacity();
$progress.set_message(format!(
" of {} records. Cache {:.2}% entries: {} capacity: {}.",
HumanCount($progress.length().unwrap()),
hit_ratio,
HumanCount(size as u64),
HumanCount(capacity.unwrap() as u64),
));
}
},
_ => {},
}
};
($progress:expr, $cache_hits:expr, $num_rows:expr) => {
use indicatif::HumanCount;
let hit_ratio = ($cache_hits as f64 / $num_rows as f64) * 100.0;
$progress.set_message(format!(
" of {} records. Cache hit ratio: {hit_ratio:.2}%",
HumanCount($progress.length().unwrap()),
));
};
}
#[cfg(all(any(feature = "fetch", feature = "geocode"), not(feature = "lite")))]
pub(crate) use update_cache_info;
pub fn get_args<T>(usage: &str, argv: &[&str]) -> CliResult<T>
where
T: DeserializeOwned,
{
Docopt::new(usage)
.and_then(|d| {
d.argv(argv.iter().copied())
.version(Some(version()))
.deserialize()
})
.map_err(From::from)
}
#[inline]
pub fn many_configs(
inps: &[PathBuf],
delim: Option<Delimiter>,
no_headers: bool,
flexible: bool,
) -> Result<Vec<Config>, String> {
let mut inps = inps
.iter()
.map(|p| p.to_str().unwrap_or("-").to_owned())
.collect::<Vec<_>>();
if inps.is_empty() {
inps.push("-".to_owned()); // stdin
}
let confs = inps
.into_iter()
.map(|p| {
Config::new(&Some(p))
.delimiter(delim)
.no_headers(no_headers)
.flexible(flexible)
})
.collect::<Vec<_>>();
errif_greater_one_stdin(&confs)?;
Ok(confs)
}
pub fn errif_greater_one_stdin(inps: &[Config]) -> Result<(), String> {
let nstd = inps.iter().filter(|inp| inp.is_stdin()).count();
if nstd > 1 {
return fail!("At most one <stdin> input is allowed.");
}
Ok(())
}
pub const fn chunk_size(nitems: usize, njobs: usize) -> usize {
if nitems < njobs {
nitems
} else {
nitems / njobs
}
}
pub const fn num_of_chunks(nitems: usize, chunk_size: usize) -> usize {
if chunk_size == 0 {
return nitems;
}
let mut n = nitems / chunk_size;
if nitems % chunk_size != 0 {
n += 1;
}
n
}
pub fn file_metadata(md: &fs::Metadata) -> (u64, u64) {
use filetime::FileTime;
let last_modified = FileTime::from_last_modification_time(md).unix_seconds() as u64;
let fsize = md.len();
(last_modified, fsize)
}
/// Check if there is enough memory to process the file.
/// Return the maximum file size that can be processed.
/// If the file is larger than the maximum file size, return an error.
/// If memcheck is true, check memory in CONSERVATIVE mode (i.e., Filesize < AVAIL memory + SWAP -
/// headroom) If memcheck is false, check memory in NORMAL mode (i.e., Filesize < TOTAL memory -
/// headroom)
pub fn mem_file_check(
path: &Path,
version_check: bool,
conservative_memcheck: bool,
) -> CliResult<i64> {
// if we're NOT calling this from the version() and the file doesn't exist,
// we don't need to check memory as file existence is checked before this function is called.
// If we do get here with a non-existent file, that means we're using stdin,
// so this check doesn't apply, so we return -1
if !path.exists() && !version_check {
return Ok(-1_i64);
}
let conservative_memcheck_work = get_envvar_flag("QSV_MEMORY_CHECK") || conservative_memcheck;
let mut sys = sysinfo::System::new();
sys.refresh_memory();
let avail_mem = sys.available_memory();
let free_swap = sys.free_swap();
let total_mem = sys.total_memory();
let mut mem_pct = env::var("QSV_FREEMEMORY_HEADROOM_PCT")
.unwrap_or_else(|_| DEFAULT_FREEMEMORY_HEADROOM_PCT.to_string())
.parse::<u8>()
.unwrap_or(DEFAULT_FREEMEMORY_HEADROOM_PCT);
// for safety, we don't want to go below 10% memory headroom
// nor above 90% memory headroom as its too memory-restrictive
mem_pct = mem_pct.clamp(10, 90);
#[allow(clippy::cast_precision_loss)]
let max_avail_mem = if conservative_memcheck_work {
((avail_mem + free_swap) as f32 * ((100 - mem_pct) as f32 / 100.0_f32)) as u64
} else {
(total_mem as f32 * ((100 - mem_pct) as f32 / 100.0_f32)) as u64
};
// if we're calling this from version(), we don't need to check the file size
if !version_check {
let file_metadata =
fs::metadata(path).map_err(|e| format!("Failed to get file size: {e}"))?;
let fsize = file_metadata.len();
if fsize > max_avail_mem {
return fail_OOM_clierror!(
"Not enough memory to process the file. qsv running in non-streaming {mode} mode. \
Total memory: {total_mem} Available memory: {avail_mem}. Free swap: {free_swap} \
Max Available memory/Max input file size: {max_avail_mem}. \
QSV_FREEMEMORY_HEADROOM_PCT: {mem_pct}%. File size: {fsize}.",
mode = if conservative_memcheck_work {
"CONSERVATIVE"
} else {
"NORMAL"
},
total_mem = indicatif::HumanBytes(total_mem),
avail_mem = indicatif::HumanBytes(avail_mem),
free_swap = indicatif::HumanBytes(free_swap),
max_avail_mem = indicatif::HumanBytes(max_avail_mem),
mem_pct = mem_pct,
fsize = indicatif::HumanBytes(fsize)
);
}
}
Ok(max_avail_mem as i64)
}
#[cfg(any(feature = "feature_capable", feature = "lite"))]
#[inline]
pub fn condense(val: Cow<[u8]>, n: Option<usize>) -> Cow<[u8]> {
match n {
None => val,
Some(n) => {
let mut is_short_utf8 = false;
if let Ok(s) = simdutf8::basic::from_utf8(&val) {
if n >= s.chars().count() {
is_short_utf8 = true;
} else {
let mut s = s.chars().take(n).collect::<String>();
s.push_str("...");
return Cow::Owned(s.into_bytes());
}
}
if is_short_utf8 || n >= (*val).len() {
// already short enough
val
} else {
// This is a non-Unicode string, so we just trim on bytes.
let mut s = val[0..n].to_vec();
s.extend(b"...".iter().copied());
Cow::Owned(s)
}
},
}
}
pub fn idx_path(csv_path: &Path) -> PathBuf {
let mut p = csv_path
.to_path_buf()
.into_os_string()
.into_string()
.unwrap();
p.push_str(".idx");
PathBuf::from(&p)
}
pub type Idx = Option<usize>;
pub fn range(start: Idx, end: Idx, len: Idx, index: Idx) -> Result<(usize, usize), String> {
match (start, end, len, index) {
(None, None, None, Some(i)) => Ok((i, i + 1)),
(_, _, _, Some(_)) => fail!("--index cannot be used with --start, --end or --len"),
(_, Some(_), Some(_), None) => {
fail!("--end and --len cannot be used at the same time.")
},
(_, None, None, None) => Ok((start.unwrap_or(0), ::std::usize::MAX)),
(_, Some(e), None, None) => {
let s = start.unwrap_or(0);
if s > e {
fail_format!(
"The end of the range ({e}) must be greater than or\nequal to the start of \
the range ({s})."
)
} else {
Ok((s, e))
}
},
(_, None, Some(l), None) => {
let s = start.unwrap_or(0);
Ok((s, s + l))
},
}
}
/// Create a directory recursively, avoiding the race conditions fixed by
/// https://github.com/rust-lang/rust/pull/39799.
#[cfg(any(feature = "feature_capable", feature = "lite"))]
fn create_dir_all_threadsafe(path: &Path) -> std::io::Result<()> {
use std::thread;
// Try 20 times. This shouldn't theoretically need to be any larger
// than the number of nested directories we need to create.
for _ in 0..20 {
match fs::create_dir_all(path) {
// This happens if a directory in `path` doesn't exist when we
// test for it, and another thread creates it before we can.
Err(ref err) if err.kind() == std::io::ErrorKind::AlreadyExists => {},
other => return other,
}
// We probably don't need to sleep at all, because the intermediate
// directory is already created. But let's attempt to back off a
// bit and let the other thread finish.
thread::sleep(std::time::Duration::from_millis(25));
}
// Try one last time, returning whatever happens.
fs::create_dir_all(path)
}
/// Represents a filename template of the form `"{}.csv"`, where `"{}"` is
/// the place to insert the part of the filename generated by `qsv`.
#[cfg(any(feature = "feature_capable", feature = "lite"))]
#[derive(Clone)]
pub struct FilenameTemplate {
prefix: String,
suffix: String,
}
#[cfg(any(feature = "feature_capable", feature = "lite"))]
impl FilenameTemplate {
/// Generate a new filename using `unique_value` to replace the `"{}"`
/// in the template.
pub fn filename(&self, unique_value: &str) -> String {
format!("{}{unique_value}{}", &self.prefix, &self.suffix)
}
/// Create a new, writable file in directory `path` with a filename
/// using `unique_value` to replace the `"{}"` in the template. Note
/// that we do not output headers; the caller must do that if
/// desired.
pub fn writer<P>(
&self,
path: P,
unique_value: &str,
) -> std::io::Result<csv::Writer<Box<dyn std::io::Write + 'static>>>
where
P: AsRef<Path>,
{
let filename = self.filename(unique_value);
let full_path = path.as_ref().join(filename);
if let Some(parent) = full_path.parent() {
// We may be called concurrently, especially by parallel `qsv
// split`, so be careful to avoid the `create_dir_all` race
// condition.
create_dir_all_threadsafe(parent)?;
}
let spath = Some(full_path.display().to_string());
Config::new(&spath).writer()
}
}
#[cfg(any(feature = "feature_capable", feature = "lite"))]
impl<'de> Deserialize<'de> for FilenameTemplate {
fn deserialize<D: Deserializer<'de>>(d: D) -> Result<FilenameTemplate, D::Error> {
let raw = String::deserialize(d)?;
let chunks = raw.split("{}").collect::<Vec<_>>();
if chunks.len() == 2 {
Ok(FilenameTemplate {
prefix: chunks[0].to_owned(),
suffix: chunks[1].to_owned(),
})
} else {
Err(D::Error::custom(
"The --filename argument must contain one '{}'.",
))
}
}
}
pub fn init_logger() -> (String, flexi_logger::LoggerHandle) {
use flexi_logger::{Cleanup, Criterion, FileSpec, Logger, Naming};
let qsv_log_env = env::var("QSV_LOG_LEVEL").unwrap_or_else(|_| "off".to_string());
let qsv_log_dir = env::var("QSV_LOG_DIR").unwrap_or_else(|_| ".".to_string());
let write_mode = if get_envvar_flag("QSV_LOG_UNBUFFERED") {
flexi_logger::WriteMode::Direct
} else {
flexi_logger::WriteMode::BufferAndFlush
};
let logger = Logger::try_with_env_or_str(qsv_log_env)
.unwrap()
.use_utc()
.log_to_file(
FileSpec::default()
.directory(qsv_log_dir)
.suppress_timestamp(),
)
.write_mode(write_mode)
.format_for_files(flexi_logger::detailed_format)
.o_append(true)
.rotate(
Criterion::Size(20_000_000), // 20 mb
Naming::Numbers,
Cleanup::KeepLogAndCompressedFiles(10, 100),
)
.start()
.unwrap();
let qsv_args: String = if log::log_enabled!(log::Level::Info) {
env::args().skip(1).collect::<Vec<_>>().join(" ")
} else {
String::new()
};
log::info!("START: {qsv_args}");
(qsv_args, logger)
}
#[cfg(feature = "self_update")]
pub fn qsv_check_for_update(check_only: bool, no_confirm: bool) -> Result<bool, String> {
use self_update::cargo_crate_version;
const GITHUB_RATELIMIT_MSG: &str =
"Github is rate-limiting self-update checks at the moment. Try again in an hour.";
if get_envvar_flag("QSV_NO_UPDATE") {
return Ok(false);
}
let bin_name = match std::env::current_exe() {
Ok(pb) => {
if let Some(fs) = pb.file_stem() {
fs.to_string_lossy().into_owned()
} else {
return fail!("Can't get the exec stem name");
}
},
Err(e) => return fail_format!("Can't get the exec path: {e}"),
};
winfo!("Checking GitHub for updates...");
let curr_version = cargo_crate_version!();
let releases = if let Ok(releases_list) =
self_update::backends::github::ReleaseList::configure()
.repo_owner("jqnatividad")
.repo_name("qsv")
.build()
{
if let Ok(releases) = releases_list.fetch() {
releases
} else {
return fail!(GITHUB_RATELIMIT_MSG);
}
} else {
return fail!(GITHUB_RATELIMIT_MSG);
};
let latest_release = &releases[0].version;
log::info!("Current version: {curr_version} Latest Release: {latest_release}");
let mut updated = false;
let Ok(latest_release_sv) = semver::Version::parse(latest_release) else {
return fail_format!("Can't parse latest release version: {latest_release}");
};
let Ok(curr_version_sv) = semver::Version::parse(curr_version) else {
return fail_format!("Can't parse current version: {curr_version}");
};
if latest_release_sv > curr_version_sv {
eprintln!("Update {latest_release} available. Current version is {curr_version}.");
eprintln!(
"Release notes: https://github.com/jqnatividad/qsv/releases/tag/{latest_release}\n"
);
if QSV_KIND.starts_with("prebuilt") && !check_only {
match self_update::backends::github::Update::configure()
.repo_owner("jqnatividad")
.repo_name("qsv")
.bin_name(&bin_name)
.show_download_progress(true)
.show_output(false)
.no_confirm(no_confirm)
.current_version(curr_version)
.verifying_keys([*include_bytes!("qsv-zipsign-public.key")])
.build()
{
Ok(update_job) => match update_job.update() {
Ok(status) => {
updated = true;
let update_status = format!(
"Update successful for {}: `{}`!",
bin_name,
status.version()
);
winfo!("{update_status}");
},
Err(e) => werr!("Update job error: {e}"),
},
Err(e) => werr!("Update builder error: {e}"),
};
} else if check_only {
winfo!("Use the --update option to upgrade {bin_name} to the latest release.");
} else {
// we don't want to overwrite manually curated/configured qsv installations.
// If QSV_KIND is not "prebuilt", just inform the user of the new release, and let them
// rebuild their qsvs the way they like it, instead of overwriting it with
// our prebuilt binaries.
winfo!(
r#"This qsv was {QSV_KIND}. self-update does not work for manually {QSV_KIND} binaries.
If you wish to update to the latest version of qsv, manually install/compile from source.
Self-update only works with prebuilt binaries released on GitHub https://github.com/jqnatividad/qsv/releases/latest"#
);
}
} else {
winfo!("Up to date ({curr_version})... no update required.");
};
if !check_only {
if let Ok(status_code) =
send_hwsurvey(&bin_name, updated, latest_release, curr_version, false)
{
log::info!("HW survey sent. Status code: {status_code}");
}
}
Ok(updated)
}
#[cfg(not(feature = "self_update"))]
pub const fn qsv_check_for_update(_check_only: bool, _no_confirm: bool) -> Result<bool, String> {
Ok(true)
}
// the qsv hwsurvey allows us to keep a better
// track of qsv's usage in the wild, so we can do a
// better job of prioritizing platforms/features we support
// no personally identifiable information is collected
#[cfg(feature = "self_update")]
fn send_hwsurvey(
bin_name: &str,
updated: bool,
latest_release: &str,
curr_version: &str,
dry_run: bool,
) -> Result<reqwest::StatusCode, String> {
use serde_json::json;
static HW_SURVEY_URL: &str =
"https://4dhmneehnl.execute-api.us-east-1.amazonaws.com/dev/qsv-hwsurvey";
let mut sys = System::new();
sys.refresh_all();
let total_mem = sys.total_memory();
let kernel_version =
sysinfo::System::kernel_version().unwrap_or_else(|| "Unknown kernel".to_string());
let long_os_version =
sysinfo::System::long_os_version().unwrap_or_else(|| "Unknown OS version".to_string());
let cpu_count = sys.cpus().len();
let physical_cpu_count = sys.physical_core_count().unwrap_or_default();
let cpu_vendor_id = sys.cpus()[0].vendor_id();
let cpu_brand = sys.cpus()[0].brand().trim();
let cpu_freq = sys.cpus()[0].frequency();
let long_id: u128 = std::time::SystemTime::now()
.duration_since(std::time::SystemTime::UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
// the id doubles as a timestamp
// we first get number of milliseconds since UNIX EPOCH
// and then cast to u64 as serde_json cannot serialize u128
let id: u64 = long_id.try_into().unwrap_or_default();
let hwsurvey_json = json!(
{
"id": id,
"variant": bin_name,
"kind": QSV_KIND,
"ver": if updated { latest_release } else { curr_version },
"updated": updated,
"prev_ver": curr_version,
"cpu_phy_cores": physical_cpu_count,
"cpu_log_cores": cpu_count,
"cpu_vendor": cpu_vendor_id,
"cpu_brand": cpu_brand,
"cpu_freq": cpu_freq,
"mem": total_mem,
"kernel": kernel_version,
"os": long_os_version,
"target": TARGET,
}
);
log::debug!("hwsurvey: {hwsurvey_json}");
let mut survey_done = false;
let mut status = reqwest::StatusCode::OK;
if dry_run {
log::info!("Survey dry run. hw survey compiled successfully, but not sent.");
} else {
let client = match reqwest::blocking::Client::builder()
.user_agent(default_user_agent())
.brotli(true)
.gzip(true)
.deflate(true)
.use_rustls_tls()
.http2_adaptive_window(true)
.build()
{
Ok(c) => c,
Err(e) => return fail_format!("Cannot build hw_survey reqwest client: {e}"),
};
match client
.post(HW_SURVEY_URL)
.body(hwsurvey_json.to_string())
.header(reqwest::header::CONTENT_TYPE, "application/json")
.header(reqwest::header::HOST, "qsv.rs")
.send()
{
Ok(resp) => {
log::debug!("hw_survey response sent: {:?}", &resp);
status = resp.status();
survey_done = status.is_success();
},
Err(e) => {
log::warn!("Cannot send hw survey: {e}");
status = reqwest::StatusCode::BAD_REQUEST;
},
};
}
if survey_done || dry_run {
Ok(status)
} else {
fail!("hw survey failed.")
}
}
#[allow(clippy::needless_pass_by_value)]
pub fn safe_header_names(
headers: &csv::StringRecord,
check_first_char: bool,
conditional: bool,
reserved_names: Option<Vec<String>>,
unsafe_prefix: &str,
keep_case: bool,
) -> (Vec<String>, u16) {
// Create "safe" var/key names - to support dynfmt/url-template, valid python vars & db-safe
// column names. Fold to lowercase if keep_case is false. Trim leading & trailing whitespace.
// Replace whitespace/non-alphanumeric) with _. If name starts with a number & check_first_char
// is true, prepend the unsafe_prefix. If a column with the same name already exists,
// append a sequence suffix (e.g. _n). Names are limited to 60 characters in length.
// Empty names are replaced with unsafe_prefix as well.
// If conditional = true & reserved_names is none, only rename the header if its not safe
let prefix = if unsafe_prefix.is_empty() {
"_"
} else {
unsafe_prefix
};
let safename_regex = regex_oncelock!(r"[^A-Za-z0-9]");
let mut changed_count = 0_u16;
let mut name_vec: Vec<String> = Vec::with_capacity(headers.len());
let mut safe_name: String;
let mut safename_always: String;
let mut safename_candidate: String;
let mut final_candidate: String;
let mut buf_wrk = String::new();
for header_name in headers {
let reserved_found = if let Some(reserved_names_vec) = reserved_names.clone() {
if keep_case {
header_name.clone_into(&mut buf_wrk);
} else {
to_lowercase_into(header_name, &mut buf_wrk);
};
reserved_names_vec
.iter()
.any(|reserved_name| reserved_name == &buf_wrk)
} else {
false
};
safe_name = if conditional && is_safe_name(header_name) && !reserved_found {
header_name.to_string()
} else {
safename_always = if header_name.is_empty() {
prefix.to_string()
} else {
safename_regex
.replace_all(header_name.trim(), "_")
.to_string()
};
if check_first_char && safename_always.as_bytes()[0].is_ascii_digit() {
safename_always = format!("{prefix}{safename_always}");
}