package org.example;
​
import java.sql.*;
import java.util.Properties;
​
public class Main {
 public static void main(String[] args) {
   final Properties properties = new Properties();
   properties.put("user", "root");
   properties.put("password", "");
   properties.put("tenant", "cnosdb");
   properties.put("useEncryption", false);
   try (
     Connection connection = DriverManager.getConnection(
       "jdbc:arrow-flight-sql://localhost:31004", properties
    );
     Statement statement = connection.createStatement())
  {
     statement.execute("CREATE TABLE IF NOT EXISTS air\n" +
                       "(\n" +
                       "   visibility DOUBLE,\n" +
                       "   temperature DOUBLE,\n" +
                       "   pressure   DOUBLE,\n" +
                       "   TAGS(station)\n" +
                       ");");
     statement.executeUpdate("INSERT INTO air (TIME, station, visibility, temperature, pressure) VALUES\n" +
                             "   (1666165200290401000, 'XiaoMaiDao', 56, 69, 77);");
     ResultSet resultSet = statement.executeQuery("select * from air limit 1;");
​
     while (resultSet.next()) {
       Timestamp column1 = resultSet.getTimestamp(1);
       String column2 = resultSet.getString(2);
       Double column3 = resultSet.getDouble(3);
       Double column4 = resultSet.getDouble(4);
       Double column5 = resultSet.getDouble(5);
​
       System.out.printf("%s %s %f %f %f", column1, column2, column3, column4, column5);
    }
  } catch (SQLException e) {
     throw new RuntimeException(e);
  }
}
}

arrow = {version = "28.0.0", features = ["prettyprint"] }
arrow-flight = {version = "28.0.0", features = ["flight-sql-experimental"]}
tokio = "1.23.0"
futures = "0.3.25"
prost-types = "0.11.2"
tonic = "0.8.3"
prost = "0.11.3"
http-auth-basic = "0.3.3"
base64 = "0.13.1"

let mut client = FlightServiceClient::connect("http://localhost:31004")
.await
.expect("connect faile");

use std::collections::HashMap;
use std::iter;
use std::sync::Arc;

use arrow::buffer::Buffer;
use arrow::datatypes::Schema;
use arrow::ipc;
use arrow::ipc::{reader, root_as_message};
use arrow_flight::flight_service_client::FlightServiceClient;
use arrow_flight::sql::{CommandStatementQuery, ProstAnyExt};
use arrow_flight::utils::flight_data_to_arrow_batch;
use arrow_flight::{FlightDescriptor, HandshakeRequest, IpcMessage};
use futures::StreamExt;

use prost::Message;
use tonic::codegen::http::header::AUTHORIZATION;
use tonic::metadata::AsciiMetadataValue;
use tonic::Request;

#[tokio::main]
async fn main() {

    let mut client = FlightServiceClient::connect("http://localhost:31004")
        .await
        .expect("connect");

    let mut req = Request::new(futures::stream::iter(iter::once(
        HandshakeRequest::default(),
    )));

    req.metadata_mut().insert(
        AUTHORIZATION.as_str(),
        AsciiMetadataValue::try_from(format!(
            "Basic {}",
            base64::encode(format!("{}:{}", "root", ""))
        ))
        .expect("metadata construct fail"),
    );

    let resp = client.handshake(req).await.expect("handshake");

    println!("handshake resp: {:?}", resp.metadata());

    let cmd = CommandStatementQuery {
        query: "select now();".to_string(),
    };
    let pack = prost_types::Any::pack(&cmd).expect("pack");
    let fd = FlightDescriptor::new_cmd(pack.encode_to_vec());

    let mut req = Request::new(fd);
    req.metadata_mut().insert(
        AUTHORIZATION.as_str(),
        resp.metadata().get(AUTHORIZATION.as_str()).unwrap().clone(),
    );
    let resp = client.get_flight_info(req).await.expect("get_flight_info");

    let flight_info = resp.into_inner();
    let schema_ref =
        Arc::new(Schema::try_from(IpcMessage(flight_info.schema)).expect("Schema::try_from"));
    println!("{}", schema_ref);

    for ep in flight_info.endpoint {
        if let Some(ticket) = ep.ticket {
            let resp = client.do_get(ticket).await.expect("do_get");
            let mut stream = resp.into_inner();
            let mut dictionaries_by_id = HashMap::new();

            let mut record_batches = Vec::new();
            while let Some(Ok(flight_data)) = stream.next().await {
                let message =
                    root_as_message(&flight_data.data_header[..]).expect("root as message");
                match message.header_type() {
                    ipc::MessageHeader::Schema => {
                        println!("a schema when messages are read",);
                    }

                    ipc::MessageHeader::RecordBatch => {
                        let record_batch = flight_data_to_arrow_batch(
                            &flight_data,
                            schema_ref.clone(),
                            &dictionaries_by_id,
                        )
                        .expect("record_batch_from_message");
                        record_batches.push(record_batch);
                    }
                    ipc::MessageHeader::DictionaryBatch => {
                        let ipc_batch = message.header_as_dictionary_batch().unwrap();

                        reader::read_dictionary(
                            &Buffer::from(flight_data.data_body),
                            ipc_batch,
                            &schema_ref,
                            &mut dictionaries_by_id,
                            &message.version(),
                        )
                        .unwrap();
                    }
                    _ => {
                        panic!("Reading types other than record batches not yet supported");
                    }
                }
            }

            println!(
                "{}",
                arrow::util::pretty::pretty_format_batches(&record_batches).expect("print")
            );
        }
    }
}