Avro and raw_string serialization (#541)

crates/arroyo-api/src/connection_tables.rs

@@ -18,8 +18,8 @@ use tracing::warn;
 use arroyo_connectors::confluent::ConfluentProfile;
 use arroyo_connectors::connector_for_type;
 use arroyo_connectors::kafka::{KafkaConfig, KafkaTable, SchemaRegistry};
-use arroyo_formats::avro;
-use arroyo_formats::json_schema::to_arrow;
+use arroyo_formats::avro::schema;
+use arroyo_formats::avro::schema::to_arrow;
 use arroyo_operator::connector::ErasedConnector;
 use arroyo_rpc::api_types::connections::{
     ConnectionProfile, ConnectionSchema, ConnectionTable, ConnectionTablePost, ConnectionType,
@@ -526,7 +526,7 @@ async fn expand_avro_schema(
         );
     }
 
-    let fields: Result<_, String> = avro::to_arrow(&name, &definition)
+    let fields: Result<_, String> = schema::to_arrow(&name, &definition)
         .map_err(|e| bad_request(format!("Invalid avro schema: {}", e)))?
         .fields
         .into_iter()

crates/arroyo-api/src/jobs.rs

@@ -12,15 +12,13 @@ use arroyo_rpc::api_types::{
 };
 use arroyo_rpc::grpc;
 use arroyo_rpc::grpc::api::{
-    CreateJobReq, JobStatus, OperatorCheckpointDetail, TaskCheckpointDetail,
-    TaskCheckpointEventType,
+    CreateJobReq, OperatorCheckpointDetail, TaskCheckpointDetail, TaskCheckpointEventType,
 };
 use arroyo_rpc::grpc::controller_grpc_client::ControllerGrpcClient;
 use arroyo_rpc::public_ids::{generate_id, IdTypes};
 use axum::extract::{Path, Query, State};
 use axum::response::sse::{Event, Sse};
 use axum::Json;
-use cornucopia_async::GenericClient;
 use cornucopia_async::Params;
 use deadpool_postgres::Transaction;
 use futures_util::stream::Stream;
@@ -63,10 +61,20 @@ pub(crate) async fn create_job<'a>(
         ));
     }
 
-    let running_jobs = get_job_statuses(&auth, client)
-        .await?
+    let running_jobs = api_queries::get_jobs()
+        .bind(client, &auth.organization_id)
+        .all()
+        .await
+        .map_err(log_and_map)?
         .iter()
-        .filter(|j| j.running_desired && j.state != "Failed" && j.state != "Finished")
+        .filter(|j| {
+            j.stop == public::StopMode::none
+                && !j
+                    .state
+                    .as_ref()
+                    .map(|s| s == "Failed" || s == "Finished")
+                    .unwrap_or(false)
+        })
         .count();
 
     if running_jobs > auth.org_metadata.max_running_jobs as usize {
@@ -111,36 +119,6 @@ pub(crate) async fn create_job<'a>(
     Ok(job_id)
 }
 
-pub(crate) async fn get_job_statuses(
-    auth: &AuthData,
-    client: &impl GenericClient,
-) -> Result<Vec<JobStatus>, ErrorResp> {
-    let res = api_queries::get_jobs()
-        .bind(client, &auth.organization_id)
-        .all()
-        .await
-        .map_err(log_and_map)?;
-
-    res.into_iter()
-        .map(|rec| {
-            Ok(JobStatus {
-                job_id: rec.id,
-                pipeline_name: rec.pipeline_name,
-                running_desired: rec.stop == public::StopMode::none,
-                state: rec.state.unwrap_or_else(|| "Created".to_string()),
-                run_id: rec.run_id.unwrap_or(0) as u64,
-                start_time: rec.start_time.map(to_micros),
-                finish_time: rec.finish_time.map(to_micros),
-                tasks: rec.tasks.map(|t| t as u64),
-                definition: rec.textual_repr,
-                udfs: serde_json::from_value(rec.udfs).map_err(log_and_map)?,
-                pipeline_id: format!("{}", rec.pipeline_id),
-                failure_message: rec.failure_message,
-            })
-        })
-        .collect()
-}
-
 pub(crate) fn get_action(state: &str, running_desired: &bool) -> (String, Option<StopType>, bool) {
     enum Progress {
         InProgress,
@@ -483,7 +461,7 @@ pub async fn get_job_output(
         .graph
         .nodes
         .iter()
-        .any(|n| n.operator.contains("PreviewSink"))
+        .any(|n| n.operator.contains("preview"))
     {
         // TODO: make this check more robust
         return Err(bad_request("Job does not have a preview sink".to_string()));

crates/arroyo-api/src/pipelines.rs

@@ -8,30 +8,30 @@ use cornucopia_async::{GenericClient, Params};
 use deadpool_postgres::{Object, Transaction};
 use http::StatusCode;
 
-use petgraph::Direction;
+use petgraph::{Direction, EdgeDirection};
 use std::collections::HashMap;
 use std::env;
 
+use petgraph::visit::NodeRef;
 use std::time::Duration;
 
 use crate::{compiler_service, connection_profiles, jobs, pipelines, types};
-use arroyo_datastream::ConnectorOp;
+use arroyo_datastream::preview_sink;
 use arroyo_rpc::api_types::pipelines::{
     Job, Pipeline, PipelinePatch, PipelinePost, PipelineRestart, QueryValidationResult, StopType,
     ValidateQueryPost,
 };
 use arroyo_rpc::api_types::udfs::{GlobalUdf, Udf};
 use arroyo_rpc::api_types::{JobCollection, PaginationQueryParams, PipelineCollection};
+use arroyo_rpc::grpc::api as api_proto;
 use arroyo_rpc::grpc::api::{
-    create_pipeline_req, ArrowProgram, CreateJobReq, CreatePipelineReq, CreateSqlJob,
+    create_pipeline_req, ArrowProgram, ConnectorOp, CreateJobReq, CreatePipelineReq, CreateSqlJob,
 };
-use arroyo_rpc::grpc::{api as api_proto, api};
 
 use arroyo_connectors::kafka::{KafkaConfig, KafkaTable, SchemaRegistry};
 use arroyo_datastream::logical::{LogicalProgram, OperatorName};
 use arroyo_df::{has_duplicate_udf_names, ArroyoSchemaProvider, CompiledSql, ParsedUdf, SqlConfig};
-use arroyo_formats::avro::arrow_to_avro_schema;
-use arroyo_formats::json::arrow_to_json_schema;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_rpc::formats::Format;
 use arroyo_rpc::grpc::compiler_grpc_client::CompilerGrpcClient;
 use arroyo_rpc::public_ids::{generate_id, IdTypes};
@@ -212,7 +212,7 @@ async fn try_register_confluent_schema(
     match config.format.clone() {
         Some(Format::Avro(mut avro)) => {
             if avro.confluent_schema_registry && avro.schema_id.is_none() {
-                let avro_schema = arrow_to_avro_schema("avro_root", schema.fields());
+                let avro_schema = ArrowSerializer::avro_schema(&*schema);
 
                 let id = schema_registry
                     .write_schema(avro_schema.canonical_form(), ConfluentSchemaType::Avro)
@@ -224,7 +224,7 @@ async fn try_register_confluent_schema(
         }
         Some(Format::Json(mut json)) => {
             if json.confluent_schema_registry && json.schema_id.is_none() {
-                let json_schema = arrow_to_json_schema(schema.fields());
+                let json_schema = ArrowSerializer::json_schema(&*schema);
 
                 let id = schema_registry
                     .write_schema(json_schema.to_string(), ConfluentSchemaType::Json)
@@ -245,26 +245,35 @@ async fn try_register_confluent_schema(
 }
 
 async fn register_schemas(compiled_sql: &mut CompiledSql) -> anyhow::Result<()> {
-    for node in compiled_sql.program.graph.node_indices() {
-        let Some(_input) = compiled_sql
+    // register schemas for sinks
+    for idx in compiled_sql
+        .program
+        .graph
+        .externals(Direction::Outgoing)
+        .collect::<Vec<_>>()
+    {
+        let edge = compiled_sql
             .program
             .graph
-            .edges_directed(node, Direction::Incoming)
+            .edges_directed(idx, EdgeDirection::Incoming)
             .next()
-        else {
-            continue;
-        };
+            .ok_or_else(|| anyhow!("no incoming edges for sink node: {:?}", idx.weight()))?;
+
+        let schema = edge.weight().schema.schema.clone();
 
-        // TODO: schema registration
-        // let Some(value_schema) = compiled_sql.schemas.get(&input.weight().value) else {
-        //     continue;
-        // };
-        //
-        // let node = compiled_sql.program.graph.node_weight_mut(node).unwrap();
-        //
-        // if let Operator::ConnectorSink(connector) = &mut node.operator {
-        //     try_register_confluent_schema(connector, value_schema).await?;
-        // }
+        let node = compiled_sql.program.graph.node_weight_mut(idx).unwrap();
+        if node.operator_name == OperatorName::ConnectorSink {
+            let mut op = ConnectorOp::decode(&node.operator_config[..]).map_err(|_| {
+                anyhow!(
+                    "failed to decode configuration for connector node {:?}",
+                    node
+                )
+            })?;
+
+            try_register_confluent_schema(&mut op, &schema).await?;
+
+            node.operator_config = op.encode_to_vec();
+        }
     }
 
     Ok(())
@@ -296,7 +305,6 @@ pub(crate) async fn create_pipeline<'a>(
                     .try_into()
                     .map_err(log_and_map)?,
                 connection_ids: vec![],
-                schemas: HashMap::new(),
             };
             text = None;
             udfs = None;
@@ -356,8 +364,7 @@ pub(crate) async fn create_pipeline<'a>(
         for node in compiled.program.graph.node_weights_mut() {
             // replace all sink connectors with websink for preview
             if node.operator_name == OperatorName::ConnectorSink {
-                node.operator_config =
-                    api::ConnectorOp::from(ConnectorOp::web_sink()).encode_to_vec();
+                node.operator_config = preview_sink().encode_to_vec();
             }
         }
     }
@@ -455,7 +462,7 @@ impl TryInto<Pipeline> for DbPipeline {
             checkpoint_interval_micros: self.checkpoint_interval_micros as u64,
             stop,
             created_at: to_micros(self.created_at),
-            graph: program.as_job_graph().into(),
+            graph: program.try_into().map_err(log_and_map)?,
             action: action.map(|a| a.into()),
             action_text,
             action_in_progress,
@@ -514,7 +521,7 @@ pub async fn validate_query(
             //optimizations::optimize(&mut program.graph);
 
             QueryValidationResult {
-                graph: Some(program.as_job_graph().into()),
+                graph: Some(program.try_into().map_err(log_and_map)?),
                 errors: None,
             }
         }

crates/arroyo-connectors/src/filesystem/sink/json.rs

@@ -1,7 +1,7 @@
 use std::{fs::File, io::Write, time::Instant};
 
 use arrow::record_batch::RecordBatch;
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_rpc::{df::ArroyoSchemaRef, formats::Format};
 
 use super::{

crates/arroyo-connectors/src/fluvio/mod.rs

@@ -1,5 +1,5 @@
 use anyhow::{anyhow, bail};
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::{Connection, Connector};
 use arroyo_operator::operator::OperatorNode;
 use arroyo_rpc::api_types::connections::{ConnectionProfile, ConnectionSchema, TestSourceMessage};

crates/arroyo-connectors/src/fluvio/sink.rs

@@ -2,7 +2,7 @@ use arrow::array::RecordBatch;
 use async_trait::async_trait;
 use fluvio::{Fluvio, FluvioConfig, TopicProducer};
 
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use tracing::info;
 
 use arroyo_operator::context::ArrowContext;

crates/arroyo-connectors/src/kafka/mod.rs

@@ -1,6 +1,6 @@
 use anyhow::{anyhow, bail};
-use arroyo_formats::serialize::ArrowSerializer;
-use arroyo_formats::ArrowDeserializer;
+use arroyo_formats::de::ArrowDeserializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::Connection;
 use arroyo_rpc::api_types::connections::{ConnectionProfile, ConnectionSchema, TestSourceMessage};
 use arroyo_rpc::df::ArroyoSchema;

crates/arroyo-connectors/src/kafka/sink/mod.rs

@@ -13,7 +13,7 @@ use rdkafka::util::Timeout;
 use rdkafka::ClientConfig;
 
 use arrow::array::RecordBatch;
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::context::ArrowContext;
 use arroyo_operator::operator::ArrowOperator;
 use arroyo_types::CheckpointBarrier;

crates/arroyo-connectors/src/kafka/sink/test.rs

@@ -7,7 +7,7 @@ use std::time::{Duration, SystemTime};
 use arrow::array::{RecordBatch, UInt32Array};
 use arrow::datatypes::Field;
 use arrow::datatypes::{DataType, Schema, SchemaRef};
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::context::ArrowContext;
 use arroyo_operator::operator::ArrowOperator;
 use arroyo_rpc::df::ArroyoSchema;

crates/arroyo-connectors/src/kinesis/mod.rs

@@ -2,7 +2,7 @@ use anyhow::{anyhow, bail, Result};
 use std::collections::HashMap;
 use typify::import_types;
 
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::Connection;
 use arroyo_rpc::api_types::connections::{ConnectionProfile, TestSourceMessage};
 use arroyo_rpc::{api_types, OperatorConfig};

crates/arroyo-connectors/src/kinesis/sink.rs

@@ -2,7 +2,7 @@ use std::time::{Duration, SystemTime};
 
 use anyhow::Result;
 use arrow::array::RecordBatch;
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::context::ArrowContext;
 use arroyo_operator::operator::ArrowOperator;
 use arroyo_types::CheckpointBarrier;

crates/arroyo-connectors/src/preview/operator.rs

@@ -55,7 +55,6 @@ impl ArrowOperator for PreviewSink {
                             .map(|nanos| from_nanos(nanos as u128))
                             .unwrap_or_else(|| SystemTime::now()),
                     ),
-                    key: value.clone(),
                     value,
                     done: false,
                 })
@@ -72,7 +71,6 @@ impl ArrowOperator for PreviewSink {
                 operator_id: ctx.task_info.operator_id.clone(),
                 subtask_index: ctx.task_info.task_index as u32,
                 timestamp: to_micros(SystemTime::now()),
-                key: "".to_string(),
                 value: "".to_string(),
                 done: true,
             })

crates/arroyo-connectors/src/redis/mod.rs

@@ -3,7 +3,7 @@ mod operator;
 use std::collections::HashMap;
 
 use anyhow::{anyhow, bail};
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::{Connection, Connector};
 use arroyo_operator::operator::OperatorNode;
 use arroyo_rpc::var_str::VarStr;

crates/arroyo-connectors/src/redis/operator/sink.rs

@@ -1,6 +1,6 @@
 use crate::redis::{ListOperation, RedisClient, RedisTable, TableType, Target};
 use arrow::array::{AsArray, RecordBatch};
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::context::{ArrowContext, ErrorReporter};
 use arroyo_operator::operator::ArrowOperator;
 use arroyo_types::CheckpointBarrier;

crates/arroyo-connectors/src/webhook/mod.rs

@@ -7,7 +7,7 @@ use std::time::SystemTime;
 use anyhow::anyhow;
 use arroyo_rpc::OperatorConfig;
 
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::Connection;
 use arroyo_rpc::api_types::connections::{
     ConnectionProfile, ConnectionSchema, ConnectionType, TestSourceMessage,

crates/arroyo-connectors/src/webhook/operator.rs

@@ -12,7 +12,7 @@ use tokio::sync::{Mutex, Semaphore};
 use tracing::warn;
 
 use crate::webhook::MAX_INFLIGHT;
-use arroyo_formats::serialize::ArrowSerializer;
+use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::context::ArrowContext;
 use arroyo_operator::operator::ArrowOperator;
 use arroyo_rpc::grpc::TableConfig;