caisin/sea-orm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add Send bound to streams.

Browse Source
This commit is contained in:
Sebastian Pütz 2022-01-19 20:43:40 +01:00
parent 2e038a7eae
commit f012ccaec3
14 changed files with 192 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.toml
View File

@ -36,7 +36,7 @@ serde = { version = "^1.0", features = ["derive"] }
serde_json = { version = "^1", optional = true }
sqlx = { version = "^0.5", optional = true }
uuid = { version = "0.8", features = ["serde", "v4"], optional = true }
ouroboros = "0.14"
ouroboros = { git = "https://github.com/sebpuetz/ouroboros", branch = "send-builder-functions" }
url = "^2.2"
once_cell = "1.8"

3
examples/send_stream_example/.env Normal file
View File

@ -0,0 +1,3 @@
HOST=127.0.0.1
PORT=8000
DATABASE_URL="postgres://postgres:password@localhost/axum_exmaple"

23
examples/send_stream_example/Cargo.toml Normal file
View File

@ -0,0 +1,23 @@
[package]
name = "sea-orm-axum-example"
version = "0.1.0"
authors = ["Sebastian Pütz <seb.puetz@gmail.com>"]
edition = "2021"
publish = false
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[workspace]
[dependencies]
tokio = { version = "1.14", features = ["full"] }
anyhow = "1"
dotenv = "0.15"
futures-util = "0.3"
serde = "1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
[dependencies.sea-orm]
path = "../../" # remove this line in your own project
# version = "^0.5.0"
features = ["macros", "mock", "sqlx-all", "runtime-tokio-rustls", "debug-print"]
default-features = false

1
examples/send_stream_example/README.md Normal file
View File

@ -0,0 +1 @@
Demonstrator for using streaming queries with `tokio::spawn` or in contexts that require `Send` futures.

29
examples/send_stream_example/src/main.rs Normal file
View File

@ -0,0 +1,29 @@
mod post;
mod setup;
use futures_util::StreamExt;
use post::Entity as Post;
use sea_orm::{prelude::*, Database};
use std::env;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
env::set_var("RUST_LOG", "debug");
tracing_subscriber::fmt::init();
dotenv::dotenv().ok();
let db_url = env::var("DATABASE_URL").expect("DATABASE_URL is not set in .env file");
let db = Database::connect(db_url)
.await
.expect("Database connection failed");
let _ = setup::create_post_table(&db);
tokio::task::spawn(async move {
let mut stream = Post::find().stream(&db).await.unwrap();
while let Some(item) = stream.next().await {
let item = item?;
println!("got something: {}", item.text);
}
Ok::<(), anyhow::Error>(())
})
.await?
}

26
examples/send_stream_example/src/post.rs Normal file
View File

@ -0,0 +1,26 @@
//! SeaORM Entity. Generated by sea-orm-codegen 0.3.2
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "posts")]
pub struct Model {
#[sea_orm(primary_key)]
#[serde(skip_deserializing)]
pub id: i32,
pub title: String,
#[sea_orm(column_type = "Text")]
pub text: String,
}
#[derive(Copy, Clone, Debug, EnumIter)]
pub enum Relation {}
impl RelationTrait for Relation {
fn def(&self) -> RelationDef {
panic!("No RelationDef")
}
}
impl ActiveModelBehavior for ActiveModel {}

33
examples/send_stream_example/src/setup.rs Normal file
View File

@ -0,0 +1,33 @@
use sea_orm::sea_query::{ColumnDef, TableCreateStatement};
use sea_orm::{error::*, sea_query, ConnectionTrait, DbConn, ExecResult};
async fn create_table(db: &DbConn, stmt: &TableCreateStatement) -> Result<ExecResult, DbErr> {
let builder = db.get_database_backend();
db.execute(builder.build(stmt)).await
}
pub async fn create_post_table(db: &DbConn) -> Result<ExecResult, DbErr> {
let stmt = sea_query::Table::create()
.table(super::post::Entity)
.if_not_exists()
.col(
ColumnDef::new(super::post::Column::Id)
.integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(
ColumnDef::new(super::post::Column::Title)
.string()
.not_null(),
)
.col(
ColumnDef::new(super::post::Column::Text)
.string()
.not_null(),
)
.to_owned();
create_table(db, &stmt).await
}

4
src/database/connection.rs
View File

@ -9,7 +9,7 @@ use std::{future::Future, pin::Pin};
#[async_trait::async_trait]
pub trait ConnectionTrait<'a>: Sync {
/// Create a stream for the [QueryResult]
type Stream: Stream<Item = Result<QueryResult, DbErr>>;
type Stream: Stream<Item = Result<QueryResult, DbErr>> + Send;
/// Fetch the database backend as specified in [DbBackend].
/// This depends on feature flags enabled.
@ -28,7 +28,7 @@ pub trait ConnectionTrait<'a>: Sync {
fn stream(
&'a self,
stmt: Statement,
) -> Pin<Box<dyn Future<Output = Result<Self::Stream, DbErr>> + 'a>>;
) -> Pin<Box<dyn Future<Output = Result<Self::Stream, DbErr>> + 'a + Send>>;
/// Execute SQL `BEGIN` transaction.
/// Returns a Transaction that can be committed or rolled back

2
src/database/db_connection.rs
View File

@ -155,7 +155,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
fn stream(
&'a self,
stmt: Statement,
) -> Pin<Box<dyn Future<Output = Result<Self::Stream, DbErr>> + 'a>> {
) -> Pin<Box<dyn Future<Output = Result<Self::Stream, DbErr>> + 'a + Send>> {
Box::pin(async move {
Ok(match self {
#[cfg(feature = "sqlx-mysql")]

2
src/database/stream/query.rs
View File

@ -24,7 +24,7 @@ pub struct QueryStream {
metric_callback: Option<crate::metric::Callback>,
#[borrows(mut conn, stmt, metric_callback)]
#[not_covariant]
stream: Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + 'this>>,
stream: Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + Send + 'this>>,
}
#[cfg(feature = "sqlx-mysql")]

103
src/database/stream/transaction.rs
View File

@ -2,7 +2,7 @@
use std::{ops::DerefMut, pin::Pin, task::Poll};
use futures::Stream;
use futures::{FutureExt, Stream};
#[cfg(feature = "sqlx-dep")]
use futures::TryStreamExt;
@ -24,7 +24,7 @@ pub struct TransactionStream<'a> {
metric_callback: Option<crate::metric::Callback>,
#[borrows(mut conn, stmt, metric_callback)]
#[not_covariant]
stream: Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + 'this>>,
stream: Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + 'this + Send>>,
}
impl<'a> std::fmt::Debug for TransactionStream<'a> {
@ -33,61 +33,72 @@ impl<'a> std::fmt::Debug for TransactionStream<'a> {
}
}
type PinStream<'s> = Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + Send + 's>>;
impl<'a> TransactionStream<'a> {
fn stream_builder<'s>(
conn: &'s mut MutexGuard<'a, InnerConnection>,
stmt: &'s Statement,
_metric_callback: &'s Option<crate::metric::Callback>,
) -> Pin<Box<dyn std::future::Future<Output = PinStream<'s>> + 's + Send>> {
async move {
match conn.deref_mut() {
#[cfg(feature = "sqlx-mysql")]
InnerConnection::MySql(c) => {
let query = crate::driver::sqlx_mysql::sqlx_query(stmt);
crate::metric::metric_ok!(_metric_callback, stmt, {
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + Send>>
})
}
#[cfg(feature = "sqlx-postgres")]
InnerConnection::Postgres(c) => {
let query = crate::driver::sqlx_postgres::sqlx_query(stmt);
crate::metric::metric_ok!(_metric_callback, stmt, {
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + Send>>
})
}
#[cfg(feature = "sqlx-sqlite")]
InnerConnection::Sqlite(c) => {
let query = crate::driver::sqlx_sqlite::sqlx_query(stmt);
crate::metric::metric_ok!(_metric_callback, stmt, {
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + Send>>
})
}
#[cfg(feature = "mock")]
InnerConnection::Mock(c) => c.fetch(stmt),
}
}
.boxed()
}
#[instrument(level = "trace", skip(metric_callback))]
pub(crate) async fn build(
conn: MutexGuard<'a, InnerConnection>,
stmt: Statement,
metric_callback: Option<crate::metric::Callback>,
) -> TransactionStream<'a> {
let stream_builder = Self::stream_builder;
TransactionStreamAsyncBuilder {
stmt,
conn,
metric_callback,
stream_builder: |conn, stmt, _metric_callback| {
Box::pin(async move {
match conn.deref_mut() {
#[cfg(feature = "sqlx-mysql")]
InnerConnection::MySql(c) => {
let query = crate::driver::sqlx_mysql::sqlx_query(stmt);
crate::metric::metric_ok!(_metric_callback, stmt, {
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>
})
}
#[cfg(feature = "sqlx-postgres")]
InnerConnection::Postgres(c) => {
let query = crate::driver::sqlx_postgres::sqlx_query(stmt);
crate::metric::metric_ok!(_metric_callback, stmt, {
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>
})
}
#[cfg(feature = "sqlx-sqlite")]
InnerConnection::Sqlite(c) => {
let query = crate::driver::sqlx_sqlite::sqlx_query(stmt);
crate::metric::metric_ok!(_metric_callback, stmt, {
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>
})
}
#[cfg(feature = "mock")]
InnerConnection::Mock(c) => c.fetch(stmt),
}
})
},
stream_builder,
}
.build()
.await

5
src/database/transaction.rs
View File

@ -354,10 +354,11 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
fn stream(
&'a self,
stmt: Statement,
) -> Pin<Box<dyn Future<Output = Result<Self::Stream, DbErr>> + 'a>> {
) -> Pin<Box<dyn Future<Output = Result<Self::Stream, DbErr>> + 'a + Send>> {
Box::pin(async move {
let conn = self.conn.lock().await;
Ok(crate::TransactionStream::build(
self.conn.lock().await,
conn,
stmt,
self.metric_callback.clone(),
)

2
src/driver/mock.rs
View File

@ -148,7 +148,7 @@ impl MockDatabaseConnection {
pub fn fetch(
&self,
statement: &Statement,
) -> Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>> {
) -> Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>> + Send>> {
match self.query_all(statement.clone()) {
Ok(v) => Box::pin(futures::stream::iter(v.into_iter().map(Ok))),
Err(e) => Box::pin(futures::stream::iter(Some(Err(e)).into_iter())),

20
src/executor/select.rs
View File

@ -273,9 +273,9 @@ where
pub async fn stream<'a: 'b, 'b, C>(
self,
db: &'a C,
) -> Result<impl Stream<Item = Result<E::Model, DbErr>> + 'b, DbErr>
) -> Result<impl Stream<Item = Result<E::Model, DbErr>> + 'b + Send, DbErr>
where
C: ConnectionTrait<'a>,
C: ConnectionTrait<'a> + Send,
{
self.into_model().stream(db).await
}
@ -329,7 +329,7 @@ where
db: &'a C,
) -> Result<impl Stream<Item = Result<(E::Model, Option<F::Model>), DbErr>> + 'b, DbErr>
where
C: ConnectionTrait<'a>,
C: ConnectionTrait<'a> + Send,
{
self.into_model().stream(db).await
}
@ -373,9 +373,9 @@ where
pub async fn stream<'a: 'b, 'b, C>(
self,
db: &'a C,
) -> Result<impl Stream<Item = Result<(E::Model, Option<F::Model>), DbErr>> + 'b, DbErr>
) -> Result<impl Stream<Item = Result<(E::Model, Option<F::Model>), DbErr>> + 'b + Send, DbErr>
where
C: ConnectionTrait<'a>,
C: ConnectionTrait<'a> + Send,
{
self.into_model().stream(db).await
}
@ -452,10 +452,11 @@ where
pub async fn stream<'a: 'b, 'b, C>(
self,
db: &'a C,
) -> Result<Pin<Box<dyn Stream<Item = Result<S::Item, DbErr>> + 'b>>, DbErr>
) -> Result<Pin<Box<dyn Stream<Item = Result<S::Item, DbErr>> + 'b + Send>>, DbErr>
where
C: ConnectionTrait<'a>,
C: ConnectionTrait<'a> + Send,
S: 'b,
S::Item: Send,
{
self.into_selector_raw(db).stream(db).await
}
@ -737,10 +738,11 @@ where
pub async fn stream<'a: 'b, 'b, C>(
self,
db: &'a C,
) -> Result<Pin<Box<dyn Stream<Item = Result<S::Item, DbErr>> + 'b>>, DbErr>
) -> Result<Pin<Box<dyn Stream<Item = Result<S::Item, DbErr>> + 'b + Send>>, DbErr>
where
C: ConnectionTrait<'a>,
C: ConnectionTrait<'a> + Send,
S: 'b,
S::Item: Send,
{
let stream = db.stream(self.stmt).await?;
Ok(Box::pin(stream.and_then(|row| {