Ensure prepared statement is freed when lookupPrimaryKeyColumnNames throws

Author: hamzahrmalik

Sources/CassandraClient/Session.swift

@@ -812,18 +812,7 @@ extension CassandraClient {
                 logger.debug("preparing: \(query)")
                 let future = self.underlying.prepare(query: query)
                 return future.asNIOFuture(eventLoop: eventLoop).flatMapThrowing { rawPointer in
-                    let pkColumns: [String]
-                    if let tableName = encryptionTable {
-                        pkColumns = try self.lookupPrimaryKeyColumnNames(tableName: tableName)
-                    } else {
-                        pkColumns = []
-                    }
-                    let prepared = CassandraClient.PreparedStatement(
-                        rawPointer: rawPointer,
-                        encryptionTable: encryptionTable,
-                        primaryKeyColumnNames: pkColumns
-                    )
-                    return prepared
+                    try self.makePreparedStatement(rawPointer: rawPointer, encryptionTable: encryptionTable)
                 }
             }
         }
@@ -892,6 +881,32 @@ extension CassandraClient {
             return names
         }
 
+        /// Builds a `PreparedStatement` taking ownership of `rawPointer`.
+        ///
+        /// If primary-key lookup throws, it does so before the `PreparedStatement` takes
+        /// ownership of the prepared pointer, so we free it here to avoid leaking it.
+        private func makePreparedStatement(
+            rawPointer: OpaquePointer,
+            encryptionTable: String?
+        ) throws -> CassandraClient.PreparedStatement {
+            let pkColumns: [String]
+            if let tableName = encryptionTable {
+                do {
+                    pkColumns = try self.lookupPrimaryKeyColumnNames(tableName: tableName)
+                } catch {
+                    cass_prepared_free(rawPointer)
+                    throw error
+                }
+            } else {
+                pkColumns = []
+            }
+            return CassandraClient.PreparedStatement(
+                rawPointer: rawPointer,
+                encryptionTable: encryptionTable,
+                primaryKeyColumnNames: pkColumns
+            )
+        }
+
         /// Resolve encryption contexts for parameters that have `context: nil` using driver schema metadata.
         /// Discovers PK columns from Cassandra's metadata cache rather than requiring EncryptionSchema registration.
         @available(macOS 15.0, iOS 18.0, visionOS 2.0, *)
@@ -1428,17 +1443,7 @@ extension CassandraClient.Session {
             logger.debug("preparing: \(query)")
             return try await self.underlying.prepare(query: query).await()
         }
-        let pkColumns: [String]
-        if let tableName = encryptionTable {
-            pkColumns = try self.lookupPrimaryKeyColumnNames(tableName: tableName)
-        } else {
-            pkColumns = []
-        }
-        return CassandraClient.PreparedStatement(
-            rawPointer: preparedRawPointer,
-            encryptionTable: encryptionTable,
-            primaryKeyColumnNames: pkColumns
-        )
+        return try self.makePreparedStatement(rawPointer: preparedRawPointer, encryptionTable: encryptionTable)
     }
 
     @available(macOS 12, iOS 15, tvOS 15, watchOS 8, *)