Skip to content

Commit 16e5c1d

Browse files
docteurkleindocteurklein
committed
conditional listen
1 parent 619977c commit 16e5c1d

5 files changed

Lines changed: 132 additions & 76 deletions

File tree

‎flake.nix‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -292,7 +292,7 @@
292292
};
293293

294294
systemd.services.httpg = {
295-
enable = true;
295+
enable = false;
296296
wantedBy = [ "default.target" ];
297297
serviceConfig = {
298298
Type = "simple";
@@ -473,7 +473,7 @@
473473

474474
system.stateVersion = "25.11";
475475

476-
networking.firewall.allowedTCPPorts = [ 5432 ];
476+
networking.firewall.allowedTCPPorts = [ 5432 6432 ];
477477
networking.useDHCP = false;
478478

479479
users.users.postgres = {
@@ -507,6 +507,21 @@
507507
};
508508
};
509509

510+
services.pgbouncer = {
511+
enable = true;
512+
openFirewall = true;
513+
settings = {
514+
databases = {
515+
httpg = "host=10.250.2.2 user=httpg";
516+
};
517+
pgbouncer = {
518+
auth_type = "any";
519+
listen_addr = "10.250.2.2";
520+
pool_mode = "transaction";
521+
};
522+
};
523+
};
524+
510525
services.postgresql = {
511526
enable = true;
512527
# enableJIT = true;

‎sql/ivm.sql‎

Lines changed: 56 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,59 @@ begin atomic
2727
);
2828
end;
2929

30-
with change (change) as (
31-
select change
32-
from pg_logical_slot_get_changes('test', null, null),
33-
jsonb_array_elements(data::jsonb->'change') change
34-
)
35-
select change->>'schema', change->>'table', change->>'kind', old, new, new - old diff
36-
from change, wal2json_v1_to_record(change), jsonb_populate_record(null::blog.post, new::jsonb) n;
30+
-- drop table if exists blog.stat;
31+
create table if not exists blog.stat (
32+
id int primary key default 1,
33+
nposts int default 0,
34+
lsn pg_lsn default '0/0'
35+
);
36+
insert into blog.stat (id, nposts)
37+
select 1, count(post_id)
38+
from blog.post
39+
on conflict (id) do nothing;
40+
41+
set session characteristics as transaction isolation level serializable;
42+
43+
do $$
44+
declare
45+
lsn_ pg_lsn;
46+
change_ jsonb;
47+
begin
48+
select confirmed_flush_lsn
49+
into lsn_
50+
from pg_replication_slots
51+
where slot_name = 'test';
52+
raise notice '%', lsn_;
53+
loop
54+
with change (lsn, change) as (
55+
select lsn, change
56+
from pg_logical_slot_peek_changes('test', null, null, 'include-types', 'false', 'add-tables', 'blog.post'),
57+
jsonb_array_elements(data::jsonb->'change') change
58+
),
59+
stat as (
60+
update blog.stat
61+
set nposts = nposts + case when change->>'kind' = 'insert' then 1 else -1 end,
62+
lsn = change.lsn
63+
from change
64+
where (change->>'schema', change->>'table') = ('blog', 'post')
65+
and id = 1
66+
and stat.lsn < change.lsn
67+
and change->>'kind' in ('insert', 'delete')
68+
)
69+
select lsn, change
70+
into lsn_, change_
71+
from change;
72+
--, change_--, change, change->>'schema', change->>'table', change->>'kind', old, new, new - old diff into found, lsn_, change_
73+
--, wal2json_v1_to_record(change), jsonb_populate_record(null::blog.post, new::jsonb) n;
74+
75+
raise notice '%', change_;
76+
commit;
77+
perform pg_replication_slot_advance('test', lsn_);
78+
if lsn_ is not null then
79+
raise notice '%', lsn_;
80+
end if;
81+
perform pg_sleep(1);
82+
end loop;
83+
84+
end;
85+
$$;

‎src/extract/query.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -378,7 +378,7 @@ use axum::{body::Body, http::{header::CONTENT_TYPE, Request}};
378378
use axum::extract::FromRequest;
379379
use conf::Conf;
380380
use tokio_postgres::Notification;
381-
use crate::{extract::query::{Param, Query}};
381+
use crate::{Listen, extract::query::{Param, Query}};
382382

383383
#[tokio::test]
384384
async fn test_json_body() {
@@ -400,8 +400,7 @@ use tokio_postgres::Notification;
400400
read_pool,
401401
write_pool,
402402
config: httpg_config.to_owned(),
403-
tx,
404-
client: Arc::new(client),
403+
listen: Some(Listen {tx, client: Arc::new(client) }),
405404
};
406405
let q = Query::from_request(req, &state).await.unwrap();
407406

@@ -430,8 +429,7 @@ use tokio_postgres::Notification;
430429
read_pool,
431430
write_pool,
432431
config: httpg_config.to_owned(),
433-
tx,
434-
client: Arc::new(client),
432+
listen: Some(Listen {tx, client: Arc::new(client) }),
435433
};
436434
let q = Query::from_request(req, &state).await.unwrap();
437435

‎src/main.rs‎

Lines changed: 43 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ struct HttpgConfig {
6363
index_sql: String,
6464
#[conf(long, env)]
6565
login_query: String,
66+
#[conf(long, env)]
67+
listen: bool,
6668
#[conf(long, env, default_value="3000")]
6769
port: u16,
6870
#[conf(flatten)]
@@ -73,13 +75,18 @@ struct HttpgConfig {
7375
pg: PostgresConfig,
7476
}
7577

78+
#[derive(Clone)]
79+
struct Listen {
80+
tx: Sender<Notification>,
81+
client: Arc<Client>,
82+
}
83+
7684
#[derive(Clone)]
7785
struct AppState {
7886
read_pool: Pool,
7987
write_pool: Pool,
8088
config: HttpgConfig,
81-
tx: Sender<Notification>,
82-
client: Arc<Client>,
89+
listen: Option<Listen>,
8390
}
8491

8592
#[tokio::main]
@@ -97,40 +104,42 @@ async fn main() -> Result<(), HttpgError> {
97104
let read_pool = httpg_config.pg.read_pool()?;
98105
let write_pool = httpg_config.pg.write_pool()?;
99106

100-
let (client, mut conn) = httpg_config.pg.connect().await?;
101-
let (tx, _rx) = tokio::sync::broadcast::channel::<Notification>(16);
102-
let mut stream = futures::stream::poll_fn(move |cx| conn.poll_message(cx));
103-
let wrapped_tx = tx.clone();
104-
tokio::spawn(async move {
105-
while let Some(Ok(m)) = stream.next().await {
106-
match m {
107-
tokio_postgres::AsyncMessage::Notice(n) => {
108-
tracing::info!("{n:#?}");
109-
},
110-
tokio_postgres::AsyncMessage::Notification(n) => {
111-
wrapped_tx.send(n).map_err(Box::new)?;
107+
let listen = if httpg_config.listen {
108+
let (client, mut conn) = httpg_config.pg.connect().await?;
109+
let (tx, _rx) = tokio::sync::broadcast::channel::<Notification>(16);
110+
let mut stream = futures::stream::poll_fn(move |cx| conn.poll_message(cx));
111+
let wrapped_tx = tx.clone();
112+
tokio::spawn(async move {
113+
while let Some(Ok(m)) = stream.next().await {
114+
match m {
115+
tokio_postgres::AsyncMessage::Notice(n) => {
116+
tracing::info!("{n:#?}");
117+
},
118+
tokio_postgres::AsyncMessage::Notification(n) => {
119+
wrapped_tx.send(n).map_err(Box::new)?;
120+
}
121+
_ => {}
112122
}
113-
_ => {}
114123
}
115-
}
116-
Ok::<_, HttpgError>(())
117-
});
124+
Ok::<_, HttpgError>(())
125+
});
126+
Some(Listen {tx, client: Arc::new(client)})
127+
} else {
128+
None
129+
};
118130

119131
let state = AppState {
120132
read_pool,
121133
write_pool,
122134
config: httpg_config.to_owned(),
123-
tx,
124-
client: Arc::new(client),
135+
listen,
125136
};
126137

127-
let app = Router::new()
138+
let mut app = Router::new()
128139
.route("/", get(index))
129140
.route("/{path}/", get(index))
130141
.route("/logout", get(logout).post(logout))
131142
.route("/{path}/logout", get(logout).post(logout))
132-
.route("/sse/{channel}", get(sse_query))
133-
.route("/{path}/sse/{channel}", get(sse_query))
134143
.route("/query", get(stream_query).post(post_query))
135144
.route("/{path}/query", get(stream_query).post(post_query))
136145
.route("/{path}/query/{cursor}", post(post_query))
@@ -142,6 +151,14 @@ async fn main() -> Result<(), HttpgError> {
142151
.route("/{path}/webpush", get(web_push).post(web_push))
143152
.route("/login", get(login).post(login))
144153
.route("/{path}/login", get(login).post(login))
154+
;
155+
if httpg_config.listen {
156+
app = app
157+
.route("/listen/{channel}", get(listen_query))
158+
.route("/{path}/listen/{channel}", get(listen_query))
159+
;
160+
}
161+
let app = app
145162
.fallback_service(ServeDir::new(httpg_config.public_dir))
146163
.with_state(state.to_owned())
147164
.layer(ServiceBuilder::new()
@@ -156,42 +173,6 @@ async fn main() -> Result<(), HttpgError> {
156173
)
157174
;
158175

159-
// tokio::spawn(async move {
160-
// let (client, mut conn) = cfg.connect().await?;
161-
162-
// let mut stream = futures::stream::poll_fn(move |cx| conn.poll_message(cx));
163-
164-
// let state = axum::extract::State(state);
165-
166-
// client.simple_query("listen web_push").await?;
167-
// // client.simple_query("listen job").await?;
168-
169-
// while let Some(Ok(m)) = stream.next().await {
170-
// match m {
171-
// tokio_postgres::AsyncMessage::Notice(n) => tracing::info!("{n:#?}"),
172-
// tokio_postgres::AsyncMessage::Notification(n) => {
173-
// match n.channel() {
174-
// "web_push" => {
175-
// let res = web_push(
176-
// state.to_owned(),
177-
// None,
178-
// Query::default()
179-
// )
180-
// .await?
181-
// .into_response()
182-
// ;
183-
// dbg!(&res);
184-
// }
185-
// _ => todo!("{n:#?}")
186-
// }
187-
// },
188-
// _ => todo!(),
189-
// }
190-
// }
191-
192-
// Ok::<(), HttpgError>(())
193-
// });
194-
195176
let addr = SocketAddr::from((
196177
[0, 0, 0, 0],
197178
httpg_config.port,
@@ -538,11 +519,12 @@ async fn stream_query(
538519
}
539520

540521
#[debug_handler]
541-
async fn sse_query(
542-
State(AppState {tx, client, ..}): State<AppState>,
522+
async fn listen_query(
523+
State(AppState {listen, ..}): State<AppState>,
543524
Path(channel): Path<String>,
544525
) -> Result<impl IntoResponse, HttpgError> {
545526

527+
let Listen {tx, client } = listen.ok_or(HttpgError::anyhow("no listen"))?;
546528
client.execute(&format!("listen {channel}"), &[]).await?;
547529

548530
Ok(Sse::new(

‎src/postgres/mod.rs‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,18 @@ use crate::{HttpgError};
1111

1212
#[derive(Clone, Debug, Conf)]
1313
pub struct PostgresConfig {
14+
#[conf(long, env, default_value="")]
15+
read_url: String,
1416
#[conf(long, env)]
1517
read_host: String,
18+
#[conf(long, env, default_value="5432")]
19+
read_port: u16,
20+
#[conf(long, env, default_value="")]
21+
write_url: String,
1622
#[conf(long, env)]
1723
write_host: String,
24+
#[conf(long, env, default_value="5432")]
25+
write_port: u16,
1826
#[conf(long, env)]
1927
user: String,
2028
#[conf(long="password-file", env="PASSWORD_FILE", value_parser = |file: &str| -> Result<_, HttpgError> { Ok(fs::read_to_string(file)?) })]
@@ -33,14 +41,18 @@ impl PostgresConfig {
3341
pub fn read_pool(&self) -> Result<Pool, HttpgError> {
3442
let mut cfg = deadpool_postgres::Config::new();
3543

44+
cfg.url = Some(self.read_url.clone());
3645
cfg.host = Some(self.read_host.clone());
46+
cfg.port = Some(self.read_port);
3747
self.rest(&mut cfg)
3848
}
3949

4050
pub fn write_pool(&self) -> Result<Pool, HttpgError> {
4151
let mut cfg = deadpool_postgres::Config::new();
4252

53+
cfg.url = Some(self.write_url.clone());
4354
cfg.host = Some(self.write_host.clone());
55+
cfg.port = Some(self.write_port);
4456
self.rest(&mut cfg)
4557
}
4658

@@ -50,6 +62,7 @@ impl PostgresConfig {
5062
.password(self.password.clone())
5163
.dbname(self.dbname.clone())
5264
.host(self.write_host.clone())
65+
.port(self.write_port)
5366
.ssl_mode(match self.ssl_mode.as_deref() {
5467
Some("require") => tokio_postgres::config::SslMode::Require,
5568
_ => tokio_postgres::config::SslMode::Prefer,
@@ -59,7 +72,6 @@ impl PostgresConfig {
5972
_ => tokio_postgres::config::ChannelBinding::Prefer,
6073
})
6174
.application_name(self.application_name.to_owned())
62-
// .options("-c statement_timeout=600ms")
6375
.to_owned()
6476
;
6577

0 commit comments

Comments
 (0)