Fix dropshot-pagination _stream client codegen (#756)

oxidecomputer · Apr 3, 2024 · 1a61bf1 · 1a61bf1
1 parent 9843d26
commit 1a61bf1
Show file tree

Hide file tree

Showing 9 changed files with 1,035 additions and 1,167 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/progenitor-impl/Cargo.toml b/progenitor-impl/Cargo.toml
@@ -29,8 +29,13 @@ unicode-ident = "1.0.12"
 [dev-dependencies]
 dropshot = { git = "https://github.com/oxidecomputer/dropshot", default-features = false }
 expectorate = "1.1"
+futures = "0.3.30"
 http = "0.2.9"
 hyper = "0.14.27"
+reqwest = "0.11.27"
 rustfmt-wrapper = "0.2.1"
 serde_yaml = "0.9"
 serde_json = "1.0.113"
+tokio = { version = "1.37.0", features = ["rt", "net"] }
+
+progenitor-client.path = "../progenitor-client"
diff --git a/progenitor-impl/src/method.rs b/progenitor-impl/src/method.rs
@@ -710,9 +710,12 @@ impl Generator {
             let step_params = method.params.iter().map(|param| {
                 if param.api_name.as_str() == "page_token" {
                     quote! { state.as_deref() }
-                } else if let OperationParameterKind::Query(_) = param.kind {
-                    // Query parameters are None; having page_token as Some(_)
-                    // is mutually exclusive with other query parameters.
+                } else if param.api_name.as_str() != "limit"
+                    && matches!(param.kind, OperationParameterKind::Query(_))
+                {
+                    // Query parameters (other than "page_token" and "limit")
+                    // are None; having page_token as Some(_) is mutually
+                    // exclusive with other query parameters.
                     quote! { None }
                 } else {
                     // Non-query parameters are passed in; this is necessary
@@ -744,15 +747,6 @@ impl Generator {
                     use futures::TryFutureExt;
                     use futures::TryStreamExt;
 
-                    // Grab the limit. This is intended to be agnostic to the
-                    // specific type for the limit input which is why it's a
-                    // bit convoluted.
-                    let final_stream_limit = limit
-                        .clone()
-                        .and_then(|x| std::num::NonZeroUsize::try_from(x).ok())
-                        .map(std::num::NonZeroUsize::get)
-                        .unwrap_or(usize::MAX);
-
                     // Execute the operation with the basic parameters
                     // (omitting page_token) to get the first page.
                     self.#operation_id( #(#first_params,)* )
@@ -800,7 +794,6 @@ impl Generator {
                             first.chain(rest)
                         })
                         .try_flatten_stream()
-                        .take(final_stream_limit)
                         .boxed()
                 }
             }
@@ -1767,7 +1760,12 @@ impl Generator {
             self.uses_futures = true;
 
             let step_params = method.params.iter().filter_map(|param| {
-                if let OperationParameterKind::Query(_) = param.kind {
+                if param.api_name.as_str() != "limit"
+                    && matches!(param.kind, OperationParameterKind::Query(_))
+                {
+                    // Query parameters (other than "limit") are None; having
+                    // page_token as Some(_), as we will during the loop below,
+                    // is mutually exclusive with other query parameters.
                     let name = format_ident!("{}", param.name);
                     Some(quote! {
                         #name: Ok(None)
@@ -1799,18 +1797,6 @@ impl Generator {
                     use futures::TryFutureExt;
                     use futures::TryStreamExt;
 
-                    // Grab the limit. This is intended to be agnostic to the
-                    // specific type for the limit input which is why it's a
-                    // bit convoluted.
-                    let limit = self
-                        .limit
-                        .clone()
-                        .ok()
-                        .flatten()
-                        .and_then(|x| std::num::NonZeroUsize::try_from(x).ok())
-                        .map(std::num::NonZeroUsize::get)
-                        .unwrap_or(usize::MAX);
-
                     // This is the builder template we'll use for iterative
                     // steps past the first; it has all query params set to
                     // None (the step will fill in page_token).
@@ -1865,7 +1851,6 @@ impl Generator {
                             first.chain(rest)
                         })
                         .try_flatten_stream()
-                        .take(limit)
                         .boxed()
                 }
             }