feat: add HTTP keep-alive support and connection statistics to ClickHouse adapter

Author: lohanidamodar

src/Usage/Adapter/ClickHouse.php

@@ -72,6 +72,12 @@ class ClickHouse extends SQL
     /** @var bool Whether to enable gzip compression for HTTP requests/responses */
     private bool $enableCompression = false;
 
+    /** @var bool Whether to enable HTTP keep-alive for connection pooling */
+    private bool $enableKeepAlive = true;
+
+    /** @var int Number of requests made using this adapter instance */
+    private int $requestCount = 0;
+
     /**
      * @param  string  $host  ClickHouse host
      * @param  string  $username  ClickHouse username (default: 'default')
@@ -155,6 +161,34 @@ public function setCompression(bool $enable): self
         return $this;
     }
 
+    /**
+     * Enable or disable HTTP keep-alive for connection pooling.
+     * When enabled, HTTP connections are reused across multiple requests, reducing latency.
+     *
+     * @param bool $enable Whether to enable keep-alive (default: true)
+     * @return self
+     */
+    public function setKeepAlive(bool $enable): self
+    {
+        $this->enableKeepAlive = $enable;
+        return $this;
+    }
+
+    /**
+     * Get connection statistics for monitoring.
+     *
+     * @return array{request_count: int, keep_alive_enabled: bool, compression_enabled: bool, query_logging_enabled: bool}
+     */
+    public function getConnectionStats(): array
+    {
+        return [
+            'request_count' => $this->requestCount,
+            'keep_alive_enabled' => $this->enableKeepAlive,
+            'compression_enabled' => $this->enableCompression,
+            'query_logging_enabled' => $this->enableQueryLogging,
+        ];
+    }
+
     /**
      * Get the query execution log.
      *
@@ -506,11 +540,21 @@ private function query(string $sql, array $params = []): string
         // Update the database header for each query (in case setDatabase was called)
         $this->client->addHeader('X-ClickHouse-Database', $this->database);
 
+        // Enable keep-alive for connection pooling
+        if ($this->enableKeepAlive) {
+            $this->client->addHeader('Connection', 'keep-alive');
+        } else {
+            $this->client->addHeader('Connection', 'close');
+        }
+
         // Enable compression if configured
         if ($this->enableCompression) {
             $this->client->addHeader('Accept-Encoding', 'gzip');
         }
 
+        // Track request count for statistics
+        $this->requestCount++;
+
         // Build multipart form data body with query and parameters
         // The Fetch client will automatically encode arrays as multipart/form-data
         $body = ['query' => $sql];
@@ -575,11 +619,21 @@ private function insert(string $table, array $data): void
         $this->client->addHeader('X-ClickHouse-Database', $this->database);
         $this->client->addHeader('Content-Type', 'application/x-ndjson');
 
+        // Enable keep-alive for connection pooling
+        if ($this->enableKeepAlive) {
+            $this->client->addHeader('Connection', 'keep-alive');
+        } else {
+            $this->client->addHeader('Connection', 'close');
+        }
+
         // Enable compression if configured
         if ($this->enableCompression) {
             $this->client->addHeader('Accept-Encoding', 'gzip');
         }
 
+        // Track request count for statistics
+        $this->requestCount++;
+
         // Join JSON strings with newlines
         $body = implode("\n", $data);
 

tests/Usage/Adapter/ClickHouseTest.php

@@ -594,4 +594,60 @@ public function testCompression(): void
         $sum = $usage->sumByPeriod('compression.test.batch', '1h');
         $this->assertIsInt($sum);
     }
+
+    /**
+     * Test connection pooling functionality
+     */
+    public function testConnectionPooling(): void
+    {
+        // Create a new adapter instance
+        $host = getenv('CLICKHOUSE_HOST') ?: 'clickhouse';
+        $username = getenv('CLICKHOUSE_USER') ?: 'default';
+        $password = getenv('CLICKHOUSE_PASSWORD') ?: 'clickhouse';
+        $port = (int) (getenv('CLICKHOUSE_PORT') ?: 8123);
+        $secure = (bool) (getenv('CLICKHOUSE_SECURE') ?: false);
+
+        $adapter = new ClickHouseAdapter($host, $username, $password, $port, $secure);
+        $adapter->setNamespace('utopia_usage_pooling_test');
+        $adapter->setTenant(1);
+
+        if ($database = getenv('CLICKHOUSE_DATABASE')) {
+            $adapter->setDatabase($database);
+        }
+
+        $usage = new Usage($adapter);
+        $usage->setup();
+
+        // Test enabling keep-alive
+        $result = $adapter->setKeepAlive(true);
+        $this->assertInstanceOf(ClickHouseAdapter::class, $result);
+
+        // Test disabling keep-alive
+        $result = $adapter->setKeepAlive(false);
+        $this->assertInstanceOf(ClickHouseAdapter::class, $result);
+
+        // Re-enable for testing
+        $adapter->setKeepAlive(true);
+
+        // Get initial stats
+        $stats = $adapter->getConnectionStats();
+        $this->assertIsArray($stats);
+        $this->assertArrayHasKey('request_count', $stats);
+        $this->assertArrayHasKey('keep_alive_enabled', $stats);
+        $this->assertArrayHasKey('compression_enabled', $stats);
+        $this->assertArrayHasKey('query_logging_enabled', $stats);
+        $this->assertTrue($stats['keep_alive_enabled']);
+
+        $initialCount = $stats['request_count'];
+
+        // Make some requests
+        $usage->log('pooling.test', 100, '1h', ['test' => 'value']);
+        $usage->find([]);
+        $usage->count([]);
+
+        // Verify request count increased
+        $newStats = $adapter->getConnectionStats();
+        $this->assertGreaterThan($initialCount, $newStats['request_count']);
+        $this->assertGreaterThanOrEqual(3, $newStats['request_count'] - $initialCount);
+    }
 }