feat(api): Limit maximum number of pipelines to 3 (#371)

Author: iambriccardo

etl-api/.sqlx/query-0f02ed5830e25fe721763c067db835db846fb41a1047310d3c8577906f469c5f.json

@@ -15,6 +15,9 @@ use crate::db::replicators::{ReplicatorsDbError, create_replicator};
 use crate::db::sources::Source;
 use crate::routes::connect_to_source_database_with_defaults;
 
+/// Maximum number of pipelines allowed per tenant.
+pub const MAX_PIPELINES_PER_TENANT: i64 = 3;
+
 pub struct Pipeline {
     pub id: i64,
     pub tenant_id: String,
@@ -47,6 +50,27 @@ pub enum PipelinesDbError {
     SlotError(#[from] slots::SlotError),
 }
 
+pub async fn count_pipelines_for_tenant<'c, E>(
+    executor: E,
+    tenant_id: &str,
+) -> Result<i64, PipelinesDbError>
+where
+    E: PgExecutor<'c>,
+{
+    let record = sqlx::query!(
+        r#"
+        select count(*) as "count!"
+        from app.pipelines
+        where tenant_id = $1
+        "#,
+        tenant_id
+    )
+    .fetch_one(executor)
+    .await?;
+
+    Ok(record.count)
+}
+
 pub async fn create_pipeline(
     txn: &mut PgTransaction<'_>,
     tenant_id: &str,

etl-api/src/db/pipelines.rs

@@ -18,7 +18,9 @@ use crate::db;
 use crate::db::destinations::{DestinationsDbError, destination_exists};
 use crate::db::destinations_pipelines::DestinationPipelinesDbError;
 use crate::db::images::ImagesDbError;
-use crate::db::pipelines::{PipelinesDbError, read_pipeline};
+use crate::db::pipelines::{
+    MAX_PIPELINES_PER_TENANT, PipelinesDbError, count_pipelines_for_tenant, read_pipeline,
+};
 use crate::db::sources::{SourcesDbError, source_exists};
 
 #[derive(Debug, Error)]
@@ -47,6 +49,9 @@ enum DestinationPipelineError {
     #[error("A pipeline already exists for this source and destination combination")]
     DuplicatePipeline,
 
+    #[error("The maximum number of pipelines ({limit}) has been reached for this project")]
+    PipelineLimitReached { limit: i64 },
+
     #[error(transparent)]
     DestinationPipelinesDb(DestinationPipelinesDbError),
 
@@ -116,6 +121,9 @@ impl ResponseError for DestinationPipelineError {
                 StatusCode::BAD_REQUEST
             }
             DestinationPipelineError::DuplicatePipeline => StatusCode::CONFLICT,
+            DestinationPipelineError::PipelineLimitReached { .. } => {
+                StatusCode::UNPROCESSABLE_ENTITY
+            }
         }
     }
 
@@ -202,6 +210,13 @@ pub async fn create_destination_and_pipeline(
         ));
     }
 
+    let pipeline_count = count_pipelines_for_tenant(txn.deref_mut(), tenant_id).await?;
+    if pipeline_count >= MAX_PIPELINES_PER_TENANT {
+        return Err(DestinationPipelineError::PipelineLimitReached {
+            limit: MAX_PIPELINES_PER_TENANT,
+        });
+    }
+
     let image = db::images::read_default_image(&**pool)
         .await?
         .ok_or(DestinationPipelineError::NoDefaultImageFound)?;

etl-api/src/routes/destinations_pipelines.rs

@@ -26,7 +26,7 @@ use crate::configs::source::StoredSourceConfig;
 use crate::db;
 use crate::db::destinations::{Destination, DestinationsDbError, destination_exists};
 use crate::db::images::{Image, ImagesDbError};
-use crate::db::pipelines::{Pipeline, PipelinesDbError};
+use crate::db::pipelines::{MAX_PIPELINES_PER_TENANT, Pipeline, PipelinesDbError};
 use crate::db::replicators::{Replicator, ReplicatorsDbError};
 use crate::db::sources::{Source, SourcesDbError, source_exists};
 use crate::k8s::http::{
@@ -103,6 +103,9 @@ pub enum PipelineError {
     #[error("The specified image id {0} does not match the default image id")]
     ImageIdNotDefault(i64),
 
+    #[error("The maximum number of pipelines ({limit}) has been reached for this project")]
+    PipelineLimitReached { limit: i64 },
+
     #[error("There was an error while looking up table information in the source database: {0}")]
     TableLookup(#[from] TableLookupError),
 
@@ -170,6 +173,7 @@ impl ResponseError for PipelineError {
                 StatusCode::BAD_REQUEST
             }
             PipelineError::DuplicatePipeline => StatusCode::CONFLICT,
+            PipelineError::PipelineLimitReached { .. } => StatusCode::UNPROCESSABLE_ENTITY,
         }
     }
 
@@ -428,6 +432,14 @@ pub async fn create_pipeline(
         return Err(PipelineError::DestinationNotFound(pipeline.destination_id));
     }
 
+    let pipeline_count =
+        db::pipelines::count_pipelines_for_tenant(txn.deref_mut(), tenant_id).await?;
+    if pipeline_count >= MAX_PIPELINES_PER_TENANT {
+        return Err(PipelineError::PipelineLimitReached {
+            limit: MAX_PIPELINES_PER_TENANT,
+        });
+    }
+
     let image = db::images::read_default_image(txn.deref_mut())
         .await?
         .ok_or(PipelineError::NoDefaultImageFound)?;

etl-api/src/routes/pipelines.rs

@@ -78,6 +78,42 @@ async fn destination_and_pipeline_can_be_created() {
     insta::assert_debug_snapshot!(response.config);
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn tenant_cannot_create_more_than_three_destinations_pipelines() {
+    init_test_tracing();
+    // Arrange
+    let app = spawn_test_app().await;
+    let tenant_id = &create_tenant(&app).await;
+    let source_id = create_source(&app, tenant_id).await;
+    create_default_image(&app).await;
+
+    for idx in 0..3 {
+        let destination_pipeline = CreateDestinationPipelineRequest {
+            destination_name: format!("BigQuery Destination {idx}"),
+            destination_config: new_destination_config(),
+            source_id,
+            pipeline_config: new_pipeline_config(),
+        };
+        let response = app
+            .create_destination_pipeline(tenant_id, &destination_pipeline)
+            .await;
+        assert!(response.status().is_success());
+    }
+
+    let destination_pipeline = CreateDestinationPipelineRequest {
+        destination_name: "BigQuery Destination 3".to_string(),
+        destination_config: new_destination_config(),
+        source_id,
+        pipeline_config: new_pipeline_config(),
+    };
+    let response = app
+        .create_destination_pipeline(tenant_id, &destination_pipeline)
+        .await;
+
+    // Assert
+    assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn destination_and_pipeline_with_another_tenants_source_cannot_be_created() {
     init_test_tracing();

etl-api/tests/destinations_pipelines.rs

@@ -221,6 +221,38 @@ async fn pipeline_can_be_created() {
     assert_eq!(response.id, 1);
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn tenant_cannot_create_more_than_three_pipelines() {
+    init_test_tracing();
+    // Arrange
+    let app = spawn_test_app().await;
+    create_default_image(&app).await;
+    let tenant_id = &create_tenant(&app).await;
+
+    for _ in 0..3 {
+        let source_id = create_source(&app, tenant_id).await;
+        let destination_id = create_destination(&app, tenant_id).await;
+        let pipeline = CreatePipelineRequest {
+            source_id,
+            destination_id,
+            config: new_pipeline_config(),
+        };
+        let response = app.create_pipeline(tenant_id, &pipeline).await;
+        assert!(response.status().is_success());
+    }
+
+    let source_id = create_source(&app, tenant_id).await;
+    let destination_id = create_destination(&app, tenant_id).await;
+    let pipeline = CreatePipelineRequest {
+        source_id,
+        destination_id,
+        config: new_pipeline_config(),
+    };
+    let response = app.create_pipeline(tenant_id, &pipeline).await;
+
+    assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn pipeline_with_another_tenants_source_cannot_be_created() {
     init_test_tracing();