diff --git a/cornucopia/src/cli.rs b/cornucopia/src/cli.rs
index e83efda3..7b0a29ee 100644
--- a/cornucopia/src/cli.rs
+++ b/cornucopia/src/cli.rs
@@ -1,6 +1,6 @@
 use clap::{Parser, Subcommand};
 
-use crate::{conn, container, error::Error, generate_live, generate_managed, CodegenSettings};
+use crate::{conn, error::Error, generate_live, generate_managed, CodegenSettings};
 
 /// Command line interface to interact with Cornucopia SQL.
 #[derive(Parser, Debug)]
@@ -65,7 +65,7 @@ pub fn run() -> Result<(), Error> {
         }
         Action::Schema { schema_files } => {
             // Run the generate command. If the command is unsuccessful, cleanup Cornucopia's container
-            if let Err(e) = generate_managed(
+            generate_managed(
                 &queries_path,
                 schema_files,
                 Some(&destination),
@@ -74,10 +74,7 @@ pub fn run() -> Result<(), Error> {
                     is_async: !sync,
                     derive_ser: serialize,
                 },
-            ) {
-                container::cleanup(podman).ok();
-                return Err(e);
-            }
+            )?;
         }
     };
     Ok(())
diff --git a/cornucopia/src/lib.rs b/cornucopia/src/lib.rs
index 6547d04f..8a631a1d 100644
--- a/cornucopia/src/lib.rs
+++ b/cornucopia/src/lib.rs
@@ -74,6 +74,18 @@ pub fn generate_managed(
     destination: Option<&str>,
     podman: bool,
     settings: CodegenSettings,
+) -> Result<String, Error> {
+    let result = generate(queries_path, podman, schema_files, settings, destination);
+    container::cleanup(podman).ok();
+    result
+}
+
+fn generate(
+    queries_path: &str,
+    podman: bool,
+    schema_files: Vec<String>,
+    settings: CodegenSettings,
+    destination: Option<&str>,
 ) -> Result<String, Error> {
     // Read
     let modules = read_query_modules(queries_path)?
@@ -85,7 +97,6 @@ pub fn generate_managed(
     load_schema(&mut client, schema_files)?;
     let prepared_modules = prepare(&mut client, modules)?;
     let generated_code = generate_internal(prepared_modules, settings);
-    container::cleanup(podman)?;
 
     if let Some(destination) = destination {
         write_generated_code(destination, &generated_code)?;