Use ADBC for Google BigQuery and update Athena code

Author: jonatanklosko

lib/assets/connection_cell/main.js

@@ -691,8 +691,9 @@ export function init(ctx, info) {
           const reader = new FileReader();
 
           reader.onload = (res) => {
-            const value = JSON.parse(res.target.result);
-            ctx.pushEvent("update_field", { field: "credentials", value });
+            // Reformat the JSON into a compact form
+            const value = JSON.stringify(JSON.parse(res.target.result));
+            ctx.pushEvent("update_field", { field: "credentials_json", value });
           };
 
           reader.readAsText(file);

lib/kino_db/connection_cell.ex

@@ -33,7 +33,7 @@ defmodule KinoDB.ConnectionCell do
       "database" => attrs["database"] || "",
       "project_id" => attrs["project_id"] || "",
       "default_dataset_id" => attrs["default_dataset_id"] || "",
-      "credentials" => attrs["credentials"] || %{},
+      "credentials_json" => attrs["credentials_json"] || "",
       "access_key_id" => attrs["access_key_id"] || "",
       "secret_access_key" => secret_access_key,
       "use_secret_access_key_secret" =>
@@ -128,7 +128,7 @@ defmodule KinoDB.ConnectionCell do
           ~w|database_path|
 
         "bigquery" ->
-          ~w|project_id default_dataset_id credentials|
+          ~w|project_id default_dataset_id credentials_json|
 
         "athena" ->
           if fields["use_secret_access_key_secret"],
@@ -267,6 +267,37 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
+  defp to_quoted(%{"type" => "bigquery"} = attrs) do
+    var = quoted_var(attrs["variable"])
+
+    opts =
+      [
+        driver: :bigquery,
+        "adbc.bigquery.sql.project_id": attrs["project_id"]
+      ] ++
+        case attrs["default_dataset_id"] do
+          "" -> []
+          dataset_id -> ["adbc.bigquery.sql.dataset_id": dataset_id]
+        end ++
+        case attrs["credentials_json"] do
+          "" ->
+            []
+
+          credentials_json ->
+            [
+              "adbc.bigquery.sql.auth_type": "adbc.bigquery.sql.auth_type.json_credential_string",
+              "adbc.bigquery.sql.auth_credentials":
+                {:sigil_S, [delimiter: ~s["""]], [{:<<>>, [], [credentials_json <> "\n"]}, []]}
+            ]
+        end
+
+    quote do
+      :ok = Adbc.download_driver!(:bigquery)
+      {:ok, db} = Kino.start_child({Adbc.Database, unquote(opts)})
+      {:ok, unquote(var)} = Kino.start_child({Adbc.Connection, database: db})
+    end
+  end
+
   defp to_quoted(%{"type" => "postgres"} = attrs) do
     quote do
       opts = unquote(trim_opts(shared_options(attrs) ++ postgres_and_mysql_options(attrs)))
@@ -291,40 +322,18 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
-  defp to_quoted(%{"type" => "bigquery"} = attrs) do
-    goth_opts_block = check_bigquery_credentials(attrs)
-
-    conn_block =
-      quote do
-        {:ok, _pid} = Kino.start_child({Goth, opts})
-
-        unquote(quoted_var(attrs["variable"])) =
-          Req.new(http_errors: :raise)
-          |> ReqBigQuery.attach(
-            goth: ReqBigQuery.Goth,
-            project_id: unquote(attrs["project_id"]),
-            default_dataset_id: unquote(attrs["default_dataset_id"])
-          )
-
-        :ok
-      end
-
-    join_quoted([goth_opts_block, conn_block])
-  end
-
   defp to_quoted(%{"type" => "athena"} = attrs) do
     quote do
       unquote(quoted_var(attrs["variable"])) =
-        Req.new(http_errors: :raise)
-        |> ReqAthena.attach(
-          format: :explorer,
+        ReqAthena.new(
           access_key_id: unquote(attrs["access_key_id"]),
           database: unquote(attrs["database"]),
           output_location: unquote(attrs["output_location"]),
           region: unquote(attrs["region"]),
           secret_access_key: unquote(quoted_access_key(attrs)),
           token: unquote(attrs["token"]),
-          workgroup: unquote(attrs["workgroup"])
+          workgroup: unquote(attrs["workgroup"]),
+          http_errors: :raise
         )
 
       :ok
@@ -354,37 +363,6 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
-  defp check_bigquery_credentials(attrs) do
-    case attrs["credentials"] do
-      %{"type" => "service_account"} ->
-        quote do
-          credentials = unquote(Macro.escape(attrs["credentials"]))
-
-          opts = [
-            name: ReqBigQuery.Goth,
-            http_client: &Req.request/1,
-            source: {:service_account, credentials}
-          ]
-        end
-
-      %{"type" => "authorized_user"} ->
-        quote do
-          credentials = unquote(Macro.escape(attrs["credentials"]))
-
-          opts = [
-            name: ReqBigQuery.Goth,
-            http_client: &Req.request/1,
-            source: {:refresh_token, credentials}
-          ]
-        end
-
-      _empty_map ->
-        quote do
-          opts = [name: ReqBigQuery.Goth, http_client: &Req.request/1]
-        end
-    end
-  end
-
   defp shared_options(attrs) do
     opts = [
       hostname: attrs["hostname"],
@@ -516,7 +494,6 @@ defmodule KinoDB.ConnectionCell do
       Code.ensure_loaded?(Postgrex) -> "postgres"
       Code.ensure_loaded?(MyXQL) -> "mysql"
       Code.ensure_loaded?(Exqlite) -> "sqlite"
-      Code.ensure_loaded?(ReqBigQuery) -> "bigquery"
       Code.ensure_loaded?(ReqAthena) -> "athena"
       Code.ensure_loaded?(ReqCH) -> "clickhouse"
       Code.ensure_loaded?(Adbc) -> "duckdb"
@@ -543,20 +520,14 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
-  defp missing_dep(%{"type" => "bigquery"}) do
-    unless Code.ensure_loaded?(ReqBigQuery) do
-      ~s|{:req_bigquery, "~> 0.1"}|
-    end
-  end
-
   defp missing_dep(%{"type" => "athena"}) do
     missing_many_deps([
-      {ReqAthena, ~s|{:req_athena, "~> 0.1"}|},
+      {ReqAthena, ~s|{:req_athena, "~> 0.3"}|},
       {Explorer, ~s|{:explorer, "~> 0.10"}|}
     ])
   end
 
-  defp missing_dep(%{"type" => adbc}) when adbc in ~w[snowflake duckdb] do
+  defp missing_dep(%{"type" => adbc}) when adbc in ~w[snowflake duckdb bigquery] do
     unless Code.ensure_loaded?(Adbc) do
       ~s|{:adbc, "~> 0.3"}|
     end
@@ -585,19 +556,6 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
-  defp join_quoted(quoted_blocks) do
-    asts =
-      Enum.flat_map(quoted_blocks, fn
-        {:__block__, _meta, nodes} -> nodes
-        node -> [node]
-      end)
-
-    case asts do
-      [node] -> node
-      nodes -> {:__block__, [], nodes}
-    end
-  end
-
   defp help_box(%{"type" => "bigquery"}) do
     if Code.ensure_loaded?(Mint.HTTP) do
       if running_on_google_metadata?() do

lib/kino_db/sql_cell.ex

@@ -23,7 +23,7 @@ defmodule KinoDB.SQLCell do
         result_variable: Kino.SmartCell.prefixed_var_name("result", attrs["result_variable"]),
         query: query,
         timeout: attrs["timeout"],
-        cache_query: attrs["cache_query"] || true,
+        cache_query: Map.get(attrs, "cache_query", true),
         data_frame_alias: Explorer.DataFrame,
         missing_dep: missing_dep(connection)
       )
@@ -165,7 +165,6 @@ defmodule KinoDB.SQLCell do
 
   defp connection_type(connection) when is_struct(connection, Req.Request) do
     cond do
-      Keyword.has_key?(connection.request_steps, :bigquery_run) -> "bigquery"
       Keyword.has_key?(connection.request_steps, :athena_run) -> "athena"
       Keyword.has_key?(connection.request_steps, :clickhouse_run) -> "clickhouse"
       true -> nil
@@ -220,55 +219,39 @@ defmodule KinoDB.SQLCell do
     to_quoted(attrs, quote(do: Tds), fn n -> "@#{n}" end)
   end
 
-  # query!/4 based that returns a Req response.
-  defp to_quoted(%{"connection" => %{"type" => "clickhouse"}} = attrs) do
-    to_quoted_query_req(attrs, quote(do: ReqCH), fn n, inner ->
-      name =
-        if String.match?(inner, ~r/[^a-z0-9_]/) do
-          "param_#{n}"
-        else
-          inner
-        end
-
-      "{#{name}:String}"
-    end)
-  end
-
   # Explorer-based
   defp to_quoted(%{"connection" => %{"type" => "snowflake"}} = attrs) do
-    to_explorer_quoted(attrs, fn n -> "?#{n}" end)
+    to_quoted_explorer(attrs, fn n -> "?#{n}" end)
   end
 
   defp to_quoted(%{"connection" => %{"type" => "duckdb"}} = attrs) do
-    to_explorer_quoted(attrs, fn n -> "?#{n}" end)
+    to_quoted_explorer(attrs, fn n -> "?#{n}" end)
   end
 
-  # Req-based
   defp to_quoted(%{"connection" => %{"type" => "bigquery"}} = attrs) do
-    to_req_quoted(attrs, fn _n -> "?" end, :bigquery)
+    to_quoted_explorer(attrs, fn _n -> "?" end)
   end
 
+  # Req-based
   defp to_quoted(%{"connection" => %{"type" => "athena"}} = attrs) do
-    to_req_quoted(attrs, fn _n -> "?" end, :athena)
+    to_quoted_req_query(attrs, quote(do: ReqAthena), fn _n -> "?" end)
   end
 
-  defp to_quoted(_ctx) do
-    quote do
-    end
-  end
+  defp to_quoted(%{"connection" => %{"type" => "clickhouse"}} = attrs) do
+    to_quoted_req_query(attrs, quote(do: ReqCH), fn n, inner ->
+      name =
+        if String.match?(inner, ~r/[^a-z0-9_]/) do
+          "param_#{n}"
+        else
+          inner
+        end
 
-  defp to_quoted_query_req(attrs, quoted_module, next) do
-    {query, params} = parameterize(attrs["query"], attrs["connection"]["type"], next)
-    opts_args = query_opts_args(attrs)
+      "{#{name}:String}"
+    end)
+  end
 
+  defp to_quoted(_ctx) do
     quote do
-      unquote(quoted_var(attrs["result_variable"])) =
-        unquote(quoted_module).query!(
-          unquote(quoted_var(attrs["connection"]["variable"])),
-          unquote(quoted_query(query)),
-          unquote(params),
-          unquote_splicing(opts_args)
-        ).body
     end
   end
 
@@ -287,32 +270,32 @@ defmodule KinoDB.SQLCell do
     end
   end
 
-  defp to_req_quoted(attrs, next, req_key) do
+  defp to_quoted_explorer(attrs, next) do
     {query, params} = parameterize(attrs["query"], attrs["connection"]["type"], next)
-    query = {quoted_query(query), params}
-    opts = query_opts_args(attrs)
-    req_opts = opts |> Enum.at(0, []) |> Keyword.put(req_key, query)
+    data_frame_alias = attrs["data_frame_alias"]
 
     quote do
       unquote(quoted_var(attrs["result_variable"])) =
-        Req.post!(
+        unquote(data_frame_alias).from_query!(
           unquote(quoted_var(attrs["connection"]["variable"])),
-          unquote(req_opts)
-        ).body
+          unquote(quoted_query(query)),
+          unquote(params)
+        )
     end
   end
 
-  defp to_explorer_quoted(attrs, next) do
+  defp to_quoted_req_query(attrs, quoted_module, next) do
     {query, params} = parameterize(attrs["query"], attrs["connection"]["type"], next)
-    data_frame_alias = attrs["data_frame_alias"]
+    opts_args = query_opts_args(attrs)
 
     quote do
       unquote(quoted_var(attrs["result_variable"])) =
-        unquote(data_frame_alias).from_query!(
+        unquote(quoted_module).query!(
           unquote(quoted_var(attrs["connection"]["variable"])),
           unquote(quoted_query(query)),
-          unquote(params)
-        )
+          unquote(params),
+          unquote_splicing(opts_args)
+        ).body
     end
   end
 
@@ -333,8 +316,9 @@ defmodule KinoDB.SQLCell do
        when timeout != nil and type in @connection_types_with_timeout,
        do: [[timeout: timeout * 1000]]
 
-  defp query_opts_args(%{"connection" => %{"type" => "athena"}, "cache_query" => cache_query}),
-    do: [[cache_query: cache_query]]
+  defp query_opts_args(%{"connection" => %{"type" => "athena"}} = attrs) do
+    [[format: :explorer] ++ if(attrs["cache_query"], do: [], else: [cache_query: false])]
+  end
 
   defp query_opts_args(%{"connection" => %{"type" => "clickhouse"}}),
     do: [[format: :explorer]]
@@ -418,7 +402,7 @@ defmodule KinoDB.SQLCell do
     end
   end
 
-  defp missing_dep(%{type: adbc}) when adbc in ~w[snowflake duckdb] do
+  defp missing_dep(%{type: adbc}) when adbc in ~w[snowflake duckdb bigquery] do
     unless Code.ensure_loaded?(Explorer) do
       ~s|{:explorer, "~> 0.10"}|
     end

mix.exs

@@ -37,7 +37,7 @@ defmodule KinoDB.MixProject do
       {:explorer, "~> 0.10", optional: true},
 
       # Those dependecies are new, so we use stricter versions
-      {:req_bigquery, "~> 0.1.0", optional: true},
+      # TODO bump to 0.3.0
       {:req_athena, "~> 0.2.0", optional: true},
       {:req_ch, "~> 0.1.0", optional: true},
 

mix.lock

@@ -30,7 +30,6 @@
   "postgrex": {:hex, :postgrex, "0.18.0", "f34664101eaca11ff24481ed4c378492fed2ff416cd9b06c399e90f321867d7e", [:mix], [{:db_connection, "~> 2.1", [hex: :db_connection, repo: "hexpm", optional: false]}, {:decimal, "~> 1.5 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "a042989ba1bc1cca7383ebb9e461398e3f89f868c92ce6671feb7ef132a252d1"},
   "req": {:hex, :req, "0.5.6", "8fe1eead4a085510fe3d51ad854ca8f20a622aae46e97b302f499dfb84f726ac", [:mix], [{:brotli, "~> 0.3.1", [hex: :brotli, repo: "hexpm", optional: true]}, {:ezstd, "~> 1.0", [hex: :ezstd, repo: "hexpm", optional: true]}, {:finch, "~> 0.17", [hex: :finch, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:mime, "~> 2.0.6 or ~> 2.1", [hex: :mime, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.0", [hex: :nimble_csv, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "cfaa8e720945d46654853de39d368f40362c2641c4b2153c886418914b372185"},
   "req_athena": {:hex, :req_athena, "0.2.0", "a88b1500105f7dfade0c6fdc28174577c72ae6d96f8ab67c5a5cb594f68c8b10", [:mix], [{:aws_credentials, "~> 0.2", [hex: :aws_credentials, repo: "hexpm", optional: true]}, {:aws_signature, "~> 0.3.0", [hex: :aws_signature, repo: "hexpm", optional: false]}, {:explorer, "~> 0.9", [hex: :explorer, repo: "hexpm", optional: true]}, {:req, "~> 0.5.0", [hex: :req, repo: "hexpm", optional: false]}, {:req_s3, "~> 0.2", [hex: :req_s3, repo: "hexpm", optional: false]}], "hexpm", "75376f2f7b6fa924b51cb7d797287a1967589ae7d7d355c50e4170aa7f750d3c"},
-  "req_bigquery": {:hex, :req_bigquery, "0.1.4", "4e1cbadb773d7272ef6c5ed5930f3f2d5df82bc0c6966f00ee71fa681e43a8c2", [:mix], [{:decimal, "~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:goth, "~> 1.3", [hex: :goth, repo: "hexpm", optional: false]}, {:req, "~> 0.3.5 or ~> 0.4", [hex: :req, repo: "hexpm", optional: false]}, {:table, "~> 0.1.1", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "9d099fbaca9eb89909d955ce7ff3dcc731f54e75bf68f99c97bbc2dbeef83ff6"},
   "req_ch": {:hex, :req_ch, "0.1.0", "a8c8710b2fc51ff6bf43d3842bdaff8269688f62ae94250dcfa92d2d37ca4279", [:mix], [{:explorer, "~> 0.10", [hex: :explorer, repo: "hexpm", optional: true]}, {:req, "~> 0.5", [hex: :req, repo: "hexpm", optional: false]}], "hexpm", "7aad3e3f4492647b431294e1fbb915c7280af8628ec6104506029b376b2b4150"},
   "req_s3": {:hex, :req_s3, "0.2.3", "ede5f4c792cf39995379307733ff4593032a876f38da29d9d7ea03881b498b51", [:mix], [{:req, "~> 0.5.6", [hex: :req, repo: "hexpm", optional: false]}], "hexpm", "31b5d52490495c8aeea7e3c5cbcec82f49035e11bdaf41f0e58ab716fefe44ca"},
   "rustler_precompiled": {:hex, :rustler_precompiled, "0.8.2", "5f25cbe220a8fac3e7ad62e6f950fcdca5a5a5f8501835d2823e8c74bf4268d5", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "63d1bd5f8e23096d1ff851839923162096364bac8656a4a3c00d1fff8e83ee0a"},