import_ndjson/
import_ndjson.rs

1// Copyright 2024 ScopeDB, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod common;
16
17use std::fs::File;
18use std::io::BufRead;
19use std::io::BufReader;
20use std::io::Read;
21
22#[tokio::main(flavor = "current_thread")]
23async fn main() -> Result<(), Box<dyn std::error::Error>> {
24    let path = std::env::args()
25        .nth(1)
26        .ok_or("usage: import_ndjson <events.ndjson>")?;
27    let mut input = BufReader::new(File::open(path)?);
28    let client = common::client()?;
29    let stream = common::write_table(&client)?.append_stream().build()?;
30    let admitted: Result<(), Box<dyn std::error::Error>> = async {
31        let mut line_number = 0;
32        loop {
33            // Bound a single input record as well as the stream's queued bytes.
34            let mut line = String::new();
35            let bytes = input
36                .by_ref()
37                .take(8 * 1024 * 1024 + 1)
38                .read_line(&mut line)?;
39            if bytes == 0 {
40                break;
41            }
42            line_number += 1;
43            if bytes > 8 * 1024 * 1024 {
44                return Err(format!("line {line_number} exceeds 8 MiB").into());
45            }
46            if line.trim().is_empty() {
47                continue;
48            }
49            let row: serde_json::Value = serde_json::from_str(&line)
50                .map_err(|error| format!("line {line_number}: {error}"))?;
51            stream.send(&row).await?;
52        }
53        Ok(())
54    }
55    .await;
56    // Settle the accepted prefix even if parsing or admission failed midway.
57    let settled = stream.shutdown().await;
58    println!("delivery: {:?}", stream.stats());
59    settled?;
60    admitted?;
61    Ok(())
62}
import_ndjson/import_ndjson.rs

import_ndjson/
import_ndjson.rs
