added CLI flag for creating API keys (tensorzero#4336)

* added CLI flag for API keys

* use preexisting constants for API key gen

Author: virajmehta

Cargo.lock

@@ -27,6 +27,8 @@ pin-project = { workspace = true }
 futures = { workspace = true }
 tokio-stream = { workspace = true }
 tensorzero-auth = { path = "../internal/tensorzero-auth" }
+sqlx_alpha = { package = "sqlx", version = "0.9.0-alpha.1", features = ["postgres", "runtime-tokio"] }
+secrecy = { workspace = true }
 
 [lints]
 workspace = true
@@ -43,4 +45,4 @@ tracing-test = "0.2"
 serde_json = { workspace = true }
 reqwest-eventsource = { workspace = true }
 futures = { workspace = true }
-secrecy = { workspace = true }
+secrecy = { workspace = true }

gateway/Cargo.toml

@@ -1,6 +1,7 @@
 use clap::{Args, Parser};
 use futures::{FutureExt, StreamExt};
 use mimalloc::MiMalloc;
+use secrecy::ExposeSecret;
 use std::fmt::Display;
 use std::future::{Future, IntoFuture};
 use std::io::ErrorKind;
@@ -12,6 +13,7 @@ use tokio::signal;
 use tokio_stream::wrappers::IntervalStream;
 use tower_http::metrics::in_flight_requests::InFlightRequestsCounter;
 
+use tensorzero_auth::constants::{DEFAULT_ORGANIZATION, DEFAULT_WORKSPACE};
 use tensorzero_core::config::{Config, ConfigFileGlob};
 use tensorzero_core::db::clickhouse::migration_manager::manual_run_clickhouse_migrations;
 use tensorzero_core::db::postgres::{manual_run_postgres_migrations, PostgresConnectionInfo};
@@ -57,6 +59,34 @@ struct MigrationCommands {
     /// Run PostgreSQL migrations manually then exit.
     #[arg(long)]
     run_postgres_migrations: bool,
+
+    /// Create an API key then exit.
+    #[arg(long)]
+    create_api_key: bool,
+}
+
+#[expect(clippy::print_stdout)]
+fn print_key(key: &secrecy::SecretString) {
+    println!("{}", key.expose_secret());
+}
+
+async fn handle_create_api_key() -> Result<(), Box<dyn std::error::Error>> {
+    // Read the Postgres URL from the environment
+    let postgres_url = std::env::var("TENSORZERO_POSTGRES_URL")
+        .map_err(|_| "TENSORZERO_POSTGRES_URL environment variable not set")?;
+
+    // Create connection pool (alpha version for tensorzero-auth)
+    let pool = sqlx_alpha::PgPool::connect(&postgres_url).await?;
+
+    // Create the key with default organization and workspace
+    let key =
+        tensorzero_auth::postgres::create_key(DEFAULT_ORGANIZATION, DEFAULT_WORKSPACE, None, &pool)
+            .await?;
+
+    // Print only the API key to stdout for easy machine parsing
+    print_key(&key);
+
+    Ok(())
 }
 
 #[tokio::main]
@@ -70,6 +100,14 @@ async fn main() {
         .expect_pretty("Failed to set up logs");
 
     let git_sha = tensorzero_core::built_info::GIT_COMMIT_HASH_SHORT.unwrap_or("unknown");
+
+    if args.migration_commands.create_api_key {
+        handle_create_api_key()
+            .await
+            .expect_pretty("Failed to create API key");
+        return;
+    }
+
     if args.migration_commands.run_clickhouse_migrations {
         manual_run_clickhouse_migrations()
             .await

gateway/src/main.rs

@@ -1,17 +1,21 @@
 #![allow(clippy::print_stdout)]
+use std::process::Stdio;
 use std::str::FromStr;
 
 use http::{Method, StatusCode};
 use serde_json::json;
 use tensorzero::test_helpers::make_embedded_gateway_with_config_and_postgres;
 use tensorzero_auth::key::TensorZeroApiKey;
 use tensorzero_core::endpoints::status::TENSORZERO_VERSION;
+use tokio::process::Command;
 
 use crate::common::start_gateway_on_random_port;
 use secrecy::ExposeSecret;
 
 mod common;
 
+const GATEWAY_PATH: &str = env!("CARGO_BIN_EXE_gateway");
+
 #[tokio::test]
 async fn test_tensorzero_auth_enabled() {
     let child_data = start_gateway_on_random_port(
@@ -302,3 +306,73 @@ async fn test_tensorzero_missing_auth() {
         assert_eq!(status, StatusCode::UNAUTHORIZED);
     }
 }
+
+#[tokio::test]
+async fn test_create_api_key_cli() {
+    // This test verifies that the --create-api-key CLI command works correctly
+    let output = Command::new(GATEWAY_PATH)
+        .args(["--create-api-key"])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .output()
+        .await
+        .unwrap();
+
+    assert!(
+        output.status.success(),
+        "CLI command failed with stderr: {}",
+        String::from_utf8_lossy(&output.stderr)
+    );
+
+    let stdout = String::from_utf8(output.stdout).unwrap();
+    let api_key = stdout.trim();
+
+    // Verify the key has the correct format
+    assert!(
+        api_key.starts_with("sk-t0-"),
+        "API key should start with 'sk-t0-', got: {api_key}"
+    );
+
+    // Verify the key can be parsed
+    let parsed_key = TensorZeroApiKey::parse(api_key);
+    assert!(
+        parsed_key.is_ok(),
+        "API key should be valid, got error: {:?}",
+        parsed_key.err()
+    );
+
+    // Verify the key works for authentication
+    let child_data = start_gateway_on_random_port(
+        "
+    [gateway.auth]
+    enabled = true
+    ",
+        None,
+    )
+    .await;
+
+    let inference_response = reqwest::Client::new()
+        .post(format!("http://{}/inference", child_data.addr))
+        .header(http::header::AUTHORIZATION, format!("Bearer {api_key}"))
+        .json(&json!({
+            "model_name": "dummy::good",
+            "input": {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "Hello, world!",
+                    }
+                ]
+            }
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    let status = inference_response.status();
+    assert_eq!(
+        status,
+        StatusCode::OK,
+        "Created API key should work for authentication"
+    );
+}