29957817213

Committed 22 Jul 2026 08:32PM UTC coverage: 87.585% (-0.2%) from 87.738%

Build # 29957817213

Build Type

push

github

Committed by

web-flow

Commit Message

feat: add Gdal process pool (#1192)

* ipc-channel based gdal process

Co-authored-by: Mika <78550479+0xmycf@users.noreply.github.com>

* use thread and transfer only tile grid data

* LazyLock

* gdal pool

* working pool, tests, lints

* revert arrow serialization for gdal ipc data

* fix end of file

* clarify field name

* update gdalsource-process binary handling

* fix lazy gdalsource-process path detection

* justfile: add run --release

* remove or downgrade debug logs in hot paths

* gdal pool: add more complex affinity calculation

* gdal process pool with more caching workers

* simplified pool again

* batching pool

* gdal pool next iteration

* more structure

* more constants, less unwrap. Add decay.

* make load_tile_data_process more reusable

* adapt multi tile source part 1

* adapt  gdal multi source. sort errors.

* fix typo

* adapt ci bench

* remove LLVM_PROFILE from worker processes.

* update backend justfile run

* rename pool, pool worker and reorder imports

* use process-wide gdal options in gdal worker process

* put "Grid" into type name + remove comments

* explain llcov related env modification

* rename gdal dataset holding struct

* calculate composite thereshold from others

* remove comment

* rename GdaProcess config fields and add docs

* restructure gdal handling and sources

* fmt

* change env_remove logic for coverage

* refactor common worker spwaning

* stac back to 0.17

* add  leader & follower feature to gdal pool

* add default viscurl params to gdal dataset params

* lints

* sequential build bins

* complete hashing, avod stalling in group read feature, serial tests

* deaktivate porp no loger valid test

* fix: remove trailing space in VSI_CACHE_SIZE default

* chore: remove commented-out dead gdal geotransform tests

* docs: explain leader/follower cancellation semantics in read_data

* fix: set VSICURL globals once at worker startup instead of per-request

* refactor: make gdal request deduplicati... (continued)

Coverage Stats

3022 of 3724 new or added lines in 38 files covered. (81.15%)

21 existing lines in 7 files now uncovered.

123306 of 140785 relevant lines covered (87.58%)

481313.37 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

59.56

/geoengine/operators/src/engine/execution_context.rs

use super::{
    CreateSpan, InitializedPlotOperator, InitializedRasterOperator, InitializedVectorOperator,
    MockQueryContext,
};
use crate::cache::shared_cache::SharedCache;
use crate::engine::{
    ChunkByteSize, RasterResultDescriptor, ResultDescriptor, VectorResultDescriptor,
};
use crate::error::Error;
use crate::machine_learning::MlModelLoadingInfo;
use crate::meta::quota::{QuotaChecker, QuotaTracking};
use crate::meta::wrapper::InitializedOperatorWrapper;
use crate::mock::MockDatasetDataSourceLoadingInfo;
use crate::source::gdal_worker_process::{GdalProcessPool, GdalProcessPoolAccess, WorkerConfig};
use crate::source::{
    GdalLoadingInfo, MultiBandGdalLoadingInfo, MultiBandGdalLoadingInfoQueryRectangle,
    OgrSourceDataset,
};
use crate::util::{Result, create_rayon_thread_pool};
use async_trait::async_trait;
use geoengine_datatypes::dataset::{DataId, NamedData};
use geoengine_datatypes::machine_learning::MlModelName;
use geoengine_datatypes::primitives::{RasterQueryRectangle, VectorQueryRectangle};
use geoengine_datatypes::raster::TilingSpecification;
use geoengine_datatypes::util::test::TestDefault;
use rayon::ThreadPool;
use serde::{Deserialize, Serialize};
use std::any::Any;
use std::collections::HashMap;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::sync::Arc;

/// A context that provides certain utility access during operator initialization
#[async_trait::async_trait]
pub trait ExecutionContext: Send
    + Sync
    + MetaDataProvider<MockDatasetDataSourceLoadingInfo, VectorResultDescriptor, VectorQueryRectangle>
    + MetaDataProvider<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>
    + MetaDataProvider<GdalLoadingInfo, RasterResultDescriptor, RasterQueryRectangle>
    + MetaDataProvider<
        MultiBandGdalLoadingInfo,
        RasterResultDescriptor,
        MultiBandGdalLoadingInfoQueryRectangle,
    > + GdalProcessPoolAccess
{
    fn thread_pool(&self) -> &Arc<ThreadPool>;
    fn tiling_specification(&self) -> TilingSpecification;

    fn wrap_initialized_raster_operator(
        &self,
        op: Box<dyn InitializedRasterOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedRasterOperator>;

    fn wrap_initialized_vector_operator(
        &self,
        op: Box<dyn InitializedVectorOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedVectorOperator>;

    fn wrap_initialized_plot_operator(
        &self,
        op: Box<dyn InitializedPlotOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedPlotOperator>;

    async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId>;

    async fn ml_model_loading_info(&self, name: &MlModelName) -> Result<MlModelLoadingInfo>;

    fn gdal_process_pool(&self) -> &Arc<GdalProcessPool> {
        self.get_gdal_pool()
    }
}

#[async_trait]
pub trait MetaDataProvider<L, R, Q>
where
    R: ResultDescriptor,
{
    async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>>;
}

#[async_trait]
pub trait MetaData<L, R, Q>: Debug + Send + Sync
where
    R: ResultDescriptor,
{
    async fn loading_info(&self, query: Q) -> Result<L>;
    async fn result_descriptor(&self) -> Result<R>;

    fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>>;
}

impl<L, R, Q> Clone for Box<dyn MetaData<L, R, Q>>
where
    R: ResultDescriptor,
{
    fn clone(&self) -> Box<dyn MetaData<L, R, Q>> {
        self.box_clone()
    }
}

pub struct MockExecutionContext {
    pub thread_pool: Arc<ThreadPool>,
    pub meta_data: HashMap<DataId, Box<dyn Any + Send + Sync>>,
    pub named_data: HashMap<NamedData, DataId>,
    pub ml_models: HashMap<MlModelName, MlModelLoadingInfo>,
    pub tiling_specification: TilingSpecification,
    pub gdal_process_pool: Arc<GdalProcessPool>,
}

impl TestDefault for MockExecutionContext {
    fn test_default() -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(0),
            meta_data: HashMap::default(),
            named_data: HashMap::default(),
            ml_models: HashMap::default(),
            tiling_specification: TilingSpecification::test_default(),
            gdal_process_pool: GdalProcessPool::new(2, 2, 2, WorkerConfig::default()),
        }
    }
}

impl MockExecutionContext {
    pub fn new_with_tiling_spec_and_tokio_handle(
        tiling_specification: TilingSpecification,
        handle: &tokio::runtime::Handle,
    ) -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(0),
            meta_data: HashMap::default(),
            named_data: HashMap::default(),
            ml_models: HashMap::default(),
            tiling_specification,
            gdal_process_pool: GdalProcessPool::new_with_tokio_handle(
                handle,
                2,
                2,
                2,
                WorkerConfig::default(),
            ),
        }
    }

    pub fn new_with_tiling_spec(tiling_specification: TilingSpecification) -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(0),
            meta_data: HashMap::default(),
            named_data: HashMap::default(),
            ml_models: HashMap::default(),
            tiling_specification,
            gdal_process_pool: GdalProcessPool::new(2, 2, 2, WorkerConfig::default()),
        }
    }

    pub fn new_with_tiling_spec_and_thread_count(
        tiling_specification: TilingSpecification,
        num_threads: usize,
    ) -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(num_threads),
            meta_data: HashMap::default(),
            named_data: HashMap::default(),
            ml_models: HashMap::default(),
            tiling_specification,
            gdal_process_pool: GdalProcessPool::new(2, 2, 2, WorkerConfig::default()),
        }
    }

    pub fn add_meta_data<L, R, Q>(
        &mut self,
        data: DataId,
        named_data: NamedData,
        meta_data: Box<dyn MetaData<L, R, Q>>,
    ) where
        L: Send + Sync + 'static,
        R: Send + Sync + 'static + ResultDescriptor,
        Q: Send + Sync + 'static,
    {
        self.meta_data.insert(
            data.clone(),
            Box::new(meta_data) as Box<dyn Any + Send + Sync>,
        );

        self.named_data.insert(named_data, data);
    }

    pub fn delete_meta_data(&mut self, named_data: &NamedData) {
        let data = self.named_data.remove(named_data);
        if let Some(data) = data {
            self.meta_data.remove(&data);
        }
    }

    pub fn mock_query_context_test_default(&self) -> MockQueryContext {
        MockQueryContext::new(
            ChunkByteSize::test_default(),
            self.tiling_specification,
            self.gdal_process_pool.clone(),
        )
    }

    pub fn mock_query_context(&self, chunk_byte_size: ChunkByteSize) -> MockQueryContext {
        MockQueryContext::new(
            chunk_byte_size,
            self.tiling_specification,
            self.gdal_process_pool.clone(),
        )
    }

    pub fn mock_query_context_with_query_extensions(
        &self,
        chunk_byte_size: ChunkByteSize,
        cache: Option<Arc<SharedCache>>,
        quota_tracking: Option<QuotaTracking>,
        quota_checker: Option<QuotaChecker>,
    ) -> MockQueryContext {
        MockQueryContext::new_with_query_extensions(
            chunk_byte_size,
            self.tiling_specification,
            self.gdal_process_pool.clone(),
            cache,
            quota_tracking,
            quota_checker,
        )
    }

    pub fn mock_query_context_with_chunk_size_and_thread_count(
        &self,
        chunk_byte_size: ChunkByteSize,
        num_threads: usize,
    ) -> MockQueryContext {
        MockQueryContext::with_chunk_size_and_thread_count(
            chunk_byte_size,
            self.tiling_specification,
            num_threads,
            self.gdal_process_pool.clone(),
        )
    }
}

impl GdalProcessPoolAccess for MockExecutionContext {
    fn get_gdal_pool(&self) -> &Arc<GdalProcessPool> {
        &self.gdal_process_pool
    }
}

#[async_trait::async_trait]
impl ExecutionContext for MockExecutionContext {
    fn thread_pool(&self) -> &Arc<ThreadPool> {
        &self.thread_pool
    }

    fn tiling_specification(&self) -> TilingSpecification {
        self.tiling_specification
    }

    fn wrap_initialized_raster_operator(
        &self,
        op: Box<dyn InitializedRasterOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedRasterOperator> {
        op
    }

    fn wrap_initialized_vector_operator(
        &self,
        op: Box<dyn InitializedVectorOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedVectorOperator> {
        op
    }

    fn wrap_initialized_plot_operator(
        &self,
        op: Box<dyn InitializedPlotOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedPlotOperator> {
        op
    }

    async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId> {
        self.named_data
            .get(data)
            .cloned()
            .ok_or_else(|| Error::UnknownDatasetName { name: data.clone() })
    }

    async fn ml_model_loading_info(&self, name: &MlModelName) -> Result<MlModelLoadingInfo> {
        self.ml_models
            .get(name)
            .cloned()
            .ok_or_else(|| Error::UnknownMlModelName { name: name.clone() })
    }
}

#[async_trait]
impl<L, R, Q> MetaDataProvider<L, R, Q> for MockExecutionContext
where
    L: 'static,
    R: 'static + ResultDescriptor,
    Q: 'static,
{
    async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>> {
        let meta_data = self
            .meta_data
            .get(id)
            .ok_or(Error::UnknownDataId)?
            .downcast_ref::<Box<dyn MetaData<L, R, Q>>>()
            .ok_or(Error::InvalidMetaDataType)?;

        Ok(meta_data.clone())
    }
}

#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StaticMetaData<L, R, Q>
where
    L: Debug + Clone + Send + Sync + 'static,
    R: Debug + Send + Sync + 'static + ResultDescriptor,
    Q: Debug + Clone + Send + Sync + 'static,
{
    pub loading_info: L,
    pub result_descriptor: R,
    #[serde(skip)]
    pub phantom: PhantomData<Q>,
}

#[async_trait]
impl<L, R, Q> MetaData<L, R, Q> for StaticMetaData<L, R, Q>
where
    L: Debug + Clone + Send + Sync + 'static,
    R: Debug + Send + Sync + 'static + ResultDescriptor,
    Q: Debug + Clone + Send + Sync + 'static,
{
    async fn loading_info(&self, _query: Q) -> Result<L> {
        Ok(self.loading_info.clone())
    }

    async fn result_descriptor(&self) -> Result<R> {
        Ok(self.result_descriptor.clone())
    }

    fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>> {
        Box::new(self.clone())
    }
}

mod db_types {
    use geoengine_datatypes::delegate_from_to_sql;
    use postgres_types::{FromSql, ToSql};

    use super::*;

    pub type MockMetaData = StaticMetaData<
        MockDatasetDataSourceLoadingInfo,
        VectorResultDescriptor,
        VectorQueryRectangle,
    >;

    #[derive(Debug, ToSql, FromSql)]
    #[postgres(name = "MockMetaData")]
    pub struct MockMetaDataDbType {
        pub loading_info: MockDatasetDataSourceLoadingInfo,
        pub result_descriptor: VectorResultDescriptor,
    }

    impl From<&MockMetaData> for MockMetaDataDbType {
        fn from(other: &MockMetaData) -> Self {
            Self {
                loading_info: other.loading_info.clone(),
                result_descriptor: other.result_descriptor.clone(),
            }
        }
    }

    impl TryFrom<MockMetaDataDbType> for MockMetaData {
        type Error = Error;

        fn try_from(other: MockMetaDataDbType) -> Result<Self, Self::Error> {
            Ok(Self {
                loading_info: other.loading_info,
                result_descriptor: other.result_descriptor,
                phantom: PhantomData,
            })
        }
    }

    pub type OgrMetaData =
        StaticMetaData<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>;

    #[derive(Debug, ToSql, FromSql)]
    #[postgres(name = "OgrMetaData")]
    pub struct OgrMetaDataDbType {
        pub loading_info: OgrSourceDataset,
        pub result_descriptor: VectorResultDescriptor,
    }

    impl From<&StaticMetaData<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>>
        for OgrMetaDataDbType
    {
        fn from(other: &OgrMetaData) -> Self {
            Self {
                loading_info: other.loading_info.clone(),
                result_descriptor: other.result_descriptor.clone(),
            }
        }
    }

    impl TryFrom<OgrMetaDataDbType> for OgrMetaData {
        type Error = Error;

        fn try_from(other: OgrMetaDataDbType) -> Result<Self, Self::Error> {
            Ok(Self {
                loading_info: other.loading_info,
                result_descriptor: other.result_descriptor,
                phantom: PhantomData,
            })
        }
    }

    delegate_from_to_sql!(MockMetaData, MockMetaDataDbType);
    delegate_from_to_sql!(OgrMetaData, OgrMetaDataDbType);
}

/// A mock execution context that wraps all operators with a statistics operator.
pub struct StatisticsWrappingMockExecutionContext {
    pub inner: MockExecutionContext,
}

impl TestDefault for StatisticsWrappingMockExecutionContext {
    fn test_default() -> Self {
        Self {
            inner: MockExecutionContext::test_default(),
        }
    }
}

impl StatisticsWrappingMockExecutionContext {
    pub fn mock_query_context_with_query_extensions(
        &self,
        chunk_byte_size: ChunkByteSize,
        cache: Option<Arc<SharedCache>>,
        quota_tracking: Option<QuotaTracking>,
        quota_checker: Option<QuotaChecker>,
    ) -> MockQueryContext {
        self.inner.mock_query_context_with_query_extensions(
            chunk_byte_size,
            cache,
            quota_tracking,
            quota_checker,
        )
    }
}

#[async_trait::async_trait]
impl ExecutionContext for StatisticsWrappingMockExecutionContext {
    fn thread_pool(&self) -> &Arc<ThreadPool> {
        &self.inner.thread_pool
    }

    fn tiling_specification(&self) -> TilingSpecification {
        self.inner.tiling_specification
    }

    fn wrap_initialized_raster_operator(
        &self,
        op: Box<dyn InitializedRasterOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedRasterOperator> {
        InitializedOperatorWrapper::new(op, span).boxed()
    }

    fn wrap_initialized_vector_operator(
        &self,
        op: Box<dyn InitializedVectorOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedVectorOperator> {
        InitializedOperatorWrapper::new(op, span).boxed()
    }

    fn wrap_initialized_plot_operator(
        &self,
        op: Box<dyn InitializedPlotOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedPlotOperator> {
        op
    }

    async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId> {
        self.inner.resolve_named_data(data).await
    }

    async fn ml_model_loading_info(&self, name: &MlModelName) -> Result<MlModelLoadingInfo> {
        self.inner.ml_model_loading_info(name).await
    }
}

impl GdalProcessPoolAccess for StatisticsWrappingMockExecutionContext {
    fn get_gdal_pool(&self) -> &Arc<GdalProcessPool> {
        self.inner.get_gdal_pool()
    }
}

#[async_trait]
impl<L, R, Q> MetaDataProvider<L, R, Q> for StatisticsWrappingMockExecutionContext
where
    L: 'static,
    R: 'static + ResultDescriptor,
    Q: 'static,
{
    async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>> {
        self.inner.meta_data(id).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use geoengine_datatypes::collections::VectorDataType;
    use geoengine_datatypes::spatial_reference::SpatialReferenceOption;

    #[tokio::test]
    async fn test() {
        let info = StaticMetaData {
            loading_info: 1_i32,
            result_descriptor: VectorResultDescriptor {
                data_type: VectorDataType::Data,
                spatial_reference: SpatialReferenceOption::Unreferenced,
                columns: Default::default(),
                time: None,
                bbox: None,
            },
            phantom: Default::default(),
        };

        let info: Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>> =
            Box::new(info);

        let info2: Box<dyn Any + Send + Sync> = Box::new(info);

        let info3 = info2
            .downcast_ref::<Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>>>()
            .unwrap();

        assert_eq!(
            info3.result_descriptor().await.unwrap(),
            VectorResultDescriptor {
                data_type: VectorDataType::Data,
                spatial_reference: SpatialReferenceOption::Unreferenced,
                columns: Default::default(),
                time: None,
                bbox: None,
            }
        );
    }
}

1	use super::{
2	CreateSpan, InitializedPlotOperator, InitializedRasterOperator, InitializedVectorOperator,
3	MockQueryContext,
4	};
5	use crate::cache::shared_cache::SharedCache;
6	use crate::engine::{
7	ChunkByteSize, RasterResultDescriptor, ResultDescriptor, VectorResultDescriptor,
8	};
9	use crate::error::Error;
10	use crate::machine_learning::MlModelLoadingInfo;
11	use crate::meta::quota::{QuotaChecker, QuotaTracking};
12	use crate::meta::wrapper::InitializedOperatorWrapper;
13	use crate::mock::MockDatasetDataSourceLoadingInfo;
14	use crate::source::gdal_worker_process::{GdalProcessPool, GdalProcessPoolAccess, WorkerConfig};
15	use crate::source::{
16	GdalLoadingInfo, MultiBandGdalLoadingInfo, MultiBandGdalLoadingInfoQueryRectangle,
17	OgrSourceDataset,
18	};
19	use crate::util::{Result, create_rayon_thread_pool};
20	use async_trait::async_trait;
21	use geoengine_datatypes::dataset::{DataId, NamedData};
22	use geoengine_datatypes::machine_learning::MlModelName;
23	use geoengine_datatypes::primitives::{RasterQueryRectangle, VectorQueryRectangle};
24	use geoengine_datatypes::raster::TilingSpecification;
25	use geoengine_datatypes::util::test::TestDefault;
26	use rayon::ThreadPool;
27	use serde::{Deserialize, Serialize};
28	use std::any::Any;
29	use std::collections::HashMap;
30	use std::fmt::Debug;
31	use std::marker::PhantomData;
32	use std::sync::Arc;
33
34	/// A context that provides certain utility access during operator initialization
35	#[async_trait::async_trait]
36	pub trait ExecutionContext: Send
37	+ Sync
38	+ MetaDataProvider<MockDatasetDataSourceLoadingInfo, VectorResultDescriptor, VectorQueryRectangle>
39	+ MetaDataProvider<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>
40	+ MetaDataProvider<GdalLoadingInfo, RasterResultDescriptor, RasterQueryRectangle>
41	+ MetaDataProvider<
42	MultiBandGdalLoadingInfo,
43	RasterResultDescriptor,
44	MultiBandGdalLoadingInfoQueryRectangle,
45	> + GdalProcessPoolAccess
46	{
47	fn thread_pool(&self) -> &Arc<ThreadPool>;
48	fn tiling_specification(&self) -> TilingSpecification;
49
50	fn wrap_initialized_raster_operator(
51	&self,
52	op: Box<dyn InitializedRasterOperator>,
53	span: CreateSpan,
54	) -> Box<dyn InitializedRasterOperator>;
55
56	fn wrap_initialized_vector_operator(
57	&self,
58	op: Box<dyn InitializedVectorOperator>,
59	span: CreateSpan,
60	) -> Box<dyn InitializedVectorOperator>;
61
62	fn wrap_initialized_plot_operator(
63	&self,
64	op: Box<dyn InitializedPlotOperator>,
65	span: CreateSpan,
66	) -> Box<dyn InitializedPlotOperator>;
67
68	async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId>;
69
70	async fn ml_model_loading_info(&self, name: &MlModelName) -> Result<MlModelLoadingInfo>;
71
NEW 72	fn gdal_process_pool(&self) -> &Arc<GdalProcessPool> {	×
NEW 73	self.get_gdal_pool()	×
NEW 74	}	×
75	}
76
77	#[async_trait]
78	pub trait MetaDataProvider<L, R, Q>
79	where
80	R: ResultDescriptor,
81	{
82	async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>>;
83	}
84
85	#[async_trait]
86	pub trait MetaData<L, R, Q>: Debug + Send + Sync
87	where
88	R: ResultDescriptor,
89	{
90	async fn loading_info(&self, query: Q) -> Result<L>;
91	async fn result_descriptor(&self) -> Result<R>;
92
93	fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>>;
94	}
95
96	impl<L, R, Q> Clone for Box<dyn MetaData<L, R, Q>>
97	where
98	R: ResultDescriptor,
99	{
100	fn clone(&self) -> Box<dyn MetaData<L, R, Q>> {	213✔
101	self.box_clone()	213✔
102	}	213✔
103	}
104
105	pub struct MockExecutionContext {
106	pub thread_pool: Arc<ThreadPool>,
107	pub meta_data: HashMap<DataId, Box<dyn Any + Send + Sync>>,
108	pub named_data: HashMap<NamedData, DataId>,
109	pub ml_models: HashMap<MlModelName, MlModelLoadingInfo>,
110	pub tiling_specification: TilingSpecification,
111	pub gdal_process_pool: Arc<GdalProcessPool>,
112	}
113
114	impl TestDefault for MockExecutionContext {
115	fn test_default() -> Self {	191✔
116	Self {	191✔
117	thread_pool: create_rayon_thread_pool(0),	191✔
118	meta_data: HashMap::default(),	191✔
119	named_data: HashMap::default(),	191✔
120	ml_models: HashMap::default(),	191✔
121	tiling_specification: TilingSpecification::test_default(),	191✔
122	gdal_process_pool: GdalProcessPool::new(2, 2, 2, WorkerConfig::default()),	191✔
123	}	191✔
124	}	191✔
125	}
126
127	impl MockExecutionContext {
NEW 128	pub fn new_with_tiling_spec_and_tokio_handle(	×
NEW 129	tiling_specification: TilingSpecification,	×
NEW 130	handle: &tokio::runtime::Handle,	×
NEW 131	) -> Self {	×
NEW 132	Self {	×
NEW 133	thread_pool: create_rayon_thread_pool(0),	×
NEW 134	meta_data: HashMap::default(),	×
NEW 135	named_data: HashMap::default(),	×
NEW 136	ml_models: HashMap::default(),	×
NEW 137	tiling_specification,	×
NEW 138	gdal_process_pool: GdalProcessPool::new_with_tokio_handle(	×
NEW 139	handle,	×
NEW 140	2,	×
NEW 141	2,	×
NEW 142	2,	×
NEW 143	WorkerConfig::default(),	×
NEW 144	),	×
NEW 145	}	×
NEW 146	}	×
147
148	pub fn new_with_tiling_spec(tiling_specification: TilingSpecification) -> Self {	141✔
149	Self {	141✔
150	thread_pool: create_rayon_thread_pool(0),	141✔
151	meta_data: HashMap::default(),	141✔
152	named_data: HashMap::default(),	141✔
153	ml_models: HashMap::default(),	141✔
154	tiling_specification,	141✔
155	gdal_process_pool: GdalProcessPool::new(2, 2, 2, WorkerConfig::default()),	141✔
156	}	141✔
157	}	141✔
158
159	pub fn new_with_tiling_spec_and_thread_count(	2✔
160	tiling_specification: TilingSpecification,	2✔
161	num_threads: usize,	2✔
162	) -> Self {	2✔
163	Self {	2✔
164	thread_pool: create_rayon_thread_pool(num_threads),	2✔
165	meta_data: HashMap::default(),	2✔
166	named_data: HashMap::default(),	2✔
167	ml_models: HashMap::default(),	2✔
168	tiling_specification,	2✔
169	gdal_process_pool: GdalProcessPool::new(2, 2, 2, WorkerConfig::default()),	2✔
170	}	2✔
171	}	2✔
172
173	pub fn add_meta_data<L, R, Q>(	78✔
174	&mut self,	78✔
175	data: DataId,	78✔
176	named_data: NamedData,	78✔
177	meta_data: Box<dyn MetaData<L, R, Q>>,	78✔
178	) where	78✔
179	L: Send + Sync + 'static,	78✔
180	R: Send + Sync + 'static + ResultDescriptor,	78✔
181	Q: Send + Sync + 'static,	78✔
182	{
183	self.meta_data.insert(	78✔
184	data.clone(),	78✔
185	Box::new(meta_data) as Box<dyn Any + Send + Sync>,	78✔
186	);
187
188	self.named_data.insert(named_data, data);	78✔
189	}	78✔
190
191	pub fn delete_meta_data(&mut self, named_data: &NamedData) {	3✔
192	let data = self.named_data.remove(named_data);	3✔
193	if let Some(data) = data {	3✔
194	self.meta_data.remove(&data);	3✔
195	}	3✔
196	}	3✔
197
198	pub fn mock_query_context_test_default(&self) -> MockQueryContext {	132✔
199	MockQueryContext::new(	132✔
200	ChunkByteSize::test_default(),	132✔
201	self.tiling_specification,	132✔
202	self.gdal_process_pool.clone(),	132✔
203	)
204	}	132✔
205
206	pub fn mock_query_context(&self, chunk_byte_size: ChunkByteSize) -> MockQueryContext {	152✔
207	MockQueryContext::new(	152✔
208	chunk_byte_size,	152✔
209	self.tiling_specification,	152✔
210	self.gdal_process_pool.clone(),	152✔
211	)
212	}	152✔
213
214	pub fn mock_query_context_with_query_extensions(	3✔
215	&self,	3✔
216	chunk_byte_size: ChunkByteSize,	3✔
217	cache: Option<Arc<SharedCache>>,	3✔
218	quota_tracking: Option<QuotaTracking>,	3✔
219	quota_checker: Option<QuotaChecker>,	3✔
220	) -> MockQueryContext {	3✔
221	MockQueryContext::new_with_query_extensions(	3✔
222	chunk_byte_size,	3✔
223	self.tiling_specification,	3✔
224	self.gdal_process_pool.clone(),	3✔
225	cache,	3✔
226	quota_tracking,	3✔
227	quota_checker,	3✔
228	)
229	}	3✔
230
231	pub fn mock_query_context_with_chunk_size_and_thread_count(	×
232	&self,	×
233	chunk_byte_size: ChunkByteSize,	×
234	num_threads: usize,	×
235	) -> MockQueryContext {	×
236	MockQueryContext::with_chunk_size_and_thread_count(	×
237	chunk_byte_size,	×
238	self.tiling_specification,	×
239	num_threads,	×
NEW 240	self.gdal_process_pool.clone(),	×
241	)
242	}	×
243	}
244
245	impl GdalProcessPoolAccess for MockExecutionContext {
NEW 246	fn get_gdal_pool(&self) -> &Arc<GdalProcessPool> {	×
NEW 247	&self.gdal_process_pool	×
NEW 248	}	×
249	}
250
251	#[async_trait::async_trait]
252	impl ExecutionContext for MockExecutionContext {
253	fn thread_pool(&self) -> &Arc<ThreadPool> {	×
254	&self.thread_pool	×
255	}	×
256
257	fn tiling_specification(&self) -> TilingSpecification {	310✔
258	self.tiling_specification	310✔
259	}	310✔
260
261	fn wrap_initialized_raster_operator(	392✔
262	&self,	392✔
263	op: Box<dyn InitializedRasterOperator>,	392✔
264	_span: CreateSpan,	392✔
265	) -> Box<dyn InitializedRasterOperator> {	392✔
266	op	392✔
267	}	392✔
268
269	fn wrap_initialized_vector_operator(	182✔
270	&self,	182✔
271	op: Box<dyn InitializedVectorOperator>,	182✔
272	_span: CreateSpan,	182✔
273	) -> Box<dyn InitializedVectorOperator> {	182✔
274	op	182✔
275	}	182✔
276
277	fn wrap_initialized_plot_operator(	55✔
278	&self,	55✔
279	op: Box<dyn InitializedPlotOperator>,	55✔
280	_span: CreateSpan,	55✔
281	) -> Box<dyn InitializedPlotOperator> {	55✔
282	op	55✔
283	}	55✔
284
285	async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId> {	105✔
286	self.named_data
287	.get(data)
288	.cloned()
289	.ok_or_else(\|\| Error::UnknownDatasetName { name: data.clone() })	×
290	}	105✔
291
292	async fn ml_model_loading_info(&self, name: &MlModelName) -> Result<MlModelLoadingInfo> {	3✔
293	self.ml_models
294	.get(name)
295	.cloned()
296	.ok_or_else(\|\| Error::UnknownMlModelName { name: name.clone() })	×
297	}	3✔
298	}
299
300	#[async_trait]
301	impl<L, R, Q> MetaDataProvider<L, R, Q> for MockExecutionContext
302	where
303	L: 'static,
304	R: 'static + ResultDescriptor,
305	Q: 'static,
306	{
307	async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>> {	105✔
308	let meta_data = self
309	.meta_data
310	.get(id)
311	.ok_or(Error::UnknownDataId)?
312	.downcast_ref::<Box<dyn MetaData<L, R, Q>>>()
313	.ok_or(Error::InvalidMetaDataType)?;
314
315	Ok(meta_data.clone())
316	}	105✔
317	}
318
319	#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
320	#[serde(rename_all = "camelCase")]
321	pub struct StaticMetaData<L, R, Q>
322	where
323	L: Debug + Clone + Send + Sync + 'static,
324	R: Debug + Send + Sync + 'static + ResultDescriptor,
325	Q: Debug + Clone + Send + Sync + 'static,
326	{
327	pub loading_info: L,
328	pub result_descriptor: R,
329	#[serde(skip)]
330	pub phantom: PhantomData<Q>,
331	}
332
333	#[async_trait]
334	impl<L, R, Q> MetaData<L, R, Q> for StaticMetaData<L, R, Q>
335	where
336	L: Debug + Clone + Send + Sync + 'static,
337	R: Debug + Send + Sync + 'static + ResultDescriptor,
338	Q: Debug + Clone + Send + Sync + 'static,
339	{
340	async fn loading_info(&self, _query: Q) -> Result<L> {	69✔
341	Ok(self.loading_info.clone())
342	}	69✔
343
344	async fn result_descriptor(&self) -> Result<R> {	53✔
345	Ok(self.result_descriptor.clone())
346	}	53✔
347
348	fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>> {	70✔
349	Box::new(self.clone())	70✔
350	}	70✔
351	}
352
353	mod db_types {
354	use geoengine_datatypes::delegate_from_to_sql;
355	use postgres_types::{FromSql, ToSql};
356
357	use super::*;
358
359	pub type MockMetaData = StaticMetaData<
360	MockDatasetDataSourceLoadingInfo,
361	VectorResultDescriptor,
362	VectorQueryRectangle,
363	>;
364
365	#[derive(Debug, ToSql, FromSql)]	×
366	#[postgres(name = "MockMetaData")]
367	pub struct MockMetaDataDbType {
368	pub loading_info: MockDatasetDataSourceLoadingInfo,
369	pub result_descriptor: VectorResultDescriptor,
370	}
371
372	impl From<&MockMetaData> for MockMetaDataDbType {
373	fn from(other: &MockMetaData) -> Self {	×
374	Self {	×
375	loading_info: other.loading_info.clone(),	×
376	result_descriptor: other.result_descriptor.clone(),	×
377	}	×
378	}	×
379	}
380
381	impl TryFrom<MockMetaDataDbType> for MockMetaData {
382	type Error = Error;
383
384	fn try_from(other: MockMetaDataDbType) -> Result<Self, Self::Error> {	×
385	Ok(Self {	×
386	loading_info: other.loading_info,	×
387	result_descriptor: other.result_descriptor,	×
388	phantom: PhantomData,	×
389	})	×
390	}	×
391	}
392
393	pub type OgrMetaData =
394	StaticMetaData<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>;
395
396	#[derive(Debug, ToSql, FromSql)]	×
397	#[postgres(name = "OgrMetaData")]
398	pub struct OgrMetaDataDbType {
399	pub loading_info: OgrSourceDataset,
400	pub result_descriptor: VectorResultDescriptor,
401	}
402
403	impl From<&StaticMetaData<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>>
404	for OgrMetaDataDbType
405	{
406	fn from(other: &OgrMetaData) -> Self {	28✔
407	Self {	28✔
408	loading_info: other.loading_info.clone(),	28✔
409	result_descriptor: other.result_descriptor.clone(),	28✔
410	}	28✔
411	}	28✔
412	}
413
414	impl TryFrom<OgrMetaDataDbType> for OgrMetaData {
415	type Error = Error;
416
417	fn try_from(other: OgrMetaDataDbType) -> Result<Self, Self::Error> {	13✔
418	Ok(Self {	13✔
419	loading_info: other.loading_info,	13✔
420	result_descriptor: other.result_descriptor,	13✔
421	phantom: PhantomData,	13✔
422	})	13✔
423	}	13✔
424	}
425
426	delegate_from_to_sql!(MockMetaData, MockMetaDataDbType);
427	delegate_from_to_sql!(OgrMetaData, OgrMetaDataDbType);
428	}
429
430	/// A mock execution context that wraps all operators with a statistics operator.
431	pub struct StatisticsWrappingMockExecutionContext {
432	pub inner: MockExecutionContext,
433	}
434
435	impl TestDefault for StatisticsWrappingMockExecutionContext {
436	fn test_default() -> Self {	×
437	Self {	×
438	inner: MockExecutionContext::test_default(),	×
439	}	×
440	}	×
441	}
442
443	impl StatisticsWrappingMockExecutionContext {
444	pub fn mock_query_context_with_query_extensions(	×
445	&self,	×
446	chunk_byte_size: ChunkByteSize,	×
447	cache: Option<Arc<SharedCache>>,	×
448	quota_tracking: Option<QuotaTracking>,	×
449	quota_checker: Option<QuotaChecker>,	×
450	) -> MockQueryContext {	×
451	self.inner.mock_query_context_with_query_extensions(	×
452	chunk_byte_size,	×
453	cache,	×
454	quota_tracking,	×
455	quota_checker,	×
456	)
457	}	×
458	}
459
460	#[async_trait::async_trait]
461	impl ExecutionContext for StatisticsWrappingMockExecutionContext {
462	fn thread_pool(&self) -> &Arc<ThreadPool> {	×
463	&self.inner.thread_pool	×
464	}	×
465
466	fn tiling_specification(&self) -> TilingSpecification {	×
467	self.inner.tiling_specification	×
468	}	×
469
470	fn wrap_initialized_raster_operator(	×
471	&self,	×
472	op: Box<dyn InitializedRasterOperator>,	×
473	span: CreateSpan,	×
474	) -> Box<dyn InitializedRasterOperator> {	×
475	InitializedOperatorWrapper::new(op, span).boxed()	×
476	}	×
477
478	fn wrap_initialized_vector_operator(	×
479	&self,	×
480	op: Box<dyn InitializedVectorOperator>,	×
481	span: CreateSpan,	×
482	) -> Box<dyn InitializedVectorOperator> {	×
483	InitializedOperatorWrapper::new(op, span).boxed()	×
484	}	×
485
486	fn wrap_initialized_plot_operator(	×
487	&self,	×
488	op: Box<dyn InitializedPlotOperator>,	×
489	_span: CreateSpan,	×
490	) -> Box<dyn InitializedPlotOperator> {	×
491	op	×
492	}	×
493
494	async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId> {	×
495	self.inner.resolve_named_data(data).await
496	}	×
497
498	async fn ml_model_loading_info(&self, name: &MlModelName) -> Result<MlModelLoadingInfo> {	×
499	self.inner.ml_model_loading_info(name).await
500	}	×
501	}
502
503	impl GdalProcessPoolAccess for StatisticsWrappingMockExecutionContext {
NEW 504	fn get_gdal_pool(&self) -> &Arc<GdalProcessPool> {	×
NEW 505	self.inner.get_gdal_pool()	×
NEW 506	}	×
507	}
508
509	#[async_trait]
510	impl<L, R, Q> MetaDataProvider<L, R, Q> for StatisticsWrappingMockExecutionContext
511	where
512	L: 'static,
513	R: 'static + ResultDescriptor,
514	Q: 'static,
515	{
516	async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>> {	×
517	self.inner.meta_data(id).await
518	}	×
519	}
520
521	#[cfg(test)]
522	mod tests {
523	use super::*;
524	use geoengine_datatypes::collections::VectorDataType;
525	use geoengine_datatypes::spatial_reference::SpatialReferenceOption;
526
527	#[tokio::test]
528	async fn test() {	1✔
529	let info = StaticMetaData {	1✔
530	loading_info: 1_i32,	1✔
531	result_descriptor: VectorResultDescriptor {	1✔
532	data_type: VectorDataType::Data,	1✔
533	spatial_reference: SpatialReferenceOption::Unreferenced,	1✔
534	columns: Default::default(),	1✔
535	time: None,	1✔
536	bbox: None,	1✔
537	},	1✔
538	phantom: Default::default(),	1✔
539	};	1✔
540
541	let info: Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>> =	1✔
542	Box::new(info);	1✔
543
544	let info2: Box<dyn Any + Send + Sync> = Box::new(info);	1✔
545
546	let info3 = info2	1✔
547	.downcast_ref::<Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>>>()	1✔
548	.unwrap();	1✔
549
550	assert_eq!(	1✔
551	info3.result_descriptor().await.unwrap(),	1✔
552	VectorResultDescriptor {	1✔
553	data_type: VectorDataType::Data,	1✔
554	spatial_reference: SpatialReferenceOption::Unreferenced,	1✔
555	columns: Default::default(),	1✔
556	time: None,	1✔
557	bbox: None,	1✔
558	}	1✔
559	);	1✔
560	}	1✔
561	}

geo-engine / geoengine / 29957817213

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous