Improved server performance

.github/workflows/main.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-latest, windows-latest]
+        os: [ubuntu-latest, windows-latest, macos-latest]
         include:
           - os: ubuntu-latest
             build_target: x86_64-unknown-linux-musl
@@ -23,6 +23,9 @@ jobs:
           - os: windows-latest
             build_target: x86_64-pc-windows-msvc
             platform_name: windows-amd64
+          - os: macos-latest
+            build_target: aarch64-apple-darwin
+            platform_name: macos-arm64
     permissions:
       contents: write
     steps:
@@ -54,10 +57,6 @@ jobs:
 
       - uses: dtolnay/rust-toolchain@stable
       
-      - uses: mbrobbel/rustfmt-check@master
-        with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          
       - name: Install rust targets
         run: rustup target add ${{ matrix.build_target }}
         
@@ -67,8 +66,12 @@ jobs:
         with:
           args: cargo build --target ${{ matrix.build_target }} --release
           
+      - name: Build MacOS
+        if: matrix.os == 'macos-latest'
+        run: cargo build --target ${{ matrix.build_target }} --release
+          
       - name: Build Windows
-        if: matrix.os != 'ubuntu-latest'
+        if: matrix.os == 'windows-latest'
         run: cargo build --target ${{ matrix.build_target }} --release
           
       - name: Upload artifact
@@ -104,12 +107,14 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-latest, windows-latest]
+        os: [ubuntu-latest, windows-latest, macos-latest]
         include:
           - os: ubuntu-latest
             platform_name: linux-amd64
           - os: windows-latest
             platform_name: windows-amd64
+          - os: macos-latest
+            platform_name: macos-arm64
     steps:
       - uses: actions/checkout@v4
       

.gitignore

@@ -1 +1,3 @@
 /target
+*.swp
+*.swo

Cargo.lock

@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "actix-codec"
@@ -1053,9 +1053,10 @@ dependencies = [
 
 [[package]]
 name = "khm"
-version = "0.2.0"
+version = "0.4.0"
 dependencies = [
  "actix-web",
+ "base64 0.21.7",
  "chrono",
  "clap",
  "env_logger",

Cargo.toml

@@ -1,8 +1,8 @@
 [package]
 name = "khm"
-version = "0.2.1"
+version = "0.4.0"
 edition = "2021"
-authors = ["AB <ab@hexor.cy>", "ChatGPT-4o"]
+authors = ["AB <ab@hexor.cy>"]
 
 [dependencies]
 actix-web = "4"
@@ -11,6 +11,7 @@ serde_json = "1.0"
 env_logger = "0.11.3"
 log = "0.4"
 regex = "1.10.5"
+base64 = "0.21"
 tokio = { version = "1", features = ["full"] }
 tokio-postgres = { version = "0.7", features = ["with-chrono-0_4"] }
 clap = { version = "4", features = ["derive"] }

README.MD

@@ -24,21 +24,23 @@ khm --server --ip 127.0.0.1 --port 8080 --db-host 127.0.0.1 --db-name khm --db-u
 To run the application in client mode, use the following command:
 
 ```bash
-khm --host http://khm.example.com:8080 --known-hosts ~/.ssh/known_hosts --in-place
+khm --host http://khm.example.com:8080/<FLOW_NAME>/ --known-hosts ~/.ssh/known_hosts --in-place
 ```
 
 ### Arguments
-- `--server`: Run in server mode (default: false).
-- `--ip`: IP address to bind the server or client to (default: 127.0.0.1).
-- `--port`: Port to bind the server or client to (default: 8080).
-- `--db-host`: Hostname or IP address of the PostgreSQL database (default: 127.0.0.1).
-- `--db-name`: Name of the PostgreSQL database (default: khm).
-- `--db-user`: Username for the PostgreSQL database (required in server mode).
-- `--db-password`: Password for the PostgreSQL database (required in server mode).
-- `--host`: Host address of the server to connect to in client mode (required in client mode).
-- `--known-hosts`: Path to the `known_hosts` file (default: ~/.ssh/known_hosts).
-- `--in-place`: Update the `known_hosts` file with keys from the server after sending keys (default: false).
-- `--flows`: Comma-separated list of flows to manage (default: default).
+Options:
+- `--server` Run in server mode
+- `--in-place` Server mode: Sync the known_hosts file with keys from the server
+- `--flows <FLOWS>...` Server mode: Comma-separated list of flows to manage [default: default]
+- `-i, --ip <IP>` Server mode: IP address to bind the server to [default: 127.0.0.1]
+- `-p, --port <PORT>` Server mode: Port to bind the server to [default: 8080]
+- `--db-host <DB_HOST>` Server mode: Hostname or IP address of the PostgreSQL database [default: 127.0.0.1]
+- `--db-name <DB_NAME>` Server mode: Name of the PostgreSQL database [default: khm]
+- `--db-user <DB_USER>` Server mode: Username for the PostgreSQL database
+- `--db-password <DB_PASSWORD>` Server mode: Password for the PostgreSQL database
+- `--host <HOST>` Client mode: Full host address of the server to connect to. Like `https://khm.example.com/<FLOW_NAME>`
+- `--known-hosts <KNOWN_HOSTS>` Client mode: Path to the known_hosts file [default: ~/.ssh/known_hosts]
+
 
 ## Installation
 

src/client.rs

@@ -1,4 +1,6 @@
+use base64::{engine::general_purpose, Engine as _};
 use log::{error, info};
+use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION};
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use std::fs::File;
@@ -48,10 +50,39 @@ fn write_known_hosts(file_path: &str, keys: &[SshKey]) -> io::Result<()> {
     Ok(())
 }
 
-async fn send_keys_to_server(host: &str, keys: Vec<SshKey>) -> Result<(), reqwest::Error> {
+async fn send_keys_to_server(
+    host: &str,
+    keys: Vec<SshKey>,
+    auth_string: &str,
+) -> Result<(), reqwest::Error> {
     let client = Client::new();
     let url = format!("{}/keys", host);
-    let response = client.post(&url).json(&keys).send().await?;
+    info!("URL: {} ", url);
+
+    let mut headers = HeaderMap::new();
+
+    if !auth_string.is_empty() {
+        let parts: Vec<&str> = auth_string.splitn(2, ':').collect();
+        if parts.len() == 2 {
+            let username = parts[0];
+            let password = parts[1];
+
+            let auth_header_value = format!("{}:{}", username, password);
+            let encoded_auth = general_purpose::STANDARD.encode(auth_header_value);
+            let auth_header = format!("Basic {}", encoded_auth);
+
+            headers.insert(AUTHORIZATION, HeaderValue::from_str(&auth_header).unwrap());
+        } else {
+            error!("Invalid auth string format. Expected 'username:password'");
+        }
+    }
+
+    let response = client
+        .post(&url)
+        .headers(headers)
+        .json(&keys)
+        .send()
+        .await?;
 
     if response.status().is_success() {
         info!("Keys successfully sent to server.");
@@ -65,22 +96,38 @@ async fn send_keys_to_server(host: &str, keys: Vec<SshKey>) -> Result<(), reqwes
     Ok(())
 }
 
-async fn get_keys_from_server(host: &str) -> Result<Vec<SshKey>, reqwest::Error> {
+async fn get_keys_from_server(
+    host: &str,
+    auth_string: &str,
+) -> Result<Vec<SshKey>, reqwest::Error> {
     let client = Client::new();
     let url = format!("{}/keys", host);
-    let response = client.get(&url).send().await?;
 
-    if response.status().is_success() {
-        let keys: Vec<SshKey> = response.json().await?;
-        info!("Received {} keys from server", keys.len());
-        Ok(keys)
-    } else {
-        error!(
-            "Failed to get keys from server. Status: {}",
-            response.status()
-        );
-        Ok(vec![])
+    let mut headers = HeaderMap::new();
+
+    if !auth_string.is_empty() {
+        let parts: Vec<&str> = auth_string.splitn(2, ':').collect();
+        if parts.len() == 2 {
+            let username = parts[0];
+            let password = parts[1];
+
+            let auth_header_value = format!("{}:{}", username, password);
+            let encoded_auth = general_purpose::STANDARD.encode(auth_header_value);
+            let auth_header = format!("Basic {}", encoded_auth);
+
+            headers.insert(AUTHORIZATION, HeaderValue::from_str(&auth_header).unwrap());
+        } else {
+            error!("Invalid auth string format. Expected 'username:password'");
+        }
     }
+
+    let response = client.get(&url).headers(headers).send().await?;
+
+    let response = response.error_for_status()?;
+
+    let keys: Vec<SshKey> = response.json().await?;
+    info!("Received {} keys from server", keys.len());
+    Ok(keys)
 }
 
 pub async fn run_client(args: crate::Args) -> std::io::Result<()> {
@@ -89,13 +136,13 @@ pub async fn run_client(args: crate::Args) -> std::io::Result<()> {
 
     let host = args.host.expect("host is required in client mode");
     info!("Client mode: Sending keys to server at {}", host);
-    send_keys_to_server(&host, keys)
+    send_keys_to_server(&host, keys, &args.basic_auth)
         .await
         .expect("Failed to send keys to server");
 
     if args.in_place {
         info!("Client mode: In-place update is enabled. Fetching keys from server.");
-        let server_keys = get_keys_from_server(&host)
+        let server_keys = get_keys_from_server(&host, &args.basic_auth)
             .await
             .expect("Failed to get keys from server");
 

src/db.rs

@@ -0,0 +1,298 @@
+use crate::server::SshKey;
+use log::info;
+use std::collections::HashMap;
+use std::collections::HashSet;
+use tokio_postgres::Client;
+
+// Структура для хранения статистики обработки ключей
+pub struct KeyInsertStats {
+    pub total: usize,                   // Общее количество полученных ключей
+    pub inserted: usize,                // Количество новых ключей
+    pub updated: usize,                 // Количество обновленных ключей
+    pub unchanged: usize,               // Количество неизмененных ключей
+    pub key_id_map: Vec<(SshKey, i32)>, // Связь ключей с их ID в базе
+}
+
+pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres::Error> {
+    info!("Checking and initializing database schema if needed");
+
+    // Check if tables exist by querying information_schema
+    let tables_exist = client
+        .query(
+            "SELECT EXISTS (
+                SELECT FROM information_schema.tables 
+                WHERE table_schema = 'public' 
+                AND table_name = 'keys'
+            ) AND EXISTS (
+                SELECT FROM information_schema.tables 
+                WHERE table_schema = 'public' 
+                AND table_name = 'flows'
+            )",
+            &[],
+        )
+        .await?
+        .get(0)
+        .map(|row| row.get::<_, bool>(0))
+        .unwrap_or(false);
+
+    if !tables_exist {
+        info!("Database schema doesn't exist. Creating tables...");
+
+        // Create the keys table
+        client
+            .execute(
+                "CREATE TABLE IF NOT EXISTS public.keys (
+                    key_id SERIAL PRIMARY KEY,
+                    host VARCHAR(255) NOT NULL,
+                    key TEXT NOT NULL,
+                    updated TIMESTAMP WITH TIME ZONE NOT NULL,
+                    CONSTRAINT unique_host_key UNIQUE (host, key)
+                )",
+                &[],
+            )
+            .await?;
+
+        // Create the flows table
+        client
+            .execute(
+                "CREATE TABLE IF NOT EXISTS public.flows (
+                    flow_id SERIAL PRIMARY KEY,
+                    name VARCHAR(255) NOT NULL,
+                    key_id INTEGER NOT NULL,
+                    CONSTRAINT fk_key
+                        FOREIGN KEY(key_id)
+                        REFERENCES public.keys(key_id)
+                        ON DELETE CASCADE,
+                    CONSTRAINT unique_flow_key UNIQUE (name, key_id)
+                )",
+                &[],
+            )
+            .await?;
+
+        // Create an index for faster lookups
+        client
+            .execute(
+                "CREATE INDEX IF NOT EXISTS idx_flows_name ON public.flows(name)",
+                &[],
+            )
+            .await?;
+
+        info!("Database schema created successfully");
+    } else {
+        info!("Database schema already exists");
+    }
+
+    Ok(())
+}
+
+pub async fn batch_insert_keys(
+    client: &Client,
+    keys: &[SshKey],
+) -> Result<KeyInsertStats, tokio_postgres::Error> {
+    if keys.is_empty() {
+        return Ok(KeyInsertStats {
+            total: 0,
+            inserted: 0,
+            updated: 0,
+            unchanged: 0,
+            key_id_map: Vec::new(),
+        });
+    }
+
+    // Подготавливаем массивы для пакетной вставки
+    let mut host_values: Vec<&str> = Vec::with_capacity(keys.len());
+    let mut key_values: Vec<&str> = Vec::with_capacity(keys.len());
+
+    for key in keys {
+        host_values.push(&key.server);
+        key_values.push(&key.public_key);
+    }
+
+    // Сначала проверяем, какие ключи уже существуют в базе
+    let mut existing_keys = HashMap::new();
+    let mut key_query = String::from("SELECT host, key, key_id FROM public.keys WHERE ");
+
+    for i in 0..keys.len() {
+        if i > 0 {
+            key_query.push_str(" OR ");
+        }
+        key_query.push_str(&format!("(host = ${} AND key = ${})", i * 2 + 1, i * 2 + 2));
+    }
+
+    let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
+        Vec::with_capacity(keys.len() * 2);
+    for i in 0..keys.len() {
+        params.push(&host_values[i]);
+        params.push(&key_values[i]);
+    }
+
+    let rows = client.query(&key_query, &params[..]).await?;
+
+    for row in rows {
+        let host: String = row.get(0);
+        let key: String = row.get(1);
+        let key_id: i32 = row.get(2);
+        existing_keys.insert((host, key), key_id);
+    }
+
+    // Определяем, какие ключи нужно вставить, а какие уже существуют
+    let mut keys_to_insert = Vec::new();
+    let mut unchanged_keys = Vec::new();
+
+    for key in keys {
+        let key_tuple = (key.server.clone(), key.public_key.clone());
+        if existing_keys.contains_key(&key_tuple) {
+            unchanged_keys.push((key.clone(), *existing_keys.get(&key_tuple).unwrap()));
+        } else {
+            keys_to_insert.push(key.clone());
+        }
+    }
+
+    let mut inserted_keys = Vec::new();
+
+    // Если есть ключи для вставки, выполняем вставку
+    if !keys_to_insert.is_empty() {
+        let mut insert_sql = String::from("INSERT INTO public.keys (host, key, updated) VALUES ");
+
+        let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new();
+        let mut param_count = 1;
+
+        for (i, key) in keys_to_insert.iter().enumerate() {
+            if i > 0 {
+                insert_sql.push_str(", ");
+            }
+            insert_sql.push_str(&format!("(${}, ${}, NOW())", param_count, param_count + 1));
+            insert_params.push(&key.server);
+            insert_params.push(&key.public_key);
+            param_count += 2;
+        }
+
+        insert_sql.push_str(" RETURNING key_id, host, key");
+
+        let inserted_rows = client.query(&insert_sql, &insert_params[..]).await?;
+
+        for row in inserted_rows {
+            let host: String = row.get(1);
+            let key_text: String = row.get(2);
+            let key_id: i32 = row.get(0);
+
+            if let Some(orig_key) = keys_to_insert
+                .iter()
+                .find(|k| k.server == host && k.public_key == key_text)
+            {
+                inserted_keys.push((orig_key.clone(), key_id));
+            }
+        }
+    }
+
+    // Сохраняем количество элементов перед объединением
+    let inserted_count = inserted_keys.len();
+    let unchanged_count = unchanged_keys.len();
+
+    // Комбинируем результаты и формируем статистику
+    let mut key_id_map = Vec::with_capacity(unchanged_count + inserted_count);
+    key_id_map.extend(unchanged_keys);
+    key_id_map.extend(inserted_keys);
+
+    let stats = KeyInsertStats {
+        total: keys.len(),
+        inserted: inserted_count,
+        updated: 0, // В этой версии мы не обновляем существующие ключи
+        unchanged: unchanged_count,
+        key_id_map,
+    };
+
+    info!(
+        "Keys stats: received={}, new={}, unchanged={}",
+        stats.total, stats.inserted, stats.unchanged
+    );
+
+    Ok(stats)
+}
+
+pub async fn batch_insert_flow_keys(
+    client: &Client,
+    flow_name: &str,
+    key_ids: &[i32],
+) -> Result<usize, tokio_postgres::Error> {
+    if key_ids.is_empty() {
+        info!("No keys to associate with flow '{}'", flow_name);
+        return Ok(0);
+    }
+
+    // Сначала проверим, какие связи уже существуют
+    let mut existing_query =
+        String::from("SELECT key_id FROM public.flows WHERE name = $1 AND key_id IN (");
+
+    for i in 0..key_ids.len() {
+        if i > 0 {
+            existing_query.push_str(", ");
+        }
+        existing_query.push_str(&format!("${}", i + 2));
+    }
+    existing_query.push_str(")");
+
+    let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
+        Vec::with_capacity(key_ids.len() + 1);
+    params.push(&flow_name);
+    for key_id in key_ids {
+        params.push(key_id);
+    }
+
+    let rows = client.query(&existing_query, &params[..]).await?;
+
+    let mut existing_associations = HashSet::new();
+    for row in rows {
+        let key_id: i32 = row.get(0);
+        existing_associations.insert(key_id);
+    }
+
+    // Фильтруем только те ключи, которые еще не связаны с потоком
+    let new_key_ids: Vec<&i32> = key_ids
+        .iter()
+        .filter(|&id| !existing_associations.contains(id))
+        .collect();
+
+    if new_key_ids.is_empty() {
+        info!(
+            "All {} keys are already associated with flow '{}'",
+            key_ids.len(),
+            flow_name
+        );
+        return Ok(0);
+    }
+
+    // Строим SQL запрос с множественными значениями только для новых связей
+    let mut sql = String::from("INSERT INTO public.flows (name, key_id) VALUES ");
+
+    for i in 0..new_key_ids.len() {
+        if i > 0 {
+            sql.push_str(", ");
+        }
+        sql.push_str(&format!("($1, ${})", i + 2));
+    }
+
+    sql.push_str(" ON CONFLICT (name, key_id) DO NOTHING");
+
+    // Подготавливаем параметры для запроса
+    let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
+        Vec::with_capacity(new_key_ids.len() + 1);
+    insert_params.push(&flow_name);
+    for key_id in &new_key_ids {
+        insert_params.push(*key_id);
+    }
+
+    // Выполняем запрос
+    let affected = client.execute(&sql, &insert_params[..]).await?;
+
+    let affected_usize = affected as usize;
+
+    info!(
+        "Added {} new key-flow associations for flow '{}' (skipped {} existing)",
+        affected_usize,
+        flow_name,
+        existing_associations.len()
+    );
+
+    Ok(affected_usize)
+}

src/main.rs

@@ -1,4 +1,5 @@
 mod client;
+mod db;
 mod server;
 
 use clap::Parser;
@@ -19,8 +20,8 @@ use log::{error, info};
     Running in server mode:\n\
     khm --server --ip 0.0.0.0 --port 1337 --db-host psql.psql.svc --db-name khm --db-user admin --db-password <SECRET> --flows work,home\n\
     \n\
-    Running in client mode to send diff and sync ~/.ssh/known_hosts with remote flow in place:\n\
-    khm --host https://khm.example.com/default/keys --known-hosts ~/.ssh/known_hosts --in-place\n\
+    Running in client mode to send diff and sync ~/.ssh/known_hosts with remote flow `work` in place:\n\
+    khm --host https://khm.example.com/work --known-hosts ~/.ssh/known_hosts --in-place\n\
     \n\
     "
 )]
@@ -94,7 +95,7 @@ struct Args {
     #[arg(
         long,
         required_if_eq("server", "false"),
-        help = "Client mode: Full host address of the server to connect to. Like https://khm.example.com/flow_name/keys"
+        help = "Client mode: Full host address of the server to connect to. Like https://khm.example.com/<FLOW_NAME>"
     )]
     host: Option<String>,
 
@@ -105,6 +106,10 @@ struct Args {
         help = "Client mode: Path to the known_hosts file"
     )]
     known_hosts: String,
+
+    /// Basic auth string for client mode. Format: user:pass
+    #[arg(long, default_value = "", help = "Client mode: Basic Auth credentials")]
+    basic_auth: String,
 }
 
 #[actix_web::main]

src/server.rs

@@ -6,6 +6,8 @@ use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
 use tokio_postgres::{Client, NoTls};
 
+use crate::db;
+
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct SshKey {
     pub server: String,
@@ -148,6 +150,8 @@ pub async fn add_keys(
         return HttpResponse::Forbidden().body("Flow ID not allowed");
     }
 
+    // Проверяем формат SSH ключей
+    let mut valid_keys = Vec::new();
     for new_key in new_keys.iter() {
         if !is_valid_ssh_key(&new_key.public_key) {
             error!("Invalid SSH key format for server: {}", new_key.server);
@@ -156,23 +160,48 @@ pub async fn add_keys(
                 new_key.server
             ));
         }
-
-        match insert_key_if_not_exists(&db_client, new_key).await {
-            Ok(key_id) => {
-                if let Err(e) = insert_flow_key(&db_client, &flow_id_str, key_id).await {
-                    error!("Failed to insert flow key into database: {}", e);
-                    return HttpResponse::InternalServerError()
-                        .body("Failed to insert flow key into database");
-                }
-            }
-            Err(e) => {
-                error!("Failed to insert key into database: {}", e);
-                return HttpResponse::InternalServerError()
-                    .body("Failed to insert key into database");
-            }
-        }
+        valid_keys.push(new_key.clone());
     }
 
+    info!(
+        "Processing batch of {} keys for flow: {}",
+        valid_keys.len(),
+        flow_id_str
+    );
+
+    // Батчевая вставка ключей с получением статистики
+    let key_stats = match crate::db::batch_insert_keys(&db_client, &valid_keys).await {
+        Ok(stats) => stats,
+        Err(e) => {
+            error!("Failed to batch insert keys into database: {}", e);
+            return HttpResponse::InternalServerError()
+                .body("Failed to batch insert keys into database");
+        }
+    };
+
+    // Если нет новых ключей, нет необходимости обновлять связи с flow
+    if key_stats.inserted > 0 {
+        // Извлекаем только ID ключей из статистики
+        let key_ids: Vec<i32> = key_stats.key_id_map.iter().map(|(_, id)| *id).collect();
+
+        // Батчевая вставка связей ключей с flow
+        if let Err(e) = crate::db::batch_insert_flow_keys(&db_client, &flow_id_str, &key_ids).await
+        {
+            error!("Failed to batch insert flow keys into database: {}", e);
+            return HttpResponse::InternalServerError()
+                .body("Failed to batch insert flow keys into database");
+        }
+
+        info!(
+            "Added flow associations for {} keys in flow '{}'",
+            key_ids.len(),
+            flow_id_str
+        );
+    } else {
+        info!("No new keys to associate with flow '{}'", flow_id_str);
+    }
+
+    // Получаем обновленные данные
     let updated_flows = match get_keys_from_db(&db_client).await {
         Ok(flows) => flows,
         Err(e) => {
@@ -188,8 +217,22 @@ pub async fn add_keys(
     let updated_flow = flows_guard.iter().find(|flow| flow.name == flow_id_str);
     if let Some(flow) = updated_flow {
         let servers: Vec<&SshKey> = flow.servers.iter().collect();
-        info!("Updated flow: {} with {} keys", flow_id_str, servers.len());
-        HttpResponse::Ok().json(servers)
+        info!(
+            "Keys summary for flow '{}': total received={}, new={}, unchanged={}, total in flow={}",
+            flow_id_str,
+            key_stats.total,
+            key_stats.inserted,
+            key_stats.unchanged,
+            servers.len()
+        );
+
+        // Добавляем статистику в HTTP заголовки ответа
+        let mut response = HttpResponse::Ok();
+        response.append_header(("X-Keys-Total", key_stats.total.to_string()));
+        response.append_header(("X-Keys-New", key_stats.inserted.to_string()));
+        response.append_header(("X-Keys-Unchanged", key_stats.unchanged.to_string()));
+
+        response.json(servers)
     } else {
         error!("Flow ID not found after update: {}", flow_id_str);
         HttpResponse::NotFound().body("Flow ID not found")
@@ -207,7 +250,17 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
         args.db_host, db_user, db_password, args.db_name
     );
 
-    let (db_client, connection) = tokio_postgres::connect(&db_conn_str, NoTls).await.unwrap();
+    info!("Connecting to database at {}", args.db_host);
+    let (db_client, connection) = match tokio_postgres::connect(&db_conn_str, NoTls).await {
+        Ok((client, conn)) => (client, conn),
+        Err(e) => {
+            error!("Failed to connect to the database: {}", e);
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::ConnectionRefused,
+                format!("Database connection error: {}", e),
+            ));
+        }
+    };
     let db_client = Arc::new(db_client);
 
     // Spawn a new thread to run the database connection
@@ -217,6 +270,15 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
         }
     });
 
+    // Initialize database schema if needed
+    if let Err(e) = db::initialize_db_schema(&db_client).await {
+        error!("Failed to initialize database schema: {}", e);
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::Other,
+            format!("Database schema initialization error: {}", e),
+        ));
+    }
+
     let mut initial_flows = match get_keys_from_db(&db_client).await {
         Ok(flows) => flows,
         Err(e) => {