diff --git a/lib/assets/connection_cell/main.js b/lib/assets/connection_cell/main.js index 3df194f..fe89f08 100644 --- a/lib/assets/connection_cell/main.js +++ b/lib/assets/connection_cell/main.js @@ -198,6 +198,73 @@ export function init(ctx, info) { ` }; + const AthenaForm = { + name: "AthenaForm", + + components: { + BaseInput: BaseInput + }, + + props: { + fields: { + type: Object, + default: {} + }, + }, + + template: ` +
+ + + +
+
+ + +
+ ` + }; + const BigQueryForm = { name: "BigQueryForm", @@ -291,7 +358,8 @@ export function init(ctx, info) { BaseSelect: BaseSelect, SQLiteForm: SQLiteForm, DefaultSQLForm: DefaultSQLForm, - BigQueryForm: BigQueryForm + BigQueryForm: BigQueryForm, + AthenaForm: AthenaForm }, template: ` @@ -328,6 +396,7 @@ export function init(ctx, info) { + @@ -342,7 +411,8 @@ export function init(ctx, info) { {label: "PostgreSQL", value: "postgres"}, {label: "MySQL", value: "mysql"}, {label: "SQLite", value: "sqlite"}, - {label: "Google BigQuery", value: "bigquery"} + {label: "Google BigQuery", value: "bigquery"}, + {label: "AWS Athena", value: "athena"} ] } }, @@ -356,6 +426,10 @@ export function init(ctx, info) { return this.fields.type === "bigquery"; }, + isAthena() { + return this.fields.type === "athena"; + }, + isDefaultDatabase() { return ["postgres", "mysql"].includes(this.fields.type); } diff --git a/lib/assets/sql_cell/main.js b/lib/assets/sql_cell/main.js index ced3514..90c0728 100644 --- a/lib/assets/sql_cell/main.js +++ b/lib/assets/sql_cell/main.js @@ -127,7 +127,8 @@ export function init(ctx, payload) { postgres: "PostgreSQL", mysql: "MySQL", sqlite: "SQLite", - bigquery: "Google BigQuery" + bigquery: "Google BigQuery", + athena: "AWS Athena" }; connectionEl.innerHTML = connections.map((connection) => ` diff --git a/lib/kino_db.ex b/lib/kino_db.ex index 17f86e3..f6ca595 100644 --- a/lib/kino_db.ex +++ b/lib/kino_db.ex @@ -1,4 +1,4 @@ -results = [Postgrex.Result, MyXQL.Result, Exqlite.Result, ReqBigQuery.Result] +results = [Postgrex.Result, MyXQL.Result, Exqlite.Result, ReqBigQuery.Result, ReqAthena.Result] for mod <- results do defimpl Kino.Render, for: mod do diff --git a/lib/kino_db/connection_cell.ex b/lib/kino_db/connection_cell.ex index 7dab03f..50b604e 100644 --- a/lib/kino_db/connection_cell.ex +++ b/lib/kino_db/connection_cell.ex @@ -25,7 +25,11 @@ defmodule KinoDB.ConnectionCell do "database" => attrs["database"] || "", "project_id" => attrs["project_id"] || "", "default_dataset_id" => attrs["default_dataset_id"] || "", - "credentials" => attrs["credentials"] || %{} + "credentials" => attrs["credentials"] || %{}, + "access_key_id" => attrs["access_key_id"] || "", + "secret_access_key" => attrs["secret_access_key"] || "", + "region" => attrs["region"] || "", + "output_location" => attrs["output_location"] || "" } {:ok, assign(ctx, fields: fields, missing_dep: missing_dep(fields))} @@ -97,6 +101,9 @@ defmodule KinoDB.ConnectionCell do "bigquery" -> ~w|project_id default_dataset_id credentials| + "athena" -> + ~w|access_key_id secret_access_key region output_location database| + type when type in ["postgres", "mysql"] -> ~w|database hostname port username password| end @@ -185,6 +192,22 @@ defmodule KinoDB.ConnectionCell do end end + defp to_quoted(%{"type" => "athena"} = attrs) do + quote do + unquote(quoted_var(attrs["variable"])) = + Req.new(http_errors: :raise) + |> ReqAthena.attach( + access_key_id: unquote(attrs["access_key_id"]), + secret_access_key: unquote(attrs["secret_access_key"]), + region: unquote(attrs["region"]), + database: unquote(attrs["database"]), + output_location: unquote(attrs["output_location"]) + ) + + :ok + end + end + defp shared_options(attrs) do quote do [ @@ -205,6 +228,7 @@ defmodule KinoDB.ConnectionCell do Code.ensure_loaded?(MyXQL) -> "mysql" Code.ensure_loaded?(Exqlite) -> "sqlite" Code.ensure_loaded?(ReqBigQuery) -> "bigquery" + Code.ensure_loaded?(ReqAthena) -> "athena" true -> "postgres" end end @@ -233,6 +257,12 @@ defmodule KinoDB.ConnectionCell do end end + defp missing_dep(%{"type" => "athena"}) do + unless Code.ensure_loaded?(ReqAthena) do + ~s|{:req_athena, github: "livebook-dev/req_athena"}| + end + end + defp missing_dep(_ctx), do: nil defp join_quoted(quoted_blocks) do diff --git a/lib/kino_db/sql_cell.ex b/lib/kino_db/sql_cell.ex index 5bbc287..465a844 100644 --- a/lib/kino_db/sql_cell.ex +++ b/lib/kino_db/sql_cell.ex @@ -121,6 +121,7 @@ defmodule KinoDB.SQLCell do defp connection_type(%{request_steps: request_steps}) do cond do Keyword.has_key?(request_steps, :bigquery_run) -> "bigquery" + Keyword.has_key?(request_steps, :athena_run) -> "athena" true -> nil end end @@ -157,18 +158,11 @@ defmodule KinoDB.SQLCell do end defp to_quoted(%{"connection" => %{"type" => "bigquery"}} = attrs) do - {query, params} = parameterize(attrs["query"], fn _n -> "?" end) - bigquery = {quoted_query(query), params} - opts = query_opts_args(attrs) - req_opts = opts |> Enum.at(0, []) |> Keyword.put(:bigquery, bigquery) + to_req_quoted(attrs, fn _n -> "?" end, :bigquery) + end - quote do - unquote(quoted_var(attrs["result_variable"])) = - Req.post!( - unquote(quoted_var(attrs["connection"]["variable"])), - unquote(req_opts) - ).body - end + defp to_quoted(%{"connection" => %{"type" => "athena"}} = attrs) do + to_req_quoted(attrs, fn _n -> "?" end, :athena) end defp to_quoted(_ctx) do @@ -191,6 +185,21 @@ defmodule KinoDB.SQLCell do end end + defp to_req_quoted(attrs, next, req_key) do + {query, params} = parameterize(attrs["query"], next) + query = {quoted_query(query), params} + opts = query_opts_args(attrs) + req_opts = opts |> Enum.at(0, []) |> Keyword.put(req_key, query) + + quote do + unquote(quoted_var(attrs["result_variable"])) = + Req.post!( + unquote(quoted_var(attrs["connection"]["variable"])), + unquote(req_opts) + ).body + end + end + defp quoted_var(nil), do: nil defp quoted_var(string), do: {String.to_atom(string), [], nil} diff --git a/mix.exs b/mix.exs index eea9f49..ca9c316 100644 --- a/mix.exs +++ b/mix.exs @@ -33,6 +33,7 @@ defmodule KinoDB.MixProject do {:myxql, "~> 0.6.2 or ~> 0.7", optional: true}, {:db_connection, "~> 2.4.2", optional: true}, {:req_bigquery, github: "livebook-dev/req_bigquery", optional: true}, + {:req_athena, github: "livebook-dev/req_athena", optional: true}, {:ex_doc, "~> 0.28", only: :dev, runtime: false} ] end diff --git a/mix.lock b/mix.lock index a5d6276..2b033c6 100644 --- a/mix.lock +++ b/mix.lock @@ -1,4 +1,5 @@ %{ + "aws_signature": {:hex, :aws_signature, "0.3.1", "67f369094cbd55ffa2bbd8cc713ede14b195fcfb45c86665cd7c5ad010276148", [:rebar3], [], "hexpm", "50fc4dc1d1f7c2d0a8c63f455b3c66ecd74c1cf4c915c768a636f9227704a674"}, "castore": {:hex, :castore, "0.1.17", "ba672681de4e51ed8ec1f74ed624d104c0db72742ea1a5e74edbc770c815182f", [:mix], [], "hexpm", "d9844227ed52d26e7519224525cb6868650c272d4a3d327ce3ca5570c12163f9"}, "connection": {:hex, :connection, "1.1.0", "ff2a49c4b75b6fb3e674bfc5536451607270aac754ffd1bdfe175abe4a6d7a68", [:mix], [], "hexpm", "722c1eb0a418fbe91ba7bd59a47e28008a189d47e37e0e7bb85585a016b2869c"}, "db_connection": {:hex, :db_connection, "2.4.2", "f92e79aff2375299a16bcb069a14ee8615c3414863a6fef93156aee8e86c2ff3", [:mix], [{:connection, "~> 1.0", [hex: :connection, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "4fe53ca91b99f55ea249693a0229356a08f4d1a7931d8ffa79289b145fe83668"}, diff --git a/test/kino_db/connection_cell_test.exs b/test/kino_db/connection_cell_test.exs index fd210dd..1bc766c 100644 --- a/test/kino_db/connection_cell_test.exs +++ b/test/kino_db/connection_cell_test.exs @@ -152,6 +152,35 @@ defmodule KinoDB.ConnectionCellTest do :ok\ """ end + + test "restores source code from attrs with Athena" do + attrs = %{ + "variable" => "db", + "type" => "athena", + "access_key_id" => "id", + "secret_access_key" => "secret", + "region" => "region", + "database" => "default", + "output_location" => "s3://my-bucket" + } + + {_kino, source} = start_smart_cell!(ConnectionCell, attrs) + + assert source == + """ + db = + Req.new(http_errors: :raise) + |> ReqAthena.attach( + access_key_id: "id", + secret_access_key: "secret", + region: "region", + database: "default", + output_location: "s3://my-bucket" + ) + + :ok\ + """ + end end test "when a field changes, broadcasts the change and sends source update" do diff --git a/test/kino_db/sql_cell_test.exs b/test/kino_db/sql_cell_test.exs index ca9a874..3a17af0 100644 --- a/test/kino_db/sql_cell_test.exs +++ b/test/kino_db/sql_cell_test.exs @@ -86,6 +86,10 @@ defmodule KinoDB.SQLCellTest do assert SQLCell.to_source(put_in(attrs["connection"]["type"], "bigquery")) == """ result = Req.post!(conn, bigquery: {\"SELECT id FROM users\", []}).body\ """ + + assert SQLCell.to_source(put_in(attrs["connection"]["type"], "athena")) == """ + result = Req.post!(conn, athena: {\"SELECT id FROM users\", []}).body\ + """ end test "uses heredoc string for a multi-line query" do @@ -142,6 +146,17 @@ defmodule KinoDB.SQLCellTest do """, []} ).body\ ''' + + assert SQLCell.to_source(put_in(attrs["connection"]["type"], "athena")) == ~s''' + result = + Req.post!(conn, + athena: + {""" + SELECT id FROM users + WHERE last_name = 'Sherlock' + """, []} + ).body\ + ''' end test "parses parameter expressions" do @@ -183,6 +198,13 @@ defmodule KinoDB.SQLCellTest do {"SELECT id FROM users WHERE id ? AND name LIKE ?", [user_id, search <> "%"]} ).body\ ''' + + assert SQLCell.to_source(put_in(attrs["connection"]["type"], "athena")) == ~s''' + result = + Req.post!(conn, + athena: {"SELECT id FROM users WHERE id ? AND name LIKE ?", [user_id, search <> "%"]} + ).body\ + ''' end test "ignores parameters inside comments" do @@ -247,6 +269,18 @@ defmodule KinoDB.SQLCellTest do """, [user_id3]} ).body\ ''' + + assert SQLCell.to_source(put_in(attrs["connection"]["type"], "athena")) == ~s''' + result = + Req.post!(conn, + athena: + {""" + SELECT id from users + -- WHERE id = {{user_id1}} + /* WHERE id = {{user_id2}} */ WHERE id = ? + """, [user_id3]} + ).body\ + ''' end test "passes timeout option when a timeout is specified" do @@ -272,6 +306,10 @@ defmodule KinoDB.SQLCellTest do assert SQLCell.to_source(put_in(attrs["connection"]["type"], "bigquery")) == """ result = Req.post!(conn, bigquery: {"SELECT id FROM users", []}, timeout: 30000).body\ """ + + assert SQLCell.to_source(put_in(attrs["connection"]["type"], "athena")) == """ + result = Req.post!(conn, athena: {"SELECT id FROM users", []}, timeout: 30000).body\ + """ end end end