@@ -9,8 +9,11 @@ use crate::context::OpContext;
9
9
use crate :: db:: datastore:: SQL_BATCH_SIZE ;
10
10
use crate :: db:: pagination:: Paginator ;
11
11
use crate :: db:: pagination:: paginated;
12
+ use crate :: db:: queries:: ALLOW_FULL_TABLE_SCAN_SQL ;
13
+ use crate :: db:: raw_query_builder:: QueryBuilder ;
12
14
use anyhow:: Context ;
13
15
use async_bb8_diesel:: AsyncRunQueryDsl ;
16
+ use async_bb8_diesel:: AsyncSimpleConnection ;
14
17
use chrono:: DateTime ;
15
18
use chrono:: Utc ;
16
19
use clickhouse_admin_types:: { KeeperId , ServerId } ;
@@ -564,6 +567,100 @@ impl DataStore {
564
567
Ok ( ( ) )
565
568
}
566
569
570
+ /// Check whether the given blueprint limit has been reached.
571
+ ///
572
+ /// This (necessarily) does a full table scan on the blueprint table up to
573
+ /// the limit, so `limit` must be relatively small to avoid performance
574
+ /// issues. Experimentally, on a Gimlet, a limit of a few thousand takes
575
+ /// under 20ms.
576
+ pub async fn check_blueprint_limit_reached (
577
+ & self ,
578
+ opctx : & OpContext ,
579
+ limit : u64 ,
580
+ ) -> Result < BlueprintLimitReachedOutput , Error > {
581
+ // The "full" table scan below is treated as a complex operation. (This
582
+ // should only be called from the blueprint planner background task,
583
+ // for which complex operations are allowed.)
584
+ opctx. check_complex_operations_allowed ( ) ?;
585
+ opctx. authorize ( authz:: Action :: Read , & authz:: BLUEPRINT_CONFIG ) . await ?;
586
+ let conn = self . pool_connection_authorized ( opctx) . await ?;
587
+
588
+ let limit_i64 = i64:: try_from ( limit) . map_err ( |e| {
589
+ Error :: invalid_value (
590
+ "limit" ,
591
+ format ! ( "limit cannot be converted to i64: {e}" ) ,
592
+ )
593
+ } ) ?;
594
+
595
+ let err = OptionalError :: new ( ) ;
596
+
597
+ let count = self
598
+ . transaction_retry_wrapper ( "blueprint_count" )
599
+ . transaction ( & conn, |conn| {
600
+ let err = err. clone ( ) ;
601
+
602
+ async move {
603
+ // We need this to call the COUNT(*) query below. But note
604
+ // that this isn't really a "full" table scan; the number of
605
+ // rows scanned is limited by the LIMIT clause.
606
+ conn. batch_execute_async ( ALLOW_FULL_TABLE_SCAN_SQL )
607
+ . await
608
+ . map_err ( |e| err. bail ( TransactionError :: Database ( e) ) ) ?;
609
+
610
+ // Rather than doing a full table scan, we use a LIMIT
611
+ // clause to limit the number of rows returned.
612
+ let mut count_star_sql = QueryBuilder :: new ( ) ;
613
+ count_star_sql
614
+ . sql (
615
+ "SELECT COUNT(*) FROM \
616
+ (SELECT 1 FROM omicron.public.blueprint \
617
+ LIMIT $1)",
618
+ )
619
+ . bind :: < diesel:: sql_types:: BigInt , _ > ( limit_i64) ;
620
+
621
+ let query =
622
+ count_star_sql. query :: < diesel:: sql_types:: BigInt > ( ) ;
623
+
624
+ // query.first_async fails with `the trait bound
625
+ // `TypedSqlQuery<BigInt>: diesel::Table` is not satisfied`.
626
+ // So we use load_async, knowing that only one row will be
627
+ // returned.
628
+ let value = query
629
+ . load_async :: < i64 > ( & conn)
630
+ . await
631
+ . map_err ( |e| err. bail ( TransactionError :: Database ( e) ) ) ?;
632
+
633
+ Ok ( value)
634
+ }
635
+ } )
636
+ . await
637
+ . map_err ( |e| match err. take ( ) {
638
+ Some ( err) => err. into ( ) ,
639
+ None => public_error_from_diesel ( e, ErrorHandler :: Server ) ,
640
+ } ) ?;
641
+
642
+ // There must be exactly one row in the returned result.
643
+ let count = * count. get ( 0 ) . ok_or_else ( || {
644
+ Error :: internal_error ( "error getting blueprint count from database" )
645
+ } ) ?;
646
+
647
+ let count = u64:: try_from ( count) . map_err ( |_| {
648
+ Error :: internal_error ( & format ! (
649
+ "error converting blueprint count {} into \
650
+ u64 (how is it negative?)",
651
+ count
652
+ ) )
653
+ } ) ?;
654
+
655
+ // Note count >= limit (and not count > limit): for a limit of 5000 we
656
+ // want to fail if it's reached 5000.
657
+ if count >= limit {
658
+ Ok ( BlueprintLimitReachedOutput :: Yes )
659
+ } else {
660
+ Ok ( BlueprintLimitReachedOutput :: No { count } )
661
+ }
662
+ }
663
+
567
664
/// Read a complete blueprint from the database
568
665
pub async fn blueprint_read (
569
666
& self ,
@@ -2602,6 +2699,12 @@ async fn insert_pending_mgs_update(
2602
2699
Ok ( ( ) )
2603
2700
}
2604
2701
2702
+ #[ derive( Clone , Copy , Debug , PartialEq , Eq ) ]
2703
+ pub enum BlueprintLimitReachedOutput {
2704
+ No { count : u64 } ,
2705
+ Yes ,
2706
+ }
2707
+
2605
2708
// Helper to process BpPendingMgsUpdateComponent rows
2606
2709
fn process_update_row < T > (
2607
2710
row : T ,
0 commit comments