i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add async request body cleanup records

Browse Source
This commit is contained in:
Entropy.Xu
2026-05-08 13:50:18 +08:00
parent fa22384f24
commit 3f29335fd6
15 changed files with 1151 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
crates/aether-data/src/backend/maintenance.rs
View File

@@ -194,6 +194,16 @@ impl DataBackends {
None => Ok(AdminSystemPurgeSummary::default()),
}
}
pub async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
match self.sql_backend() {
Some(backend) => backend.purge_admin_request_bodies_batch(batch_size).await,
None => Ok(AdminSystemPurgeSummary::default()),
}
}
}
impl PostgresBackend {
@@ -494,4 +504,15 @@ impl<'a> SqlBackendRef<'a> {
Self::Sqlite(sqlite) => sqlite.purge_admin_system_data(target).await,
}
}
async fn purge_admin_request_bodies_batch(
self,
batch_size: usize,
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
match self {
Self::Postgres(postgres) => postgres.purge_admin_request_bodies_batch(batch_size).await,
Self::Mysql(mysql) => mysql.purge_admin_request_bodies_batch(batch_size).await,
Self::Sqlite(sqlite) => sqlite.purge_admin_request_bodies_batch(batch_size).await,
}
}
}

400
crates/aether-data/src/backend/system.rs
View File

@@ -90,6 +90,20 @@ impl PostgresBackend {
Ok(summary)
}
pub async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
if batch_size == 0 {
return Ok(AdminSystemPurgeSummary::default());
}
let mut tx = self.pool().begin().await.map_postgres_err()?;
let mut summary = AdminSystemPurgeSummary::default();
purge_postgres_request_bodies_batch(&mut tx, batch_size, &mut summary).await?;
tx.commit().await.map_postgres_err()?;
Ok(summary)
}
pub async fn find_system_config_value(
&self,
key: &str,
@@ -195,6 +209,20 @@ impl MysqlBackend {
Ok(summary)
}
pub async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
if batch_size == 0 {
return Ok(AdminSystemPurgeSummary::default());
}
let mut tx = self.pool().begin().await.map_sql_err()?;
let mut summary = AdminSystemPurgeSummary::default();
purge_mysql_request_bodies_batch(&mut tx, batch_size, &mut summary).await?;
tx.commit().await.map_sql_err()?;
Ok(summary)
}
pub async fn find_system_config_value(
&self,
key: &str,
@@ -335,6 +363,20 @@ impl SqliteBackend {
Ok(summary)
}
pub async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
if batch_size == 0 {
return Ok(AdminSystemPurgeSummary::default());
}
let mut tx = self.pool().begin().await.map_sql_err()?;
let mut summary = AdminSystemPurgeSummary::default();
purge_sqlite_request_bodies_batch(&mut tx, batch_size, &mut summary).await?;
tx.commit().await.map_sql_err()?;
Ok(summary)
}
pub async fn find_system_config_value(
&self,
key: &str,
@@ -1506,6 +1548,301 @@ WHERE user_id IN ({non_admin_users})
Ok(())
}
async fn purge_postgres_request_bodies_batch(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
batch_size: usize,
summary: &mut AdminSystemPurgeSummary,
) -> Result<(), DataLayerError> {
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
pg_execute_batch_if_table(
tx,
"usage_body_blobs",
"usage_body_blobs",
r#"
WITH doomed AS (
SELECT body_ref
FROM public.usage_body_blobs
ORDER BY body_ref ASC
LIMIT $1
)
DELETE FROM public.usage_body_blobs AS blobs
USING doomed
WHERE blobs.body_ref = doomed.body_ref
"#,
summary,
limit,
)
.await?;
pg_execute_batch_if_table(
tx,
"usage",
"usage_body_fields_cleaned",
r#"
WITH batch AS (
SELECT request_id
FROM public.usage
WHERE request_body IS NOT NULL
OR response_body IS NOT NULL
OR provider_request_body IS NOT NULL
OR client_response_body IS NOT NULL
OR request_body_compressed IS NOT NULL
OR response_body_compressed IS NOT NULL
OR provider_request_body_compressed IS NOT NULL
OR client_response_body_compressed IS NOT NULL
ORDER BY created_at_unix_ms ASC, request_id ASC
LIMIT $1
)
UPDATE public.usage AS usage_rows
SET request_body = NULL,
response_body = NULL,
provider_request_body = NULL,
client_response_body = NULL,
request_body_compressed = NULL,
response_body_compressed = NULL,
provider_request_body_compressed = NULL,
client_response_body_compressed = NULL
FROM batch
WHERE usage_rows.request_id = batch.request_id
"#,
summary,
limit,
)
.await?;
pg_execute_batch_if_table(
tx,
"usage_http_audits",
"usage_http_audit_body_refs_cleaned",
r#"
WITH batch AS (
SELECT request_id
FROM public.usage_http_audits
WHERE request_body_ref IS NOT NULL
OR provider_request_body_ref IS NOT NULL
OR response_body_ref IS NOT NULL
OR client_response_body_ref IS NOT NULL
OR request_body_state IS NOT NULL
OR provider_request_body_state IS NOT NULL
OR response_body_state IS NOT NULL
OR client_response_body_state IS NOT NULL
OR body_capture_mode <> 'none'
ORDER BY request_id ASC
LIMIT $1
)
UPDATE public.usage_http_audits AS audits
SET request_body_ref = NULL,
provider_request_body_ref = NULL,
response_body_ref = NULL,
client_response_body_ref = NULL,
request_body_state = NULL,
provider_request_body_state = NULL,
response_body_state = NULL,
client_response_body_state = NULL,
body_capture_mode = 'none',
updated_at = NOW()
FROM batch
WHERE audits.request_id = batch.request_id
"#,
summary,
limit,
)
.await?;
Ok(())
}
async fn purge_mysql_request_bodies_batch(
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
batch_size: usize,
summary: &mut AdminSystemPurgeSummary,
) -> Result<(), DataLayerError> {
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
mysql_execute_batch_if_table(
tx,
"usage_body_blobs",
"usage_body_blobs",
r#"
DELETE FROM usage_body_blobs
WHERE body_ref IN (
SELECT body_ref FROM (
SELECT body_ref
FROM usage_body_blobs
ORDER BY body_ref ASC
LIMIT ?
) AS doomed
)
"#,
summary,
limit,
)
.await?;
mysql_execute_batch_if_table(
tx,
"usage",
"usage_body_fields_cleaned",
r#"
UPDATE `usage`
SET request_body = NULL,
response_body = NULL,
provider_request_body = NULL,
client_response_body = NULL,
request_body_compressed = NULL,
response_body_compressed = NULL,
provider_request_body_compressed = NULL,
client_response_body_compressed = NULL
WHERE request_id IN (
SELECT request_id FROM (
SELECT request_id
FROM `usage`
WHERE request_body IS NOT NULL
OR response_body IS NOT NULL
OR provider_request_body IS NOT NULL
OR client_response_body IS NOT NULL
OR request_body_compressed IS NOT NULL
OR response_body_compressed IS NOT NULL
OR provider_request_body_compressed IS NOT NULL
OR client_response_body_compressed IS NOT NULL
ORDER BY created_at_unix_ms ASC, request_id ASC
LIMIT ?
) AS batch
)
"#,
summary,
limit,
)
.await?;
mysql_execute_batch_if_table(
tx,
"usage_http_audits",
"usage_http_audit_body_refs_cleaned",
r#"
UPDATE usage_http_audits
SET request_body_ref = NULL,
provider_request_body_ref = NULL,
response_body_ref = NULL,
client_response_body_ref = NULL,
request_body_state = NULL,
provider_request_body_state = NULL,
response_body_state = NULL,
client_response_body_state = NULL,
body_capture_mode = 'none'
WHERE request_id IN (
SELECT request_id FROM (
SELECT request_id
FROM usage_http_audits
WHERE request_body_ref IS NOT NULL
OR provider_request_body_ref IS NOT NULL
OR response_body_ref IS NOT NULL
OR client_response_body_ref IS NOT NULL
OR request_body_state IS NOT NULL
OR provider_request_body_state IS NOT NULL
OR response_body_state IS NOT NULL
OR client_response_body_state IS NOT NULL
OR body_capture_mode <> 'none'
ORDER BY request_id ASC
LIMIT ?
) AS batch
)
"#,
summary,
limit,
)
.await?;
Ok(())
}
async fn purge_sqlite_request_bodies_batch(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
batch_size: usize,
summary: &mut AdminSystemPurgeSummary,
) -> Result<(), DataLayerError> {
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
sqlite_execute_batch_if_table(
tx,
"usage_body_blobs",
"usage_body_blobs",
r#"
DELETE FROM usage_body_blobs
WHERE body_ref IN (
SELECT body_ref
FROM usage_body_blobs
ORDER BY body_ref ASC
LIMIT ?
)
"#,
summary,
limit,
)
.await?;
sqlite_execute_batch_if_table(
tx,
"usage",
"usage_body_fields_cleaned",
r#"
UPDATE "usage"
SET request_body = NULL,
response_body = NULL,
provider_request_body = NULL,
client_response_body = NULL,
request_body_compressed = NULL,
response_body_compressed = NULL,
provider_request_body_compressed = NULL,
client_response_body_compressed = NULL
WHERE request_id IN (
SELECT request_id
FROM "usage"
WHERE request_body IS NOT NULL
OR response_body IS NOT NULL
OR provider_request_body IS NOT NULL
OR client_response_body IS NOT NULL
OR request_body_compressed IS NOT NULL
OR response_body_compressed IS NOT NULL
OR provider_request_body_compressed IS NOT NULL
OR client_response_body_compressed IS NOT NULL
ORDER BY created_at_unix_ms ASC, request_id ASC
LIMIT ?
)
"#,
summary,
limit,
)
.await?;
sqlite_execute_batch_if_table(
tx,
"usage_http_audits",
"usage_http_audit_body_refs_cleaned",
r#"
UPDATE usage_http_audits
SET request_body_ref = NULL,
provider_request_body_ref = NULL,
response_body_ref = NULL,
client_response_body_ref = NULL,
request_body_state = NULL,
provider_request_body_state = NULL,
response_body_state = NULL,
client_response_body_state = NULL,
body_capture_mode = 'none'
WHERE request_id IN (
SELECT request_id
FROM usage_http_audits
WHERE request_body_ref IS NOT NULL
OR provider_request_body_ref IS NOT NULL
OR response_body_ref IS NOT NULL
OR client_response_body_ref IS NOT NULL
OR request_body_state IS NOT NULL
OR provider_request_body_state IS NOT NULL
OR response_body_state IS NOT NULL
OR client_response_body_state IS NOT NULL
OR body_capture_mode <> 'none'
ORDER BY request_id ASC
LIMIT ?
)
"#,
summary,
limit,
)
.await?;
Ok(())
}
async fn pg_delete_table(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
table: &str,
@@ -1586,6 +1923,27 @@ async fn pg_execute_if_table(
Ok(())
}
async fn pg_execute_batch_if_table(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
table: &str,
key: &str,
sql: &str,
summary: &mut AdminSystemPurgeSummary,
limit: i64,
) -> Result<(), DataLayerError> {
if !pg_table_exists(tx, checked_sql_identifier(table)?).await? {
return Ok(());
}
let rows = sqlx::query(sql)
.bind(limit)
.execute(&mut **tx)
.await
.map_postgres_err()?
.rows_affected();
summary.add(key, rows);
Ok(())
}
async fn pg_table_exists(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
table: &str,
@@ -1678,6 +2036,27 @@ async fn mysql_execute_if_table(
Ok(())
}
async fn mysql_execute_batch_if_table(
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
table: &str,
key: &str,
sql: &str,
summary: &mut AdminSystemPurgeSummary,
limit: i64,
) -> Result<(), DataLayerError> {
if !mysql_table_exists(tx, checked_sql_identifier(table)?).await? {
return Ok(());
}
let rows = sqlx::query(sql)
.bind(limit)
.execute(&mut **tx)
.await
.map_sql_err()?
.rows_affected();
summary.add(key, rows);
Ok(())
}
async fn mysql_table_exists(
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
table: &str,
@@ -1773,6 +2152,27 @@ async fn sqlite_execute_if_table(
Ok(())
}
async fn sqlite_execute_batch_if_table(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
table: &str,
key: &str,
sql: &str,
summary: &mut AdminSystemPurgeSummary,
limit: i64,
) -> Result<(), DataLayerError> {
if !sqlite_table_exists(tx, checked_sql_identifier(table)?).await? {
return Ok(());
}
let rows = sqlx::query(sql)
.bind(limit)
.execute(&mut **tx)
.await
.map_sql_err()?
.rows_affected();
summary.add(key, rows);
Ok(())
}
async fn sqlite_table_exists(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
table: &str,

6
crates/aether-data/src/repository/system.rs
View File

@@ -41,6 +41,12 @@ impl AdminSystemPurgeSummary {
*self.affected.entry(key.into()).or_insert(0) += count;
}
pub fn merge(&mut self, other: &Self) {
for (key, count) in &other.affected {
self.add(key.clone(), *count);
}
}
pub fn total(&self) -> u64 {
self.affected.values().copied().sum()
}