Repository: libsql/libsql
File: libsql-server/src/rpc/proxy.rs (Rust)

Code before:
        let req = req.into_inner();
        let pgm = crate::connection::program::Program::try_from(req.pgm.unwrap())
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        let client_id = Uuid::from_str(&req.client_id).unwrap();

        let connection_maker = self
            .namespaces
            .with(ctx.namespace().clone(), |ns| ns.db.connection_maker())
            .await
            .map_err(|e| {
                if let crate::error::Error::NamespaceDoesntExist(_) = e {
                    tonic::Status::failed_precondition(NAMESPACE_DOESNT_EXIST)
                } else {
                    tonic::Status::internal(e.to_string())
                }
            })?;

        let conn = {
            let lock = self.clients.upgradable_read().await;
            match lock.get(&client_id) {
                Some(conn) => conn.clone(),
                None => {
                    tracing::debug!("connected: {client_id}");
                    match connection_maker.create().await {
                        Ok(conn) => {
                            let conn = Arc::new(TimeoutConnection::new(conn));
                            let mut lock = RwLockUpgradableReadGuard::upgrade(lock).await;
                            lock.insert(client_id, conn.clone());
                            conn
                        }
                        Err(e) => {
                            return Err(tonic::Status::new(tonic::Code::Internal, e.to_string()))
                        }
                    }
                }
            }
        };

        tracing::debug!("executing request for {client_id}");

        let builder = ExecuteResultsBuilder::default();
        let builder = conn
            .execute_program(pgm, ctx, builder, None)
            .await
            // TODO: this is no necessarily a permission denied error!
            .map_err(|e| tonic::Status::new(tonic::Code::PermissionDenied, e.to_string()))?;

        Ok(tonic::Response::new(builder.into_ret()))
    }

    //TODO: also handle cleanup on peer disconnect

Review comment (suggestion):
I removed this since this one actually didn't make much sense, we allow this in the stream_exec. I added an error message for the describe one which should catch schema changes.

Code after:
        let req = req.into_inner();
        let pgm = crate::connection::program::Program::try_from(req.pgm.unwrap())
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        let client_id = Uuid::from_str(&req.client_id).unwrap();

        let connection_maker = self
            .namespaces
            .with(ctx.namespace().clone(), |ns| ns.db.connection_maker())
            .await
            .map_err(|e| {
                if let crate::error::Error::NamespaceDoesntExist(_) = e {
                    tonic::Status::failed_precondition(NAMESPACE_DOESNT_EXIST)
                } else {
                    tonic::Status::internal(e.to_string())
                }
            })?;

        let conn = {
            let lock = self.clients.upgradable_read().await;
            match lock.get(&client_id) {
                Some(conn) => conn.clone(),
                None => {
                    tracing::debug!("connected: {client_id}");
                    match connection_maker.create().await {
                        Ok(conn) => {
                            if !conn.is_primary() {
                                return Err(tonic::Status::failed_precondition(
                                    "cannot run schema migration against a replica from a replica",
                                ));
                            }

                            let conn = Arc::new(TimeoutConnection::new(conn));
                            let mut lock = RwLockUpgradableReadGuard::upgrade(lock).await;
                            lock.insert(client_id, conn.clone());
                            conn
                        }
                        Err(e) => {
                            return Err(tonic::Status::new(tonic::Code::Internal, e.to_string()))
                        }
                    }
                }
            }
        };

        tracing::debug!("executing request for {client_id}");

        let builder = ExecuteResultsBuilder::default();
        let builder = conn
            .execute_program(pgm, ctx, builder, None)
            .await
            // TODO: this is no necessarily a permission denied error!