From a3c53aaabcc2e6b04ddeb545a42f7be5cef69a35 Mon Sep 17 00:00:00 2001
From: Ahmed Riza <dr.riza@gmail.com>
Date: Fri, 20 Oct 2023 21:45:38 +0100
Subject: [PATCH] Update base64 lib and fix compilation failure in
 flight_sql.rs

---
 ballista/core/src/serde/scheduler/to_proto.rs | 10 ++++++++--
 ballista/scheduler/Cargo.toml                 |  2 +-
 ballista/scheduler/src/flight_sql.rs          |  4 +++-
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/ballista/core/src/serde/scheduler/to_proto.rs b/ballista/core/src/serde/scheduler/to_proto.rs
index 6ceb1dd6e..9c92e6019 100644
--- a/ballista/core/src/serde/scheduler/to_proto.rs
+++ b/ballista/core/src/serde/scheduler/to_proto.rs
@@ -158,12 +158,18 @@ impl TryInto<protobuf::OperatorMetric> for &MetricValue {
             }),
             MetricValue::StartTimestamp(timestamp) => Ok(protobuf::OperatorMetric {
                 metric: Some(operator_metric::Metric::StartTimestamp(
-                    timestamp.value().map(|m| m.timestamp_nanos()).unwrap_or(0),
+                    timestamp
+                        .value()
+                        .map(|m| m.timestamp_nanos_opt().unwrap())
+                        .unwrap_or(0),
                 )),
             }),
             MetricValue::EndTimestamp(timestamp) => Ok(protobuf::OperatorMetric {
                 metric: Some(operator_metric::Metric::EndTimestamp(
-                    timestamp.value().map(|m| m.timestamp_nanos()).unwrap_or(0),
+                    timestamp
+                        .value()
+                        .map(|m| m.timestamp_nanos_opt().unwrap())
+                        .unwrap_or(0),
                 )),
             }),
         }
diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml
index c7ee94bc2..9883ef09d 100644
--- a/ballista/scheduler/Cargo.toml
+++ b/ballista/scheduler/Cargo.toml
@@ -47,7 +47,7 @@ arrow-flight = { workspace = true }
 async-recursion = "1.0.0"
 async-trait = "0.1.41"
 ballista-core = { path = "../core", version = "0.11.0", features = ["s3"] }
-base64 = { version = "0.13", default-features = false }
+base64 = { version = "0.21", default-features = false }
 clap = { version = "3", features = ["derive", "cargo"] }
 configure_me = { workspace = true }
 dashmap = "5.4.0"
diff --git a/ballista/scheduler/src/flight_sql.rs b/ballista/scheduler/src/flight_sql.rs
index 942930a17..2c3d02ce5 100644
--- a/ballista/scheduler/src/flight_sql.rs
+++ b/ballista/scheduler/src/flight_sql.rs
@@ -35,6 +35,7 @@ use arrow_flight::{
     Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
     HandshakeResponse, Location, Ticket,
 };
+use base64::{engine::general_purpose, Engine as _};
 use log::{debug, error, warn};
 use std::convert::TryFrom;
 use std::pin::Pin;
@@ -504,7 +505,8 @@ impl FlightSqlService for FlightSqlServiceImpl {
             )))?;
         }
         let base64 = &authorization[basic.len()..];
-        let bytes = base64::decode(base64)
+        let bytes = general_purpose::STANDARD
+            .decode(base64)
             .map_err(|_| Status::invalid_argument("authorization not parsable"))?;
         let str = String::from_utf8(bytes)
             .map_err(|_| Status::invalid_argument("authorization not parsable"))?;