     {:ok, state}
   end
 
-  def handle_prepare(%Query{} = query, _opts, %{sock: sock} = state) do
+  def handle_prepare(%Query{} = query, _opts, %__MODULE__{socket: socket} = state) do
     payload =
       Frame.new(:prepare)
       |> Protocol.encode_request(query)
       |> Frame.encode()
 
-    with :ok <- :gen_tcp.send(sock, payload),
-         {:ok, %Frame{} = frame} = Utils.recv_frame(sock) do
+    with :ok <- :gen_tcp.send(socket, payload),
+         {:ok, %Frame{} = frame} = Utils.recv_frame(socket) do
       {:ok, Protocol.decode_response(frame, query), state}
     else
       {:error, reason} ->

    {:ok, query, state}
  end

  def disconnect(_exception, %{sock: sock}) do
    :ok = :gen_tcp.close(sock)
  end

  defp connect(host, port) do
    with {:error, reason} <- :gen_tcp.connect(host, port, @default_sock_opts, @default_timeout),
         do: {:error, Error.new("TCP connect", reason)}
  end
end

  alias __MODULE__.{Error, Utils}

  @default_timeout 5_000
  @default_socket_opts [packet: :raw, mode: :binary, active: false]

  defstruct [:socket]

  def connect(opts) do
    host = Keyword.fetch!(opts, :host) |> to_char_list()
    port = Keyword.get(opts, :port, 9042)

    with {:ok, socket} <- connect(host, port),
         {:ok, options} <- Utils.request_options(socket),
         :ok <- Utils.startup_connection(socket, options),
         do: {:ok, %__MODULE__{socket: socket}}
  end

  def checkout(state) do
