use crate::core::SupabaseClient;
use crate::error::{Result, SupaError};
use futures_util::{SinkExt, StreamExt};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_json::{json, Value};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::sleep;
use tokio_stream::Stream;
use tokio_tungstenite::tungstenite::Message;
#[derive(Clone)]
pub struct RealtimeClient {
pub(crate) client: SupabaseClient,
}
impl RealtimeClient {
pub(crate) fn new(client: SupabaseClient) -> Self {
Self { client }
}
pub fn channel(&self, topic: &str) -> RealtimeChannelBuilder {
RealtimeChannelBuilder::new(self.client.clone(), topic)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PostgresEvent {
Insert,
Update,
Delete,
All,
}
impl ToString for PostgresEvent {
fn to_string(&self) -> String {
match self {
PostgresEvent::Insert => "INSERT".to_string(),
PostgresEvent::Update => "UPDATE".to_string(),
PostgresEvent::Delete => "DELETE".to_string(),
PostgresEvent::All => "*".to_string(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
Connecting,
Connected,
Reconnecting,
Closed,
}
enum ChannelCommand {
Broadcast {
event: String,
payload: Value,
},
Track {
payload: Value,
},
Untrack,
Close,
}
pub struct RealtimeChannel {
topic: String,
rx: mpsc::UnboundedReceiver<Result<RealtimeMessage>>,
cmd_tx: mpsc::UnboundedSender<ChannelCommand>,
}
impl Stream for RealtimeChannel {
type Item = Result<RealtimeMessage>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.rx.poll_recv(cx)
}
}
impl RealtimeChannel {
pub fn topic(&self) -> &str {
&self.topic
}
pub fn broadcast(&self, event: &str, payload: Value) -> Result<()> {
self.cmd_tx
.send(ChannelCommand::Broadcast {
event: event.to_string(),
payload,
})
.map_err(|_| SupaError::RealtimeError {
message: "Channel closed".to_string(),
})
}
pub fn track(&self, payload: Value) -> Result<()> {
self.cmd_tx
.send(ChannelCommand::Track { payload })
.map_err(|_| SupaError::RealtimeError {
message: "Channel closed".to_string(),
})
}
pub fn untrack(&self) -> Result<()> {
self.cmd_tx
.send(ChannelCommand::Untrack)
.map_err(|_| SupaError::RealtimeError {
message: "Channel closed".to_string(),
})
}
pub fn close(&self) -> Result<()> {
self.cmd_tx
.send(ChannelCommand::Close)
.map_err(|_| SupaError::RealtimeError {
message: "Channel already closed".to_string(),
})
}
}
pub struct RealtimeChannelBuilder {
client: SupabaseClient,
topic: String,
postgres_changes: Vec<Value>,
}
impl RealtimeChannelBuilder {
pub fn new(client: SupabaseClient, topic: &str) -> Self {
Self {
client,
topic: topic.to_string(),
postgres_changes: Vec::new(),
}
}
pub fn on_postgres_changes<S1, S2, S3>(
mut self,
event: PostgresEvent,
schema: S1,
table: Option<S2>,
filter: Option<S3>,
) -> Self
where
S1: Into<String>,
S2: Into<String>,
S3: Into<String>,
{
let mut config = json!({
"event": event.to_string(),
"schema": schema.into(),
});
if let Some(t) = table {
config
.as_object_mut()
.unwrap()
.insert("table".to_string(), json!(t.into()));
}
if let Some(f) = filter {
config
.as_object_mut()
.unwrap()
.insert("filter".to_string(), json!(f.into()));
}
self.postgres_changes.push(config);
self
}
pub async fn subscribe(self) -> Result<RealtimeChannel> {
let (tx, rx) = mpsc::unbounded_channel();
let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel();
let client = self.client.clone();
let topic = self.topic.clone();
let mut postgres_changes_config = Vec::new();
for cfg in &self.postgres_changes {
postgres_changes_config.push(json!({
"event": cfg["event"],
"schema": cfg["schema"],
"table": cfg.get("table"),
"filter": cfg.get("filter")
}));
}
let mut config = json!({});
if !postgres_changes_config.is_empty() {
config.as_object_mut().unwrap().insert(
"postgres_changes".to_string(),
json!(postgres_changes_config),
);
}
config.as_object_mut().unwrap().insert(
"broadcast".to_string(),
json!({ "ack": false, "self": false }),
);
config
.as_object_mut()
.unwrap()
.insert("presence".to_string(), json!({ "key": "" }));
let config_clone = config.clone();
tokio::spawn(async move {
let mut retry_count = 0;
let base_delay = client.inner.config.retry_base_delay_ms;
loop {
match connect_and_listen(&client, &topic, &config_clone, &tx, &mut cmd_rx).await {
Ok(_) => {
retry_count = 0;
}
Err(e) => {
let _ = tx.send(Err(SupaError::RealtimeError {
message: format!("Realtime disconnected: {}. Reconnecting...", e),
}));
}
}
retry_count += 1;
let delay = base_delay * 2u64.pow(retry_count.min(9) as u32);
sleep(Duration::from_millis(delay)).await;
}
});
Ok(RealtimeChannel {
topic: self.topic,
rx,
cmd_tx,
})
}
}
async fn connect_and_listen(
client: &SupabaseClient,
topic: &str,
config: &Value,
tx: &mpsc::UnboundedSender<Result<RealtimeMessage>>,
user_cmd_rx: &mut mpsc::UnboundedReceiver<ChannelCommand>,
) -> Result<()> {
let url = client.inner.url.clone();
let scheme = match url.scheme() {
"https" => "wss",
"http" => "ws",
_ => "wss",
};
let host = url.host_str().unwrap_or_default();
let port = url.port_or_known_default().unwrap_or(443);
let ws_url = format!(
"{}://{}:{}/realtime/v1/websocket?apikey={}&vsn=1.0.0",
scheme, host, port, client.inner.key
);
let (ws_stream, _) = tokio_tungstenite::connect_async(&ws_url)
.await
.map_err(|e| SupaError::RealtimeError {
message: format!("Connection failed: {}", e),
})?;
let (mut write, mut read) = ws_stream.split();
let (internal_cmd_tx, mut internal_cmd_rx) = mpsc::channel::<Message>(10);
let writer_handle = tokio::spawn(async move {
while let Some(msg) = internal_cmd_rx.recv().await {
if let Err(_) = write.send(msg).await {
break;
}
}
});
let join_ref = format!("{}", rand::random::<u64>());
let access_token = {
let lock = client.inner.session.read().unwrap();
lock.as_ref()
.map(|s| s.access_token.clone())
.unwrap_or_else(|| client.inner.key.clone())
};
let join_msg = json!({
"topic": topic,
"event": "phx_join",
"payload": {
"config": config,
"access_token": access_token
},
"ref": join_ref
});
internal_cmd_tx
.send(Message::Text(join_msg.to_string()))
.await
.map_err(|e| SupaError::RealtimeError {
message: format!("Failed to send join: {}", e),
})?;
let hb_cmd_tx = internal_cmd_tx.clone();
let hb_handle = tokio::spawn(async move {
loop {
sleep(Duration::from_secs(30)).await;
let msg = json!({
"topic": "phoenix",
"event": "heartbeat",
"payload": {},
"ref": format!("{}", rand::random::<u64>())
});
if hb_cmd_tx
.send(Message::Text(msg.to_string()))
.await
.is_err()
{
break;
}
}
});
loop {
tokio::select! {
msg_res = read.next() => {
match msg_res {
Some(Ok(msg)) => {
match msg {
Message::Text(text) => {
if let Ok(parsed) = serde_json::from_str::<RealtimeMessage>(&text) {
if parsed.event == "phx_reply" {
continue;
}
if parsed.event == "phx_close" {
break;
}
if parsed.event == "phx_error" {
break;
}
if tx.send(Ok(parsed)).is_err() {
break;
}
}
}
Message::Close(_) => break,
_ => {}
}
}
Some(Err(_)) => break, None => break, }
}
cmd = user_cmd_rx.recv() => {
match cmd {
Some(ChannelCommand::Broadcast { event, payload }) => {
let msg = json!({
"topic": topic,
"event": "broadcast",
"payload": {
"event": event,
"payload": payload
},
"ref": format!("{}", rand::random::<u64>())
});
if internal_cmd_tx.send(Message::Text(msg.to_string())).await.is_err() {
break;
}
}
Some(ChannelCommand::Track { payload }) => {
let msg = json!({
"topic": topic,
"event": "presence",
"payload": {
"type": "track",
"event": "track",
"payload": payload
},
"ref": format!("{}", rand::random::<u64>())
});
if internal_cmd_tx.send(Message::Text(msg.to_string())).await.is_err() {
break;
}
}
Some(ChannelCommand::Untrack) => {
let msg = json!({
"topic": topic,
"event": "presence",
"payload": {
"type": "untrack",
"event": "untrack"
},
"ref": format!("{}", rand::random::<u64>())
});
if internal_cmd_tx.send(Message::Text(msg.to_string())).await.is_err() {
break;
}
}
Some(ChannelCommand::Close) => {
let leave_msg = json!({
"topic": topic,
"event": "phx_leave",
"payload": {},
"ref": format!("{}", rand::random::<u64>())
});
let _ = internal_cmd_tx.send(Message::Text(leave_msg.to_string())).await;
return Ok(());
}
None => break }
}
}
}
hb_handle.abort();
writer_handle.abort();
Err(SupaError::RealtimeError {
message: "Connection ended".into(),
})
}
#[derive(Debug, Serialize, Deserialize)]
pub struct RealtimeMessage {
pub topic: String,
pub event: String,
pub payload: Value,
#[serde(rename = "ref")]
pub ref_: Option<String>,
}
impl RealtimeMessage {
pub fn is_postgres_change(&self) -> bool {
self.event == "postgres_changes"
|| self.event == "INSERT"
|| self.event == "UPDATE"
|| self.event == "DELETE"
}
pub fn is_presence(&self) -> bool {
self.event == "presence_state" || self.event == "presence_diff"
}
pub fn is_broadcast(&self) -> bool {
self.event == "broadcast"
}
pub fn as_insert<T: DeserializeOwned>(&self) -> Result<T> {
self.extract_record("INSERT")
}
pub fn as_update<T: DeserializeOwned>(&self) -> Result<T> {
self.extract_record("UPDATE")
}
pub fn as_delete<T: DeserializeOwned>(&self) -> Result<T> {
self.extract_record("DELETE")
}
fn extract_record<T: DeserializeOwned>(&self, expected_type: &str) -> Result<T> {
let type_ = self
.payload
.get("type")
.and_then(|v| v.as_str())
.unwrap_or_default();
if !type_.is_empty() && type_ != expected_type {
return Err(SupaError::RealtimeError {
message: format!("Expected type {}, got {}", expected_type, type_),
});
}
let record_key = if expected_type == "DELETE" {
"old_record"
} else {
"record"
};
let record = self.payload.get(record_key);
match record {
Some(val) if !val.is_null() => {
serde_json::from_value(val.clone()).map_err(|e| SupaError::RealtimeError {
message: format!("Deserialization failed: {}", e),
})
}
_ => {
let fallback = self
.payload
.get("record")
.or_else(|| self.payload.get("old_record"));
if let Some(val) = fallback {
if !val.is_null() {
return serde_json::from_value(val.clone()).map_err(|e| {
SupaError::RealtimeError {
message: format!("Deserialization failed (fallback): {}", e),
}
});
}
}
Err(SupaError::RealtimeError {
message: format!("No {} found in payload", record_key),
})
}
}
}
}