databricks_sql 0.1.0

data/lib/databricks_sql/client.rb

@@ -0,0 +1,477 @@
+# frozen_string_literal: true
+
+require "json"
+require "net/http"
+require "uri"
+
+module DatabricksSql
+  class Client
+    DEFAULT_STATEMENT_PATH = "/api/2.0/sql/statements"
+
+    TERMINAL_STATUSES = %w[SUCCEEDED FAILED CANCELED CLOSED].freeze
+
+    attr_reader :host, :warehouse_id
+
+    def initialize(host: nil, token: nil, warehouse_id: nil, timeout: nil, open_timeout: nil, statement_path: nil)
+      global_config = DatabricksSql.configuration
+
+      @host = normalize_host(host || global_config.host)
+      @token = (token || global_config.token).to_s
+      @warehouse_id = (warehouse_id || global_config.warehouse_id).to_s
+      @timeout = timeout || global_config.timeout || 30
+      @open_timeout = open_timeout || global_config.open_timeout || 10
+      @statement_path = statement_path || global_config.statement_path || DEFAULT_STATEMENT_PATH
+      @external_link_require_https = global_config.external_link_require_https != false
+      @external_link_allowed_hosts = normalize_allowed_hosts(global_config.external_link_allowed_hosts)
+      validate_configuration!
+    end
+
+    def execute_statement(statement:, parameters: nil, format: "JSON_ARRAY", disposition: "INLINE", catalog: nil,
+      schema: nil, column_schema: nil, wait_timeout: "10s", on_wait_timeout: "CONTINUE",
+      poll_interval: 1.0, max_wait: 300, request_timeout: nil, cancel_on_timeout: false,
+      auto_fetch_chunks: true, byte_limit: nil, row_limit: nil, query_tags: nil)
+      response = execute_statement_async(
+        statement: statement,
+        parameters: parameters,
+        format: format,
+        disposition: disposition,
+        catalog: catalog,
+        schema: schema,
+        wait_timeout: wait_timeout,
+        on_wait_timeout: on_wait_timeout,
+        request_timeout: request_timeout,
+        byte_limit: byte_limit,
+        row_limit: row_limit,
+        query_tags: query_tags
+      )
+
+      statement_id = response.fetch("statement_id")
+      if terminal_status?(statement_status(response))
+        response = hydrate_terminal_response(
+          response,
+          statement_id: statement_id,
+          request_timeout: request_timeout
+        )
+
+        return build_result(response, format: format, disposition: disposition, column_schema: column_schema,
+                                      auto_fetch_chunks: auto_fetch_chunks,
+                                      request_timeout: request_timeout)
+      end
+
+      wait_for_statement(
+        statement_id: statement_id,
+        format: format,
+        disposition: disposition,
+        column_schema: column_schema,
+        poll_interval: poll_interval,
+        max_wait: max_wait,
+        request_timeout: request_timeout,
+        cancel_on_timeout: cancel_on_timeout,
+        auto_fetch_chunks: auto_fetch_chunks
+      )
+    end
+
+    def execute_statement_async(statement:, parameters: nil, format: "JSON_ARRAY", disposition: "INLINE", catalog: nil,
+                                schema: nil, wait_timeout: "10s", on_wait_timeout: "CONTINUE", request_timeout: nil,
+                                byte_limit: nil, row_limit: nil, query_tags: nil)
+      payload = {
+        statement: statement,
+        warehouse_id: warehouse_id,
+        disposition: disposition,
+        format: format,
+        wait_timeout: wait_timeout,
+        on_wait_timeout: on_wait_timeout
+      }
+      payload[:parameters] = parameters if parameters
+      payload[:catalog] = catalog if catalog
+      payload[:schema] = schema if schema
+      payload[:byte_limit] = byte_limit unless byte_limit.nil?
+      payload[:row_limit] = row_limit unless row_limit.nil?
+      payload[:query_tags] = query_tags if query_tags
+
+      request_json(:post, @statement_path, payload: payload, request_timeout: request_timeout)
+    end
+
+    def get_statement(statement_id:, request_timeout: nil)
+      path = "#{@statement_path}/#{statement_id}"
+      request_json(:get, path, request_timeout: request_timeout)
+    end
+
+    def cancel_statement(statement_id:, request_timeout: nil)
+      path = "#{@statement_path}/#{statement_id}/cancel"
+      request_json(:post, path, payload: {}, request_timeout: request_timeout)
+    rescue HTTPError
+      nil
+    end
+
+    def get_statement_chunk(statement_id:, chunk_index:, row_offset: nil, request_timeout: nil)
+      path = "#{@statement_path}/#{statement_id}/result/chunks/#{chunk_index}"
+      path = "#{path}?row_offset=#{row_offset}" unless row_offset.nil?
+      request_json(:get, path, request_timeout: request_timeout)
+    end
+
+    def get_statement_chunk_by_link(next_chunk_internal_link:, request_timeout: nil)
+      request_json(:get, next_chunk_internal_link, request_timeout: request_timeout)
+    end
+
+    def fetch_next_chunk(statement_id:, next_chunk_internal_link:, columns:, disposition: "INLINE", column_schema: nil,
+                         request_timeout: nil)
+      payload = get_statement_chunk_by_link(
+        next_chunk_internal_link: next_chunk_internal_link,
+        request_timeout: request_timeout
+      )
+      rows = chunk_rows(payload, columns, disposition: disposition, column_schema: column_schema,
+                                          request_timeout: request_timeout)
+
+      Result.new(
+        statement_id: statement_id,
+        status: "SUCCEEDED",
+        disposition: normalize_disposition(disposition),
+        format: "JSON_ARRAY",
+        columns: columns,
+        rows: rows,
+        manifest: nil,
+        raw_response: payload,
+        next_chunk_internal_link: extract_next_chunk_internal_link(payload)
+      )
+    end
+
+    def wait_for_statement(statement_id:, format: "JSON_ARRAY", disposition: "INLINE", column_schema: nil,
+                 poll_interval: 1.0, max_wait: 300, request_timeout: nil, cancel_on_timeout: false,
+                 auto_fetch_chunks: true)
+      started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+      loop do
+        response = get_statement(statement_id: statement_id, request_timeout: request_timeout)
+
+        if terminal_status?(statement_status(response))
+          response = hydrate_terminal_response(
+            response,
+            statement_id: statement_id,
+            request_timeout: request_timeout
+          )
+
+          return build_result(response, format: format, disposition: disposition, column_schema: column_schema,
+                                        auto_fetch_chunks: auto_fetch_chunks,
+                                        request_timeout: request_timeout)
+        end
+
+        elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at
+        if elapsed >= max_wait
+          cancel_statement(statement_id: statement_id, request_timeout: request_timeout) if cancel_on_timeout
+          raise TimeoutError.new("Polling exceeded max_wait of #{max_wait}s for statement #{statement_id}", max_wait)
+        end
+
+        sleep(poll_interval)
+      end
+    end
+
+    def inspect
+      "#<#{self.class} " \
+        "@external_link_allowed_hosts=#{@external_link_allowed_hosts.inspect}, " \
+        "@external_link_require_https=#{@external_link_require_https.inspect}, " \
+        "@host=#{@host.inspect}, " \
+        "@open_timeout=#{@open_timeout.inspect}, " \
+        "@statement_path=#{@statement_path.inspect}, " \
+        "@timeout=#{@timeout.inspect}, " \
+        "@token=\"[REDACTED]\", " \
+        "@warehouse_id=#{@warehouse_id.inspect}>"
+    end
+
+    private
+
+    def validate_configuration!
+      raise ConfigurationError, "host is required" if @host.empty?
+      raise ConfigurationError, "token is required" if @token.empty?
+      raise ConfigurationError, "warehouse_id is required" if @warehouse_id.empty?
+      raise ConfigurationError, "timeout must be positive" unless @timeout.to_f.positive?
+      raise ConfigurationError, "open_timeout must be positive" unless @open_timeout.to_f.positive?
+
+      uri = parse_uri!(@host, field_name: "host")
+      raise ConfigurationError, "host must use HTTPS" unless uri.scheme == "https"
+    end
+
+    def normalize_host(host)
+      host.to_s.strip.sub(%r{/+\z}, "")
+    end
+
+    def terminal_status?(status)
+      TERMINAL_STATUSES.include?(status.to_s.upcase)
+    end
+
+    def statement_status(payload)
+      state = payload.dig("status", "state")
+      raise ParseError, "Missing required field: status.state" if state.to_s.empty?
+
+      state.to_s.upcase
+    end
+
+    def build_result(payload, format:, disposition:, column_schema:, auto_fetch_chunks:, request_timeout:)
+      status = statement_status(payload)
+      statement_id = payload["statement_id"]
+
+      if %w[FAILED CANCELED CLOSED].include?(status)
+        message = payload.dig("status", "error", "message") ||
+                  "Statement #{statement_id} ended with status #{status}"
+        raise ExecutionError.new(message, statement_id: statement_id, status: status)
+      end
+
+      rows, columns = resolve_rows_and_columns(
+        payload,
+        disposition: disposition,
+        column_schema: column_schema,
+        request_timeout: request_timeout
+      )
+
+      next_chunk_internal_link = extract_next_chunk_internal_link(payload)
+
+      if auto_fetch_chunks
+        rows, next_chunk_internal_link = append_remaining_chunks(
+          rows,
+          columns,
+          next_chunk_internal_link: next_chunk_internal_link,
+          disposition: disposition,
+          column_schema: column_schema,
+          request_timeout: request_timeout
+        )
+      end
+
+      Result.new(
+        statement_id: statement_id,
+        status: status,
+        disposition: normalize_disposition(disposition),
+        format: format.to_s.upcase,
+        columns: columns,
+        rows: rows,
+        manifest: payload["manifest"],
+        raw_response: payload,
+        next_chunk_internal_link: next_chunk_internal_link
+      )
+    end
+
+    def resolve_rows_and_columns(payload, disposition:, column_schema:, request_timeout:)
+      effective_disposition = infer_disposition(payload, disposition)
+      columns = extract_columns(payload)
+
+      rows = if effective_disposition == "EXTERNAL_LINK"
+               fetch_external_rows(payload, request_timeout: request_timeout)
+             else
+               inline_rows(payload, columns)
+             end
+
+      rows = normalize_rows(rows, columns)
+      rows = TypeCoercer.coerce_rows(rows, column_schema)
+
+      [rows, columns]
+    end
+
+    def infer_disposition(payload, fallback)
+      result = payload["result"] || {}
+      value = result["disposition"] || result["result_disposition"] || payload["disposition"] || fallback
+      normalize_disposition(value)
+    end
+
+    def normalize_disposition(value)
+      disposition = value.to_s.upcase
+      return "EXTERNAL_LINK" if disposition == "EXTERNAL_LINKS"
+
+      disposition
+    end
+
+    def extract_next_chunk_internal_link(payload)
+      payload.dig("result", "next_chunk_internal_link") ||
+        payload.dig("result", "external_links", 0, "next_chunk_internal_link") ||
+        payload["next_chunk_internal_link"] ||
+        payload.dig("external_links", 0, "next_chunk_internal_link")
+    end
+
+    def append_remaining_chunks(rows, columns, next_chunk_internal_link:, disposition:, column_schema:,
+request_timeout:)
+      combined_rows = rows.dup
+      next_link = next_chunk_internal_link
+
+      until next_link.to_s.empty?
+        payload = get_statement_chunk_by_link(next_chunk_internal_link: next_link, request_timeout: request_timeout)
+        combined_rows.concat(
+          chunk_rows(payload, columns, disposition: disposition, column_schema: column_schema,
+                                       request_timeout: request_timeout)
+        )
+        next_link = extract_next_chunk_internal_link(payload)
+      end
+
+      [combined_rows, next_link]
+    end
+
+    def chunk_rows(payload, columns, disposition:, column_schema:, request_timeout:)
+      effective_disposition = normalize_disposition(disposition)
+
+      rows = if effective_disposition == "EXTERNAL_LINK"
+               fetch_external_rows(payload, request_timeout: request_timeout)
+             else
+               data = payload["data_array"] || payload.dig("result", "data_array") || []
+
+               if data.first.is_a?(Hash)
+                 data
+               else
+                 Array(data).map { |row| Array(columns).zip(Array(row)).to_h }
+               end
+             end
+
+      rows = normalize_rows(rows, columns)
+      TypeCoercer.coerce_rows(rows, column_schema)
+    end
+
+    def hydrate_terminal_response(payload, statement_id:, request_timeout:)
+      return payload unless statement_status(payload) == "SUCCEEDED"
+      return payload if payload["result"].is_a?(Hash)
+
+      get_statement(statement_id: statement_id, request_timeout: request_timeout)
+    end
+
+    def inline_rows(payload, columns)
+      result = payload["result"] || {}
+      data = result["data_array"] || payload["data_array"] || []
+
+      return [] unless data.is_a?(Array)
+
+      if data.first.is_a?(Hash)
+        data
+      else
+        data.map do |row|
+          Array(columns).zip(Array(row)).to_h
+        end
+      end
+    end
+
+    def fetch_external_rows(payload, request_timeout:)
+      result = payload["result"] || {}
+      url = result["external_link"] ||
+            result["external_link_url"] ||
+            result.dig("external_links", 0, "external_link") ||
+            result.dig("external_links", 0, "external_link_url") ||
+            payload.dig("external_links", 0, "external_link") ||
+            payload.dig("external_links", 0, "external_link_url") ||
+            payload["external_link"]
+
+      raise ParseError, "EXTERNAL_LINK disposition returned without a download URL" if url.to_s.empty?
+
+      ExternalLinkHandler.fetch_and_parse(
+        url,
+        timeout: request_timeout || @timeout,
+        require_https: @external_link_require_https,
+        allowed_hosts: @external_link_allowed_hosts
+      )
+    end
+
+    def normalize_allowed_hosts(hosts)
+      return nil if hosts.nil?
+
+      Array(hosts).map { |host| host.to_s.strip.downcase }.reject(&:empty?).uniq
+    end
+
+    def parse_uri!(value, field_name:)
+      uri = URI.parse(value)
+      raise ConfigurationError, "#{field_name} must be an absolute URL" if uri.scheme.nil? || uri.host.nil?
+
+      uri
+    rescue URI::InvalidURIError => e
+      raise ConfigurationError, "#{field_name} is not a valid URL: #{e.message}"
+    end
+
+    def normalize_rows(rows, columns)
+      return [] unless rows.is_a?(Array)
+      return rows if rows.empty?
+
+      if rows.first.is_a?(Hash)
+        rows
+      else
+        rows.map { |row| Array(columns).zip(Array(row)).to_h }
+      end
+    end
+
+    def extract_columns(payload)
+      result = payload["result"] || {}
+      schema_columns = result.dig("schema", "columns") || payload.dig("manifest", "schema", "columns") || []
+
+      columns = schema_columns.map do |column|
+        column["name"] || column["column_name"]
+      end.compact
+
+      return columns unless columns.empty?
+
+      first_row = (result["data_array"] || payload["data_array"] || []).first
+      return first_row.keys if first_row.is_a?(Hash)
+
+      []
+    end
+
+    def request_json(method, path, payload: nil, request_timeout: nil)
+      uri = URI.parse("#{@host}#{path}")
+      http = Net::HTTP.new(uri.host, uri.port)
+      http.use_ssl = uri.scheme == "https"
+      http.open_timeout = @open_timeout
+      http.read_timeout = request_timeout || @timeout
+
+      request = build_request(method, uri, payload)
+      response = http.request(request)
+      validate_http_response!(response)
+      parse_json(response.body)
+    rescue Timeout::Error, Net::OpenTimeout, Net::ReadTimeout
+      timeout = request_timeout || @timeout
+      raise TimeoutError.new("HTTP request timed out after #{timeout}s", timeout)
+    rescue SocketError => e
+      raise ConnectionError, "Connection error: #{e.message}"
+    rescue JSON::ParserError => e
+      raise ParseError, "Invalid JSON response: #{e.message}"
+    end
+
+    def build_request(method, uri, payload)
+      request = case method.to_sym
+                when :get
+                  Net::HTTP::Get.new(uri.request_uri)
+                when :post
+                  Net::HTTP::Post.new(uri.request_uri)
+                else
+                  raise ConfigurationError, "Unsupported HTTP method: #{method}"
+                end
+
+      request["Authorization"] = "Bearer #{@token}"
+      request["Content-Type"] = "application/json"
+      request["User-Agent"] = "databricks_sql/#{DatabricksSql::VERSION}"
+      request.body = JSON.generate(payload) if payload
+      request
+    end
+
+    def validate_http_response!(response)
+      code = response.code.to_i
+      return if code.between?(200, 299)
+
+      message = "Databricks API request failed with HTTP #{code}"
+
+      case code
+      when 401
+        raise AuthenticationError.new(message, status_code: code, response_body: response.body)
+      when 403
+        raise AuthorizationError.new(message, status_code: code, response_body: response.body)
+      when 404
+        raise NotFoundError.new(message, status_code: code, response_body: response.body)
+      when 429
+        retry_after = response["Retry-After"]&.to_i
+        raise RateLimitError.new(
+          message,
+          status_code: code,
+          response_body: response.body,
+          retry_after_seconds: retry_after
+        )
+      when 500..599
+        raise ServerError.new(message, status_code: code, response_body: response.body)
+      else
+        raise HTTPError.new(message, status_code: code, response_body: response.body)
+      end
+    end
+
+    def parse_json(body)
+      JSON.parse(body.to_s)
+    end
+  end
+end

data/lib/databricks_sql/configuration.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module DatabricksSql
+  class Configuration
+    attr_accessor :host,
+                  :token,
+                  :warehouse_id,
+                  :timeout,
+                  :open_timeout,
+                  :statement_path,
+                  :external_link_require_https,
+                  :external_link_allowed_hosts
+
+    def initialize
+      @timeout = 30
+      @open_timeout = 10
+      @statement_path = Client::DEFAULT_STATEMENT_PATH
+      @external_link_require_https = true
+      @external_link_allowed_hosts = nil
+    end
+  end
+end

data/lib/databricks_sql/errors.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module DatabricksSql
+  class Error < StandardError; end
+
+  class ConfigurationError < Error; end
+  class ConnectionError < Error; end
+  class ParseError < Error; end
+
+  class TimeoutError < Error
+    attr_reader :timeout_seconds
+
+    def initialize(message, timeout_seconds)
+      @timeout_seconds = timeout_seconds
+      super(message)
+    end
+  end
+
+  class HTTPError < Error
+    attr_reader :status_code, :response_body
+
+    def initialize(message, status_code:, response_body: nil)
+      @status_code = status_code
+      @response_body = response_body
+      super(message)
+    end
+  end
+
+  class AuthenticationError < HTTPError; end
+  class AuthorizationError < HTTPError; end
+  class NotFoundError < HTTPError; end
+
+  class RateLimitError < HTTPError
+    attr_reader :retry_after_seconds
+
+    def initialize(message, status_code:, response_body: nil, retry_after_seconds: nil)
+      @retry_after_seconds = retry_after_seconds
+      super(message, status_code: status_code, response_body: response_body)
+    end
+  end
+
+  class ServerError < HTTPError; end
+
+  class ExecutionError < Error
+    attr_reader :statement_id, :status
+
+    def initialize(message, statement_id:, status:)
+      @statement_id = statement_id
+      @status = status
+      super(message)
+    end
+  end
+end

data/lib/databricks_sql/external_link_handler.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require "csv"
+require "json"
+require "net/http"
+require "uri"
+
+module DatabricksSql
+  class ExternalLinkHandler
+    class << self
+      def fetch_and_parse(url, timeout: 30, require_https: true, allowed_hosts: nil)
+        uri = URI.parse(url)
+        validate_external_uri!(uri, require_https: require_https, allowed_hosts: allowed_hosts)
+
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = uri.scheme == "https"
+        http.open_timeout = timeout
+        http.read_timeout = timeout
+
+        response = http.request(Net::HTTP::Get.new(uri.request_uri))
+        validate_download_response!(response)
+        parse_body(response.body, response["Content-Type"], url)
+      rescue Timeout::Error, Net::OpenTimeout, Net::ReadTimeout
+        raise TimeoutError.new("Timed out downloading external link after #{timeout}s", timeout)
+      rescue SocketError => e
+        raise ConnectionError, "Unable to connect to external link host: #{e.message}"
+      rescue URI::InvalidURIError => e
+        raise ConfigurationError, "Invalid external link URL: #{e.message}"
+      end
+
+      private
+
+      def validate_external_uri!(uri, require_https:, allowed_hosts:)
+        raise ConfigurationError, "External link must be an absolute URL" if uri.scheme.nil? || uri.host.nil?
+
+        raise ConfigurationError, "External link must use HTTPS" if require_https && uri.scheme != "https"
+
+        return if allowed_hosts.nil? || allowed_hosts.empty?
+
+        host = uri.host.to_s.downcase
+        return if allowed_hosts.include?(host)
+
+        raise ConfigurationError, "External link host is not in the allowed list"
+      end
+
+      def validate_download_response!(response)
+        return if response.code.to_i.between?(200, 299)
+
+        raise HTTPError.new(
+          "External link download failed with HTTP #{response.code}",
+          status_code: response.code.to_i,
+          response_body: response.body
+        )
+      end
+
+      def parse_body(body, content_type, url)
+        ct = content_type.to_s.downcase
+
+        return parse_json(body) if ct.include?("json")
+        return parse_csv(body) if ct.include?("csv") || url.downcase.end_with?(".csv")
+
+        # Fallback strategy for unknown content types.
+        parse_json(body)
+      rescue JSON::ParserError
+        parse_csv(body)
+      rescue CSV::MalformedCSVError, ArgumentError => e
+        raise ParseError, "Unable to parse external link payload: #{e.message}"
+      end
+
+      def parse_json(body)
+        parsed = JSON.parse(body)
+
+        case parsed
+        when Array
+          parsed
+        when Hash
+          parsed.fetch("rows", [parsed])
+        else
+          raise ParseError, "Unsupported JSON payload type for external link"
+        end
+      end
+
+      def parse_csv(body)
+        rows = CSV.parse(body, headers: true)
+        rows.map(&:to_h)
+      end
+    end
+  end
+end

data/lib/databricks_sql/result.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module DatabricksSql
+  class Result
+    attr_reader :statement_id, :status, :disposition, :format, :columns, :rows, :manifest, :raw_response,
+                :next_chunk_internal_link
+
+    def initialize(statement_id:, status:, disposition:, format:, columns:, rows:, manifest:, raw_response:,
+      next_chunk_internal_link: nil)
+      @statement_id = statement_id
+      @status = status
+      @disposition = disposition
+      @format = format
+      @columns = columns
+      @rows = rows
+      @manifest = manifest
+      @raw_response = raw_response
+      @next_chunk_internal_link = next_chunk_internal_link
+    end
+
+    def success?
+      status == "SUCCEEDED"
+    end
+  end
+end