Preserve usage data in system imports

crates/aether-data/src/repository/auth/postgres.rs

@@ -668,6 +668,40 @@ RETURNING
   is_standalone
 "#;
 
+const SET_API_KEY_USAGE_TOTALS_SQL: &str = r#"
+UPDATE api_keys
+SET
+  total_requests = $2,
+  total_tokens = $3,
+  total_cost_usd = $4,
+  updated_at = NOW()
+WHERE id = $1
+RETURNING
+  user_id,
+  id AS api_key_id,
+  key_hash,
+  key_encrypted,
+  name,
+  allowed_providers,
+  allowed_api_formats,
+  allowed_models,
+  ip_rules,
+  rate_limit,
+  concurrent_limit,
+  force_capabilities,
+  feature_settings,
+  is_active,
+  CAST(EXTRACT(EPOCH FROM expires_at) AS BIGINT) AS expires_at_unix_secs,
+  auto_delete_on_expiry,
+  total_requests,
+  COALESCE(total_tokens, 0)::BIGINT AS total_tokens,
+  COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
+  CAST(EXTRACT(EPOCH FROM last_used_at) AS BIGINT) AS last_used_at_unix_secs,
+  CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
+  CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
+  is_standalone
+"#;
+
 const SET_USER_API_KEY_LOCKED_SQL: &str = r#"
 UPDATE api_keys
 SET
@@ -1436,6 +1470,24 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
             .find(|record| record.user_id == user_id && !record.is_standalone))
     }
 
+    async fn set_api_key_usage_totals(
+        &self,
+        api_key_id: &str,
+        total_requests: u64,
+        total_tokens: u64,
+        total_cost_usd: f64,
+    ) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
+        let row = sqlx::query(SET_API_KEY_USAGE_TOTALS_SQL)
+            .bind(api_key_id)
+            .bind(total_requests as i64)
+            .bind(total_tokens as i64)
+            .bind(total_cost_usd)
+            .fetch_optional(&self.pool)
+            .await
+            .map_postgres_err()?;
+        row.as_ref().map(map_auth_api_key_export_row).transpose()
+    }
+
     async fn delete_user_api_key(
         &self,
         user_id: &str,