How to Properly Proxy DatabaseConnection to Intercept SQL Execution (Including Transactions)?

[oulover]

I am learning to explore the implementation of sea orm distributed transactions. First, I will implement an AT pattern similar to Seata. If I can create something useful, I will open source it.
Now we have encountered a problem:
I hope to act as a proxy for Database Connection to intercept SQL (including SQL in transactions) and perform some additional operations,
1.DatabaseConnectionProxy(DatabaseConnection);
2.impl ConnectionTrait for DatabaseConnectionProxy
3.impl TransactionTrait for DatabaseConnectionProxy
But the following calling method:

db.transaction(
        |txn| {
            Box::pin(async move {
                let am = SystemUsersActiveModel{
                    id: ActiveValue::Set(9),
                    username: ActiveValue::Set("test".to_string()),
                };
                SystemUsers::insert(am).exec(txn).await
            })
        }
    ).await?;

What should I do if I don't know how to execute the methods related to impl ConnectTrait for DataconnectionProxy? use the ProxyDatabaseTrait?

use crate::XID;
use sea_orm::sqlx::{Column, Row};
use sea_orm::{error::*, AccessMode, ConnectOptions, ConnectionTrait, Database,
              DatabaseConnection, DatabaseTransaction, DbBackend, ExecResult,
              IsolationLevel, QueryResult, QueryStream, Statement, StreamTrait,
              TransactionError, TransactionTrait};
use std::ops::Deref;
use std::{future::Future, pin::Pin};

use sqlparser::dialect::MySqlDialect;
use sqlparser::parser::Parser;


#[derive(Clone, Debug)]
pub struct DatabaseConnectionProxy(DatabaseConnection);
impl DatabaseConnectionProxy {
    pub fn new(db: DatabaseConnection) -> Self {
        DatabaseConnectionProxy(db)
    }

    pub async fn connect<C>(opt: C) -> Result<DatabaseConnectionProxy, DbErr>
    where
        C: Into<ConnectOptions>,
    {
        Ok(DatabaseConnectionProxy(Database::connect(opt).await?))
    }
}
impl Deref for DatabaseConnectionProxy {
    type Target = DatabaseConnection;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}


impl Default for DatabaseConnectionProxy {
    fn default() -> Self {
        Self(DatabaseConnection::Disconnected)
    }
}

impl DatabaseConnectionProxy {
    async fn query_as_json2(&self, sql: &str) -> Result<serde_json::Value, DbErr> {

        let stmt = Statement::from_string(self.get_database_backend(), sql.to_owned());
        let query_results = self.query_all(stmt).await?;

        let mut rows = Vec::new();
        for result in query_results {
            let row = result.try_as_mysql_row().unwrap();
            for my_clo in row.columns() {
                println!("query_as_json -------{:?}", my_clo);
            }
        }

        Ok(serde_json::Value::Array(rows))
    }

    async fn query_as_json(&self, sql: &str) -> Result<serde_json::Value, DbErr> {
        let stmt = Statement::from_string(self.get_database_backend(), sql.to_owned());
        let query_results = self.query_all(stmt).await?;

        let mut rows = Vec::new();
        for result in query_results {
            let row = result.try_as_mysql_row().ok_or_else(|| {
                DbErr::Custom("Not a MySQL row".into())
            })?;

            let mut map = serde_json::Map::new();

            for col in row.columns().iter() {

                let col_name_str = col.name().to_string();
                let index = col.ordinal();
              
                let value = match row.try_get::<String, _>(index) {
                    Ok(s) => serde_json::Value::String(s),
                    Err(_) => match row.try_get::<i64, _>(index) {
                        Ok(n) => serde_json::Value::Number(n.into()),
                        Err(_) => match row.try_get::<f64, _>(index) {
                            Ok(f) => serde_json::Value::from(f),
                            Err(_) => match row.try_get::<bool, _>(index) {
                                Ok(b) => serde_json::Value::Bool(b),
                                Err(_) => serde_json::Value::Null,
                            },
                        },
                    },
                };

                map.insert(col_name_str, value);
            }

            rows.push(serde_json::Value::Object(map));
        }

        Ok(serde_json::Value::Array(rows))
    }


    async fn process_update(&self, update: &sqlparser::ast::Statement) -> Result<(), DbErr> {
        if let sqlparser::ast::Statement::Update {
            table,
            assignments, from, selection, returning,
            or
        } = update {
            let table_name = table.relation.to_string();
            let where_clause = selection.as_ref().map(|e| e.to_string()).unwrap_or_default();

            for assignment in assignments {
                println!("assignments-------------{:?}", assignment);
            }
            let vec_str: Vec<String> = assignments.iter().map(|it| it.target.to_string()).collect();

            println!("assignments-----CLO--------{:?}", vec_str);

            // 获取before_image
            let select_sql = format!("SELECT * FROM {} WHERE {}", table_name, where_clause);
            println!("before select_sql-------{}", select_sql);
            let before = self.query_as_json(&select_sql).await;
            match &before {
                Ok(data) => {
                    let old = serde_json::to_string(&data).unwrap_or_default();
                    println!("before old-------{}", old);


                    let mut sql = format!("UPDATE {} SET ", table);
                    if let Some(first_row) = data.as_array().and_then(|a| a.first()) {
                        for (key, value) in first_row.as_object().unwrap() {
                            if vec_str.contains(key) {
                                sql.push_str(&format!("{} = {}, ", key, value));
                            }
                        }
                        sql.truncate(sql.len() - 2);
                        sql.push_str(format!(" WHERE {}", where_clause).as_str());
                    }
                    println!("back sql-------{}", sql);
                }
                Err(e) => {
                    eprintln!("{}", e);
                }
            }
            println!("before is-------{:?}", before);
        }
        Ok(())
    }
}


#[async_trait::async_trait]
impl ConnectionTrait for DatabaseConnectionProxy {
    fn get_database_backend(&self) -> DbBackend {
        self.0.get_database_backend()
    }

    async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
        println!("Proxy--execute--{:?}", stmt);
        self.0.execute(stmt).await
    }


    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
        println!("Proxy--sql:--{sql}");
        let xid = XID.try_with(|o| o.clone()).unwrap_or_default();
        println!("Proxy--execute_unprepared23--xid--{:?}", xid);

        let dialect = MySqlDialect {};

        let ast = Parser::parse_sql(&dialect, sql);

        match &ast {
            Ok(statements) => {
                for statement in statements {
                    if let sqlparser::ast::Statement::Update { table: _, assignments: _, from: _, selection: _, returning: _, or: _ } = statement {
                      
                        let r = self.process_update(&statement).await;
                        match r {
                            Ok(_) => {}
                            Err(e) => { eprintln!("process_update error: {}", e) }
                        }
                    } else if let sqlparser::ast::Statement::Insert(i) = statement {
                     
                    } else if let sqlparser::ast::Statement::Delete(i) = statement {
                        
                    }

                    println!("{:#?}", statement);
                }
            }
            Err(e) => eprintln!("Parse error: {}", e),
        }
        let execute_result = self.0.execute_unprepared(sql).await;

        execute_result
    }

    async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
        self.0.query_one(stmt).await
    }

    async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
        self.0.query_all(stmt).await
    }
}

#[async_trait::async_trait]
impl StreamTrait for DatabaseConnectionProxy {
    type Stream<'a> = QueryStream;

    fn stream<'a>(
        &'a self,
        stmt: Statement,
    ) -> Pin<Box<dyn Future<Output=Result<Self::Stream<'a>, DbErr>> + 'a + Send>> {
        self.0.stream(stmt)
    }
}

#[async_trait::async_trait]
impl TransactionTrait for DatabaseConnectionProxy {
    async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
        self.0.begin().await
    }


    async fn begin_with_config(
        &self,
        isolation_level: Option<IsolationLevel>,
        access_mode: Option<AccessMode>,
    ) -> Result<DatabaseTransaction, DbErr> {
        self.0.begin_with_config(isolation_level, access_mode).await
    }

    async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
    where
        F: for<'c> FnOnce(
            &'c DatabaseTransaction,
        ) -> Pin<Box<dyn Future<Output=Result<T, E>> + Send + 'c>>
        + Send,
        T: Send,
        E: std::error::Error + Send,
    {
        // how to handle sql?
        self.0.transaction(callback).await
    }

    async fn transaction_with_config<F, T, E>(
        &self,
        callback: F,
        isolation_level: Option<IsolationLevel>,
        access_mode: Option<AccessMode>,
    ) -> Result<T, TransactionError<E>>
    where
        F: for<'c> FnOnce(
            &'c DatabaseTransaction,
        ) -> Pin<Box<dyn Future<Output=Result<T, E>> + Send + 'c>>
        + Send,
        T: Send,
        E: std::error::Error + Send,
    {
        self.0.transaction_with_config(callback, isolation_level, access_mode).await
    }
}

Use as follows:

 let db = DatabaseConnectionProxy::connect("mysql://root:[email protected]:3306/test")
        .await?; 
db.transaction(
        |txn| {
            Box::pin(async move {
                let am = SystemUsersActiveModel{
                    id: ActiveValue::Set(9),
                    username: ActiveValue::Set("test".to_string()),
                };
                SystemUsers::insert(am).exec(txn).await
            })
        }
    ).await?;

Problem: The transaction has been executed, but it has not entered the Database Connection Proxy. execute

---

ProxyDatabaseTrait：

use sea_orm::{ConnectOptions, Database, DatabaseConnection, DbErr, ProxyDatabaseTrait, ProxyExecResult, ProxyRow, Statement, TransactionTrait};
use std::fmt::Debug;
use std::sync::Arc;

#[derive(Debug)]
pub struct DatabaseTransactionProxyNew {
    inner: DatabaseConnection,
}
#[async_trait::async_trait]
impl ProxyDatabaseTrait for DatabaseTransactionProxyNew {
    async fn query(&self, statement: Statement) -> Result<Vec<ProxyRow>, DbErr> {
        println!("ProxyDatabaseTrait-----query--------A{:?}", statement);
        Ok(vec![])
    }

    async fn execute(&self, statement: Statement) -> Result<ProxyExecResult, DbErr> {
        println!("ProxyDatabaseTrait-----query--------B{:?}", statement);
        Ok(ProxyExecResult { last_insert_id: 0, rows_affected: 0 })
    }


    async fn begin(&self) {
        println!("ProxyDatabaseTrait-----begin--------"); // ok
        let database_transaction = self.inner.begin().await.unwrap();

    }

    async fn commit(&self) {
        println!("ProxyDatabaseTrait-----commit--------"); // no

    }

    async fn rollback(&self) {
        println!("ProxyDatabaseTrait-----rollback--------"); // no

    }

    async fn ping(&self) -> Result<(), DbErr> {
        self.inner.ping().await
    }
}
impl DatabaseTransactionProxyNew {
    pub async fn new(opt: DatabaseConnection) -> Result<DatabaseTransactionProxyNew, DbErr>
    {
        Ok(DatabaseTransactionProxyNew {
            inner: opt,
        })
    }
    pub async fn connect<C>(opt: C) -> Result<DatabaseTransactionProxyNew, DbErr>
    where
        C: Into<ConnectOptions>,
    {
        Ok(DatabaseTransactionProxyNew {
            inner: Database::connect(opt).await?,
        })
    }
}

Use the following:

 let dc =Database::connect("mysql://root:UDMUSJv5iRq4R%@120.46.194.67:3306/rust_seata").await?;
    let proxy =  DatabaseTransactionProxyNew::new(dc) .await?;
    let db_proxy =   Database::connect_proxy(DatabaseBackend::MySql,Arc::new(Box::new(proxy))).await?;

    // begin ok
    db_proxy.transaction(
        |txn| {
            Box::pin(async move {
                let am = SystemUsersActiveModel{
                    id: ActiveValue::Set(9),
                    username: ActiveValue::Set("test".to_string()),
                };
                SystemUsers::insert(am).exec(txn).await
            })
        }
    ).await?;
    // rollback error: DatabaseConnection::Disconnected => Err(conn_err("Disconnected").into()),

[oulover]

#2573