defmodule CirroConnect do use WebSockex alias CirroConnect.MessageRegister, as: Register @moduledoc "Cirro WebSocket-based SQL connector - Copyright Cirro Inc, 2018" @timeout 60_000 @protocol_default "wss://" @query_path "/websockets/query" @doc "Connect to Cirro" def connect(url, user, password) do Register.start() case WebSockex.start_link(finalize_url(url), __MODULE__, :ok, [{:server_name_indication, :disable}]) do {:ok, wsconn} -> finalize_connection(wsconn, user, password) {:error, error} -> {:error, error} end end @doc "Connect to Cirro and return a pipeable connection" def connect!(url, user, password) do {:ok, state} = connect(url, user, password) state end @doc "Execute a rowless SQL statement" def exec(query, {wsconn, authtoken}, options \\ %{}) do dispatch(:execute, wsconn, authtoken, query, options) end @doc "Execute a query, returning status, rows and metadata" def query(query, {wsconn, authtoken}, options \\ %{}) do dispatch(:query, wsconn, authtoken, query, options) end @doc "Execute a query, returning rows and matadata without status" def query!(query, {wsconn, authtoken}, options \\ %{}) do dispatch(:query, wsconn, authtoken, query, options) |> map() end @doc "Fetch the next fetchsize batch of results" def next({wsconn, authtoken}, id) do wssend(wsconn, %{id: id, authtoken: authtoken, command: :next}) await_results(wsconn) end @doc "Cancel a query" def cancel({wsconn, authtoken}, id) do wssend(wsconn, %{id: id, authtoken: authtoken, command: :cancel}) {:ok, {wsconn, authtoken}} end @doc "Fetch the task table" def tasks({wsconn, authtoken}) do dispatch(:tasks, wsconn, authtoken, nil, %{}) |> rows() end @doc "Fetch the connections table" def connections({wsconn, authtoken}) do dispatch(:connections, wsconn, authtoken, nil, %{}) |> rows() end @doc "Forward monitoring events to the given process - options can contain restrictions for session_id and event_type" def monitor(pid, {wsconn, authtoken}, options \\ %{}) do wssend( wsconn, %{id: Register.next_id(), authtoken: authtoken, command: :monitor, options: options}, pid ) end @doc "Close a named connection" def close({wsconn, authtoken}, %{name: name}) do wssend( wsconn, %{ id: Register.next_id(), authtoken: authtoken, command: :close, options: %{ name: name } } ) {:ok, {wsconn, authtoken}} end @doc "Close the connection to Cirro" def close({wsconn, _}) do if is_connected({wsconn, nil}) do WebSockex.send_frame(wsconn, :close) end {:ok, "closed"} end @doc "Handle incorrect close gracefully" def close(nil) do {:error, "invalid connection"} end @doc "Handle abrupt termination of web socket" def terminate(reason, state) do IO.puts("\nCirroConnect WebSocket Terminating:\n#{inspect reason}\n\n#{inspect state}\n") exit(:normal) end @doc "Convert a query's output to a List of Maps of column => value" def map({:ok, results}) do %{"meta" => meta, "rows" => rows} = results colnames = meta |> Enum.map( fn %{"name" => name} -> name |> String.downcase |> String.to_atom end ) rows |> Enum.map( fn row -> Enum.zip(colnames, row) |> Enum.into(%{}) end ) end def map({:error, results}) do {:error, results} end @doc "Return just the rows from a query" def rows({:ok, results}) do %{"rows" => rows} = results rows end def rows({:error, results}) do {:error, results} end @doc "Convert a query's output to a list of lists, where the first row contains the column names" def table({:ok, results}) do %{"meta" => meta, "rows" => rows} = results [Enum.map(meta, fn (x) -> x["name"] end) | rows] end def table({:error, results}) do {:error, results} end @doc "Are we connected?" def is_connected({wsconn, _authtoken}) do case Process.whereis(:cirro_message_register) do nil -> false _ -> Process.alive?(wsconn) end end def is_connected(nil) do false end @doc "Process all monitoring events (example)" def watch_events do receive do {:cirro_monitor, event} -> {:ok, event} |> IO.inspect() end watch_events() end ## ## Inner workingnesses ## defp finalize_url(url) do prefix = if (url =~ "s://"), do: url, else: @protocol_default <> url prefix <> @query_path end defp finalize_connection(wsconn, user, password) do authenticate(wsconn, user, password) receive do {:cirro_connect, %{"error" => true, "message" => error_message}} -> {:error, error_message} {:cirro_connect, %{"error" => false} = response} -> {:ok, {wsconn, response["task"]["authtoken"]}} after @timeout -> {:error, "Timed out waiting for authentication response"} end end defp dispatch(calltype, wsconn, authtoken, query, options) do wssend( wsconn, %{id: Register.next_id(), authtoken: authtoken, command: calltype, statement: to_string(query), options: options} ) await_results(wsconn) end defp await_results(_wsconn) do receive do {:cirro_connect, %{"error" => true, "message" => error_message}} -> {:error, error_message} {:cirro_connect, %{"cancelled" => true, "message" => error_message}} -> {:error, error_message} {:cirro_connect, %{"error" => false} = response} -> {:ok, response} {:error, error} -> {:error, error} {:error} -> {:error, "Unknown error"} after @timeout -> {:error, "Timed out waiting for response"} end end defp authenticate(wsconn, user, password) do wssend( wsconn, %{ id: Register.next_id(), command: "authenticate", options: %{ user: user, password_encrypted: :base64.encode(password) } } ) end defp wssend(wsconn, message, pid \\ self()) do Register.put(message.id, pid) WebSockex.send_frame(wsconn, {:text, Poison.encode! message}) end def handle_frame({:text, text}, state) do response = Poison.decode! text id = response["task"]["id"] case Register.get(id) do {:ok, caller} -> dispatch(id, caller, response) {:ok, state} :error -> {:ok, state} end end def handle_frame(:close, state) do {:close, state} end def handle_disconnect(_, state) do {:ok, state} end defp dispatch(id, caller, response) do case response["task"]["command"] do "monitor" -> # IO.puts( # "MONITOR: " <> inspect( # {:ok, response} # |> rows # ) # ) send(caller, {:cirro_monitor, response}) _ -> Register.delete(id) send(caller, {:cirro_connect, response}) end end end