databricks_sql 0.1.0

data/spec/databricks_sql/client_spec.rb

@@ -0,0 +1,533 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe DatabricksSql::Client do
+  before do
+    DatabricksSql.reset_configuration!
+  end
+
+  let(:host) { "https://dbc.example.com" }
+  let(:token) { "dapi-token" }
+  let(:warehouse_id) { "warehouse-123" }
+  let(:client) { described_class.new(host: host, token: token, warehouse_id: warehouse_id) }
+  let(:statement) { "SELECT id, name FROM users" }
+
+  describe "#initialize" do
+    it "raises when host does not use HTTPS" do
+      expect do
+        described_class.new(host: "http://dbc.example.com", token: token, warehouse_id: warehouse_id)
+      end.to raise_error(DatabricksSql::ConfigurationError, /HTTPS/)
+    end
+
+    it "raises when warehouse_id is missing" do
+      expect do
+        described_class.new(host: host, token: token, warehouse_id: "")
+      end.to raise_error(DatabricksSql::ConfigurationError)
+    end
+
+    it "uses global configured values when args are omitted" do
+      DatabricksSql.configure do |config|
+        config.host = host
+        config.token = token
+        config.warehouse_id = warehouse_id
+      end
+
+      configured_client = described_class.new
+      expect(configured_client.host).to eq(host)
+      expect(configured_client.warehouse_id).to eq(warehouse_id)
+    end
+
+    it "redacts token in inspect output" do
+      inspected = client.inspect
+
+      expect(inspected).to include("[REDACTED]")
+      expect(inspected).not_to include(token)
+    end
+  end
+
+  describe "#execute_statement" do
+    it "sends official optional parameters in statement payload" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements")
+        .with do |request|
+          payload = JSON.parse(request.body)
+          payload["byte_limit"] == 1000 &&
+            payload["row_limit"] == 2 &&
+            payload["query_tags"] == [{ "key" => "team", "value" => "finance" }]
+        end
+        .to_return(
+          status: 200,
+          headers: { "Content-Type" => "application/json" },
+          body: JSON.generate(
+            {
+              "statement_id" => "stmt-params",
+              "status" => { "state" => "SUCCEEDED" },
+              "result" => {
+                "disposition" => "INLINE",
+                "schema" => { "columns" => [{ "name" => "id" }] },
+                "data_array" => [[1]]
+              }
+            }
+          )
+        )
+
+      result = client.execute_statement(
+        statement: statement,
+        byte_limit: 1000,
+        row_limit: 2,
+        query_tags: [{ key: "team", value: "finance" }]
+      )
+
+      expect(result.rows).to eq([{ "id" => 1 }])
+    end
+
+    it "automatically fetches remaining INLINE chunks" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-multi-inline",
+            "status" => { "state" => "SUCCEEDED" },
+            "manifest" => {
+              "schema" => { "columns" => [{ "name" => "id" }] }
+            },
+            "result" => {
+              "chunk_index" => 0,
+              "data_array" => [[1], [2]],
+              "next_chunk_internal_link" => "/api/2.0/sql/statements/stmt-multi-inline/result/chunks/1?row_offset=2"
+            }
+          }
+        )
+      )
+
+      stub_request(:get, "#{host}/api/2.0/sql/statements/stmt-multi-inline/result/chunks/1?row_offset=2").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "chunk_index" => 1,
+            "data_array" => [[3], [4]]
+          }
+        )
+      )
+
+      result = client.execute_statement(statement: statement)
+
+      expect(result.rows).to eq([
+                                  { "id" => 1 },
+                                  { "id" => 2 },
+                                  { "id" => 3 },
+                                  { "id" => 4 }
+                                ])
+      expect(result.next_chunk_internal_link).to be_nil
+    end
+
+    it "allows manual next chunk retrieval" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-manual",
+            "status" => { "state" => "SUCCEEDED" },
+            "manifest" => {
+              "schema" => { "columns" => [{ "name" => "id" }] }
+            },
+            "result" => {
+              "chunk_index" => 0,
+              "data_array" => [[1]],
+              "next_chunk_internal_link" => "/api/2.0/sql/statements/stmt-manual/result/chunks/1?row_offset=1"
+            }
+          }
+        )
+      )
+
+      stub_request(:get, "#{host}/api/2.0/sql/statements/stmt-manual/result/chunks/1?row_offset=1").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "chunk_index" => 1,
+            "data_array" => [[2]]
+          }
+        )
+      )
+
+      result = client.execute_statement(statement: statement, auto_fetch_chunks: false)
+      expect(result.rows).to eq([{ "id" => 1 }])
+      expect(result.next_chunk_internal_link).to eq(
+        "/api/2.0/sql/statements/stmt-manual/result/chunks/1?row_offset=1"
+      )
+
+      next_chunk = client.fetch_next_chunk(
+        statement_id: result.statement_id,
+        next_chunk_internal_link: result.next_chunk_internal_link,
+        columns: result.columns,
+        disposition: result.disposition
+      )
+
+      expect(next_chunk.rows).to eq([{ "id" => 2 }])
+      expect(next_chunk.next_chunk_internal_link).to be_nil
+    end
+
+    it "fetches statement details when submit response succeeds without result payload" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-doc-shape",
+            "status" => { "state" => "SUCCEEDED" }
+          }
+        )
+      )
+
+      stub_request(:get, "#{host}/api/2.0/sql/statements/stmt-doc-shape").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-doc-shape",
+            "status" => { "state" => "SUCCEEDED" },
+            "manifest" => {
+              "format" => "JSON_ARRAY",
+              "schema" => {
+                "columns" => [
+                  { "name" => "l_orderkey", "position" => 0, "type_name" => "LONG", "type_text" => "BIGINT" },
+                  {
+                    "name" => "l_extendedprice",
+                    "position" => 1,
+                    "type_name" => "DECIMAL",
+                    "type_precision" => 18,
+                    "type_scale" => 2,
+                    "type_text" => "DECIMAL(18,2)"
+                  },
+                  { "name" => "l_shipdate", "position" => 2, "type_name" => "DATE", "type_text" => "DATE" }
+                ]
+              }
+            },
+            "result" => {
+              "chunk_index" => 0,
+              "data_array" => [
+                ["2", "71433.16", "1997-01-28"],
+                ["7", "86152.02", "1996-01-15"]
+              ],
+              "row_count" => 2,
+              "row_offset" => 0
+            }
+          }
+        )
+      )
+
+      result = client.execute_statement(statement: statement)
+
+      expect(result.rows).to eq([
+                                  {
+                                    "l_orderkey" => "2",
+                                    "l_extendedprice" => "71433.16",
+                                    "l_shipdate" => "1997-01-28"
+                                  },
+                                  {
+                                    "l_orderkey" => "7",
+                                    "l_extendedprice" => "86152.02",
+                                    "l_shipdate" => "1996-01-15"
+                                  }
+                                ])
+    end
+
+    it "returns INLINE rows" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-inline",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "INLINE",
+              "schema" => { "columns" => [{ "name" => "id" }, { "name" => "name" }] },
+              "data_array" => [[1, "Ana"], [2, "Bia"]]
+            }
+          }
+        )
+      )
+
+      result = client.execute_statement(statement: statement)
+
+      expect(result).to be_a(DatabricksSql::Result)
+      expect(result.success?).to be(true)
+      expect(result.rows).to eq([
+                                  { "id" => 1, "name" => "Ana" },
+                                  { "id" => 2, "name" => "Bia" }
+                                ])
+    end
+
+    it "downloads and parses EXTERNAL_LINK results" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-external",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "EXTERNAL_LINK",
+              "external_link" => "https://files.example.com/result.csv"
+            }
+          }
+        )
+      )
+
+      stub_request(:get, "https://files.example.com/result.csv").to_return(
+        status: 200,
+        headers: { "Content-Type" => "text/csv" },
+        body: "id,name\n1,Ana\n2,Bia\n"
+      )
+
+      result = client.execute_statement(statement: statement, disposition: "EXTERNAL_LINK")
+
+      expect(result.rows).to eq([
+                                  { "id" => "1", "name" => "Ana" },
+                                  { "id" => "2", "name" => "Bia" }
+                                ])
+    end
+
+    it "downloads and parses EXTERNAL_LINKS results" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-external-links",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "EXTERNAL_LINKS",
+              "external_link" => "https://files.example.com/result-links.csv"
+            }
+          }
+        )
+      )
+
+      stub_request(:get, "https://files.example.com/result-links.csv").to_return(
+        status: 200,
+        headers: { "Content-Type" => "text/csv" },
+        body: "id,name\n3,Caio\n"
+      )
+
+      result = client.execute_statement(statement: statement, disposition: "EXTERNAL_LINKS")
+
+      expect(result.rows).to eq([
+                                  { "id" => "3", "name" => "Caio" }
+                                ])
+    end
+
+    it "automatically fetches remaining EXTERNAL_LINKS chunks" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-external-multi",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "EXTERNAL_LINKS",
+              "external_links" => [
+                {
+                  "external_link" => "https://files.example.com/chunk0.csv",
+                  "next_chunk_internal_link" => "/api/2.0/sql/statements/stmt-external-multi/result/chunks/1"
+                }
+              ]
+            }
+          }
+        )
+      )
+
+      stub_request(:get, "https://files.example.com/chunk0.csv").to_return(
+        status: 200,
+        headers: { "Content-Type" => "text/csv" },
+        body: "id,name\n1,Ana\n"
+      )
+
+      stub_request(:get, "#{host}/api/2.0/sql/statements/stmt-external-multi/result/chunks/1").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "external_links" => [
+              {
+                "external_link" => "https://files.example.com/chunk1.csv"
+              }
+            ]
+          }
+        )
+      )
+
+      stub_request(:get, "https://files.example.com/chunk1.csv").to_return(
+        status: 200,
+        headers: { "Content-Type" => "text/csv" },
+        body: "id,name\n2,Bia\n"
+      )
+
+      result = client.execute_statement(statement: statement, disposition: "EXTERNAL_LINKS")
+
+      expect(result.rows).to eq([
+                                  { "id" => "1", "name" => "Ana" },
+                                  { "id" => "2", "name" => "Bia" }
+                                ])
+      expect(result.next_chunk_internal_link).to be_nil
+    end
+
+    it "rejects EXTERNAL_LINK over HTTP by default" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-external-http",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "EXTERNAL_LINK",
+              "external_link" => "http://files.example.com/result.csv"
+            }
+          }
+        )
+      )
+
+      expect do
+        client.execute_statement(statement: statement, disposition: "EXTERNAL_LINK")
+      end.to raise_error(DatabricksSql::ConfigurationError, /HTTPS/)
+    end
+
+    it "rejects EXTERNAL_LINK host not in allowlist" do
+      DatabricksSql.configure do |config|
+        config.host = host
+        config.token = token
+        config.warehouse_id = warehouse_id
+        config.external_link_allowed_hosts = ["trusted.example.com"]
+      end
+
+      secure_client = described_class.new
+
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-external-blocked",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "EXTERNAL_LINK",
+              "external_link" => "https://files.example.com/result.csv"
+            }
+          }
+        )
+      )
+
+      expect do
+        secure_client.execute_statement(statement: statement, disposition: "EXTERNAL_LINK")
+      end.to raise_error(DatabricksSql::ConfigurationError, /allowed list/)
+    end
+
+    it "allows EXTERNAL_LINK host in allowlist" do
+      DatabricksSql.configure do |config|
+        config.host = host
+        config.token = token
+        config.warehouse_id = warehouse_id
+        config.external_link_allowed_hosts = ["files.example.com"]
+      end
+
+      secure_client = described_class.new
+
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-external-allowed",
+            "status" => { "state" => "SUCCEEDED" },
+            "result" => {
+              "disposition" => "EXTERNAL_LINK",
+              "external_link" => "https://files.example.com/result.csv"
+            }
+          }
+        )
+      )
+
+      stub_request(:get, "https://files.example.com/result.csv").to_return(
+        status: 200,
+        headers: { "Content-Type" => "text/csv" },
+        body: "id,name\n1,Ana\n"
+      )
+
+      result = secure_client.execute_statement(statement: statement, disposition: "EXTERNAL_LINK")
+      expect(result.rows).to eq([{ "id" => "1", "name" => "Ana" }])
+    end
+
+    it "raises mapped error for non-2xx" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(status: 401, body: "unauthorized")
+
+      expect do
+        client.execute_statement(statement: statement)
+      end.to raise_error(DatabricksSql::AuthenticationError)
+    end
+  end
+
+  describe "asynchronous polling" do
+    it "submits async and waits until terminal state" do
+      stub_request(:post, "#{host}/api/2.0/sql/statements").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate(
+          {
+            "statement_id" => "stmt-async",
+            "status" => { "state" => "PENDING" }
+          }
+        )
+      )
+
+      stub_request(:get, "#{host}/api/2.0/sql/statements/stmt-async")
+        .to_return(
+          {
+            status: 200,
+            headers: { "Content-Type" => "application/json" },
+            body: JSON.generate({ "statement_id" => "stmt-async", "status" => { "state" => "RUNNING" } })
+          },
+          {
+            status: 200,
+            headers: { "Content-Type" => "application/json" },
+            body: JSON.generate(
+              {
+                "statement_id" => "stmt-async",
+                "status" => { "state" => "SUCCEEDED" },
+                "result" => {
+                  "disposition" => "INLINE",
+                  "schema" => { "columns" => [{ "name" => "id" }] },
+                  "data_array" => [[7]]
+                }
+              }
+            )
+          }
+        )
+
+      submission = client.execute_statement_async(statement: statement)
+      expect(submission.fetch("statement_id")).to eq("stmt-async")
+
+      result = client.wait_for_statement(statement_id: "stmt-async", max_wait: 5, poll_interval: 0)
+      expect(result.status).to eq("SUCCEEDED")
+      expect(result.rows).to eq([{ "id" => 7 }])
+    end
+
+    it "raises timeout error when polling exceeds max_wait" do
+      stub_request(:get, "#{host}/api/2.0/sql/statements/stmt-timeout").to_return(
+        status: 200,
+        headers: { "Content-Type" => "application/json" },
+        body: JSON.generate({ "statement_id" => "stmt-timeout", "status" => { "state" => "RUNNING" } })
+      )
+
+      expect do
+        client.wait_for_statement(statement_id: "stmt-timeout", max_wait: 0, poll_interval: 0)
+      end.to raise_error(DatabricksSql::TimeoutError)
+    end
+  end
+end

data/spec/databricks_sql_spec.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe DatabricksSql do
+  before do
+    DatabricksSql.reset_configuration!
+  end
+
+  it "has a version number" do
+    expect(DatabricksSql::VERSION).not_to be_nil
+  end
+
+  it "supports global configure with DatabricksSql" do
+    DatabricksSql.configure do |config|
+      config.host = "https://dbc.example.com"
+      config.token = "dapi-token"
+      config.warehouse_id = "wh-1"
+    end
+
+    client = DatabricksSql::Client.new
+    expect(client.host).to eq("https://dbc.example.com")
+    expect(client.warehouse_id).to eq("wh-1")
+  end
+
+  it "supports Databricks.configure alias" do
+    Databricks.configure do |config|
+      config.host = "https://dbc.alias.example.com"
+      config.token = "dapi-alias"
+      config.warehouse_id = "wh-alias"
+    end
+
+    client = DatabricksSql::Client.new
+    expect(client.host).to eq("https://dbc.alias.example.com")
+    expect(client.warehouse_id).to eq("wh-alias")
+  end
+end

data/spec/spec_helper.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "databricks_sql"
+require "webmock/rspec"
+
+RSpec.configure do |config|
+  config.disable_monkey_patching!
+  config.warnings = true
+  config.order = :random
+  Kernel.srand config.seed
+
+  config.before do
+    WebMock.disable_net_connect!(allow_localhost: true)
+  end
+end

metadata

@@ -0,0 +1,77 @@
+--- !ruby/object:Gem::Specification
+name: databricks_sql
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Lairton Mendes
+bindir: exe
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: csv
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '3.3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '3.3'
+description: Execute SQL statements against Databricks SQL Warehouse with synchronous
+  and asynchronous polling support, inline results, and external link downloads.
+email:
+- lairton.mendes@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- ".rspec"
+- CHANGELOG.md
+- CODE_OF_CONDUCT.md
+- LICENSE.txt
+- README.md
+- Rakefile
+- lib/databricks_sql.rb
+- lib/databricks_sql/client.rb
+- lib/databricks_sql/configuration.rb
+- lib/databricks_sql/errors.rb
+- lib/databricks_sql/external_link_handler.rb
+- lib/databricks_sql/result.rb
+- lib/databricks_sql/type_coercer.rb
+- lib/databricks_sql/version.rb
+- sig/databricks_sql.rbs
+- spec/databricks_sql/client_spec.rb
+- spec/databricks_sql_spec.rb
+- spec/spec_helper.rb
+homepage: https://github.com/lairtonmendes/databricks_sql#readme
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/lairtonmendes/databricks_sql#readme
+  source_code_uri: https://github.com/lairtonmendes/databricks_sql
+  changelog_uri: https://github.com/lairtonmendes/databricks_sql/blob/main/CHANGELOG.md
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.2.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.3
+specification_version: 4
+summary: Ruby client for Databricks SQL Statements API
+test_files: []