spark-connect 0.2.0

data/lib/spark_connect/catalog.rb

@@ -0,0 +1,190 @@
+# frozen_string_literal: true
+
+module SparkConnect
+  # The catalog interface for inspecting and managing databases, tables,
+  # functions, and the query cache. Returned by {SparkSession#catalog}. Mirrors
+  # PySpark's `Catalog`.
+  #
+  # Methods that return rows ({#list_databases}, {#list_tables}, ...) return
+  # arrays of {Row}; predicate methods return booleans.
+  #
+  # @example
+  #   spark.catalog.list_tables.each { |t| puts t["name"] }
+  #   spark.catalog.table_exists("my_table")  #=> true
+  class Catalog
+    Proto = SparkConnect::Proto
+    C = Proto::Catalog
+
+    # @param session [SparkSession]
+    def initialize(session)
+      @session = session
+    end
+
+    # @return [String] the current default catalog.
+    def current_catalog
+      scalar(C.new(current_catalog: Proto::CurrentCatalog.new))
+    end
+
+    # Set the current catalog. @return [void]
+    def set_current_catalog(name)
+      run(C.new(set_current_catalog: Proto::SetCurrentCatalog.new(catalog_name: name.to_s)))
+    end
+
+    # @return [Array<Row>] all catalogs.
+    def list_catalogs
+      rows(C.new(list_catalogs: Proto::ListCatalogs.new))
+    end
+
+    # @return [String] the current default database.
+    def current_database
+      scalar(C.new(current_database: Proto::CurrentDatabase.new))
+    end
+
+    # Set the current database. @return [void]
+    def set_current_database(name)
+      run(C.new(set_current_database: Proto::SetCurrentDatabase.new(db_name: name.to_s)))
+    end
+
+    # @return [Array<Row>] all databases.
+    def list_databases
+      rows(C.new(list_databases: Proto::ListDatabases.new))
+    end
+
+    # @param db_name [String, nil] restrict to a database.
+    # @return [Array<Row>] tables (and views).
+    def list_tables(db_name = nil)
+      lt = Proto::ListTables.new
+      lt.db_name = db_name if db_name
+      rows(C.new(list_tables: lt))
+    end
+
+    # @return [Array<Row>] functions registered in the catalog.
+    def list_functions(db_name = nil)
+      lf = Proto::ListFunctions.new
+      lf.db_name = db_name if db_name
+      rows(C.new(list_functions: lf))
+    end
+
+    # @param table_name [String]
+    # @return [Array<Row>] the columns of a table.
+    def list_columns(table_name, db_name = nil)
+      lc = Proto::ListColumns.new(table_name: table_name.to_s)
+      lc.db_name = db_name if db_name
+      rows(C.new(list_columns: lc))
+    end
+
+    # @return [Boolean] whether a table or view exists.
+    def table_exists(table_name, db_name = nil)
+      te = Proto::TableExists.new(table_name: table_name.to_s)
+      te.db_name = db_name if db_name
+      scalar(C.new(table_exists: te)) == true
+    end
+
+    # @return [Boolean] whether a database exists.
+    def database_exists(db_name)
+      scalar(C.new(database_exists: Proto::DatabaseExists.new(db_name: db_name.to_s))) == true
+    end
+
+    # @return [Boolean] whether a function exists.
+    def function_exists(function_name, db_name = nil)
+      fe = Proto::FunctionExists.new(function_name: function_name.to_s)
+      fe.db_name = db_name if db_name
+      scalar(C.new(function_exists: fe)) == true
+    end
+
+    # Drop a session-local temporary view. @return [Boolean]
+    def drop_temp_view(view_name)
+      scalar(C.new(drop_temp_view: Proto::DropTempView.new(view_name: view_name.to_s))) == true
+    end
+
+    # Drop a global temporary view. @return [Boolean]
+    def drop_global_temp_view(view_name)
+      scalar(C.new(drop_global_temp_view: Proto::DropGlobalTempView.new(view_name: view_name.to_s))) == true
+    end
+
+    # @return [Boolean] whether the table is cached.
+    def cached?(table_name)
+      scalar(C.new(is_cached: Proto::IsCached.new(table_name: table_name.to_s))) == true
+    end
+
+    # Cache a table in memory. @return [void]
+    def cache_table(table_name)
+      run(C.new(cache_table: Proto::CacheTable.new(table_name: table_name.to_s)))
+    end
+
+    # Remove a table from the cache. @return [void]
+    def uncache_table(table_name)
+      run(C.new(uncache_table: Proto::UncacheTable.new(table_name: table_name.to_s)))
+    end
+
+    # Clear all cached tables. @return [void]
+    def clear_cache
+      run(C.new(clear_cache: Proto::ClearCache.new))
+    end
+
+    # Invalidate and refresh cached metadata for a table. @return [void]
+    def refresh_table(table_name)
+      run(C.new(refresh_table: Proto::RefreshTable.new(table_name: table_name.to_s)))
+    end
+
+    # Recover all partitions of a table. @return [void]
+    def recover_partitions(table_name)
+      run(C.new(recover_partitions: Proto::RecoverPartitions.new(table_name: table_name.to_s)))
+    end
+
+    # Create a managed table and return a {DataFrame} over it.
+    #
+    # @param table_name [String]
+    # @param path [String, nil]
+    # @param source [String, nil] the data source/format.
+    # @param schema [Types::StructType, nil]
+    # @param description [String, nil]
+    # @param options [Hash{String=>String}]
+    # @return [DataFrame]
+    def create_table(table_name, path: nil, source: nil, schema: nil, description: nil, options: {})
+      ct = Proto::CreateTable.new(table_name: table_name.to_s, options: stringify(options))
+      ct.path = path if path
+      ct.source = source if source
+      ct.description = description if description
+      ct.schema = schema.to_proto if schema
+      catalog_df(C.new(create_table: ct)).collect # eagerly create the table
+      @session.table(table_name.to_s)
+    end
+
+    # Create a table backed by data at `path` (an external/unmanaged table).
+    #
+    # @return [DataFrame]
+    def create_external_table(table_name, path: nil, source: nil, schema: nil, options: {})
+      ct = Proto::CreateExternalTable.new(table_name: table_name.to_s, options: stringify(options))
+      ct.path = path if path
+      ct.source = source if source
+      ct.schema = schema.to_proto if schema
+      catalog_df(C.new(create_external_table: ct)).collect # eagerly create the table
+      @session.table(table_name.to_s)
+    end
+
+    private
+
+    def stringify(options)
+      options.to_h { |k, v| [k.to_s, v.to_s] }
+    end
+
+    def catalog_df(catalog)
+      DataFrame.new(@session, PlanBuilder.relation(@session, catalog: catalog))
+    end
+
+    def rows(catalog)
+      catalog_df(catalog).collect
+    end
+
+    def scalar(catalog)
+      row = catalog_df(catalog).collect.first
+      row&.[](0)
+    end
+
+    def run(catalog)
+      catalog_df(catalog).collect
+      nil
+    end
+  end
+end

data/lib/spark_connect/channel_builder.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+require "grpc"
+require "uri"
+
+module SparkConnect
+  # Parses a Spark Connect connection string (`sc://...`) and builds the gRPC
+  # stub, credentials, and per-request metadata.
+  #
+  # The connection string grammar mirrors the official Spark Connect clients:
+  #
+  #     sc://host[:port][/;param=value;param=value...]
+  #
+  # Recognised parameters:
+  # * `token`      - bearer token; implies TLS and adds an `authorization` header
+  # * `user_id`    - the Spark user id
+  # * `user_agent` - client user agent (default `spark-connect-ruby/<version>`)
+  # * `use_ssl`    - `true`/`false`; force TLS on or off
+  # * `session_id` - reuse a specific server-side session id (UUID)
+  #
+  # Any parameter whose name begins with `x-` is forwarded verbatim as gRPC
+  # request metadata.
+  #
+  # @example
+  #   cb = SparkConnect::ChannelBuilder.new("sc://localhost:15002")
+  #   cb.host       #=> "localhost"
+  #   cb.port       #=> 15002
+  class ChannelBuilder
+    DEFAULT_PORT = 15_002
+    PARAM_PREFIX = "x-"
+
+    # @return [String]
+    attr_reader :host
+    # @return [Integer]
+    attr_reader :port
+    # @return [Hash{String=>String}] raw connection parameters.
+    attr_reader :params
+    # @return [String, nil]
+    attr_reader :token
+    # @return [String, nil]
+    attr_reader :user_id
+    # @return [String, nil]
+    attr_reader :session_id
+
+    # @param url [String] an `sc://` connection string.
+    def initialize(url)
+      raise ConnectionError, "Connection string must not be nil" if url.nil?
+      raise ConnectionError, "Connection string must start with 'sc://', got: #{url.inspect}" unless url.start_with?("sc://")
+
+      body = url.delete_prefix("sc://")
+      endpoint, _, param_str = body.partition("/")
+      @params = parse_params(param_str)
+      parse_endpoint(endpoint)
+
+      @token = @params["token"]
+      @user_id = @params["user_id"]
+      @session_id = @params["session_id"]
+      @use_ssl = parse_bool(@params["use_ssl"]) || !@token.nil?
+    end
+
+    # @return [Boolean] whether the channel uses TLS.
+    def ssl?
+      @use_ssl
+    end
+
+    # @return [String] the gRPC target, e.g. `"localhost:15002"`.
+    def target
+      "#{host}:#{port}"
+    end
+
+    # @return [String] the effective user agent.
+    def user_agent
+      @params["user_agent"] || "spark-connect-ruby/#{SparkConnect::VERSION}"
+    end
+
+    # Per-request gRPC metadata derived from the connection string (bearer token
+    # plus any `x-*` parameters).
+    #
+    # @return [Hash{String=>String}]
+    def metadata
+      md = {}
+      md["authorization"] = "Bearer #{@token}" if @token
+      @params.each { |k, v| md[k] = v if k.start_with?(PARAM_PREFIX) }
+      md
+    end
+
+    # Build the gRPC stub for the {Spark::Connect::SparkConnectService}.
+    #
+    # @param channel_args [Hash] extra gRPC channel arguments.
+    # @return [Spark::Connect::SparkConnectService::Stub]
+    def build_stub(channel_args: {})
+      creds = @use_ssl ? GRPC::Core::ChannelCredentials.new : :this_channel_is_insecure
+      args = { "grpc.primary_user_agent" => user_agent }.merge(channel_args)
+      Proto::SparkConnectService::Stub.new(target, creds, channel_args: args)
+    end
+
+    private
+
+    Proto = SparkConnect::Proto
+
+    def parse_endpoint(endpoint)
+      raise ConnectionError, "Missing host in connection string" if endpoint.nil? || endpoint.empty?
+
+      host, sep, port = endpoint.rpartition(":")
+      if sep.empty?
+        @host = endpoint
+        @port = DEFAULT_PORT
+      else
+        @host = host
+        @port = Integer(port, exception: false) ||
+                raise(ConnectionError, "Invalid port in connection string: #{port.inspect}")
+      end
+    end
+
+    def parse_params(param_str)
+      params = {}
+      param_str.split(";").each do |kv|
+        next if kv.empty?
+
+        key, sep, value = kv.partition("=")
+        raise ConnectionError, "Malformed parameter (expected key=value): #{kv.inspect}" if sep.empty?
+
+        params[URI.decode_www_form_component(key)] = URI.decode_www_form_component(value)
+      end
+      params
+    end
+
+    def parse_bool(value)
+      return nil if value.nil?
+
+      %w[true 1 yes].include?(value.to_s.downcase)
+    end
+  end
+end

data/lib/spark_connect/client.rb

@@ -0,0 +1,264 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module SparkConnect
+  # The low-level Spark Connect client. Wraps the gRPC stub and exposes the four
+  # core RPC families used by the high-level API: {#execute_plan},
+  # {#execute_command}, {#analyze}, and {#config}. Higher layers
+  # ({SparkSession}, {DataFrame}) never touch the stub directly.
+  #
+  # Transient transport failures (e.g. `GRPC::Unavailable`) are retried with
+  # exponential backoff and jitter before any response data has been observed.
+  class SparkConnectClient
+    Proto = SparkConnect::Proto
+
+    # Accumulated result of an `ExecutePlan` stream.
+    #
+    # @!attribute [r] arrow_batches
+    #   @return [Array<String>] each element is one Arrow IPC stream chunk.
+    # @!attribute [r] schema
+    #   @return [Spark::Connect::DataType, nil] result schema, if returned.
+    # @!attribute [r] metrics
+    #   @return [Spark::Connect::ExecutePlanResponse::Metrics, nil]
+    # @!attribute [r] observed_metrics
+    #   @return [Array] observed (named) metrics.
+    # @!attribute [r] sql_command_result
+    #   @return [Spark::Connect::Relation, nil] relation produced by a SQL command.
+    ExecuteResult = Struct.new(
+      :arrow_batches, :schema, :metrics, :observed_metrics, :sql_command_result, :row_count,
+      :write_stream_result, :streaming_query_result, :streaming_manager_result, :checkpoint_relation,
+      :pipeline_command_result, :pipeline_events
+    )
+
+    # @return [String] the client-side session id (UUID v4).
+    attr_reader :session_id
+    # @return [String] the user agent / client type.
+    attr_reader :client_type
+    # @return [ChannelBuilder]
+    attr_reader :channel_builder
+
+    # @param channel_builder [ChannelBuilder]
+    # @param session_id [String, nil] reuse a session id, otherwise generated.
+    # @param max_retries [Integer]
+    # @param retry_base_delay [Float] base backoff in seconds.
+    def initialize(channel_builder, session_id: nil, max_retries: 10, retry_base_delay: 0.05, max_retry_delay: 10.0)
+      @channel_builder = channel_builder
+      @stub = channel_builder.build_stub
+      @metadata = channel_builder.metadata
+      @session_id = session_id || channel_builder.session_id || SecureRandom.uuid
+      @client_type = channel_builder.user_agent
+      @user_context = Proto::UserContext.new(user_id: channel_builder.user_id || "")
+      @max_retries = max_retries
+      @retry_base_delay = retry_base_delay
+      @max_retry_delay = max_retry_delay
+      @server_side_session_id = nil
+      @tags = []
+    end
+
+    # Operation tags attached to every subsequent execution (used with
+    # {#interrupt} `type: :tag`).
+    # @return [Array<String>]
+    attr_reader :tags
+
+    # Add an operation tag. Tags must be non-empty and contain no commas.
+    # @return [void]
+    def add_tag(tag)
+      tag = tag.to_s
+      raise IllegalArgumentError, "Tag must not be empty" if tag.empty?
+      raise IllegalArgumentError, "Tag must not contain ','" if tag.include?(",")
+
+      @tags << tag unless @tags.include?(tag)
+      nil
+    end
+
+    # Remove an operation tag. @return [void]
+    def remove_tag(tag)
+      @tags.delete(tag.to_s)
+      nil
+    end
+
+    # Remove all operation tags. @return [void]
+    def clear_tags
+      @tags.clear
+      nil
+    end
+
+    # Execute a relation plan and accumulate the streamed response.
+    #
+    # @param relation [Spark::Connect::Relation]
+    # @return [ExecuteResult]
+    def execute_plan(relation)
+      execute(PlanBuilder.root_plan(relation))
+    end
+
+    # Execute a command plan (side-effecting, e.g. write/SQL DML).
+    #
+    # @param command [Spark::Connect::Command]
+    # @return [ExecuteResult]
+    def execute_command(command)
+      execute(PlanBuilder.command_plan(command))
+    end
+
+    # Run an `AnalyzePlan` request.
+    #
+    # @param analyze_kw [Hash] exactly one `analyze` oneof keyword, e.g.
+    #   `schema:`, `explain:`, `tree_string:`, `is_local:`, `spark_version:`.
+    # @return [Spark::Connect::AnalyzePlanResponse]
+    def analyze(**analyze_kw)
+      req = Proto::AnalyzePlanRequest.new(
+        session_id: @session_id,
+        user_context: @user_context,
+        client_type: @client_type,
+        **analyze_kw
+      )
+      with_retries { @stub.analyze_plan(req, metadata: @metadata) }
+    end
+
+    # Run a `Config` request.
+    #
+    # @param operation [Spark::Connect::ConfigRequest::Operation]
+    # @return [Spark::Connect::ConfigResponse]
+    def config(operation)
+      req = Proto::ConfigRequest.new(
+        session_id: @session_id,
+        user_context: @user_context,
+        client_type: @client_type,
+        operation: operation
+      )
+      with_retries { @stub.config(req, metadata: @metadata) }
+    end
+
+    # Interrupt running operations.
+    #
+    # @param type [Symbol] `:all`, `:tag`, or `:operation_id`.
+    # @param value [String, nil] the tag or operation id when applicable.
+    # @return [Spark::Connect::InterruptResponse]
+    def interrupt(type: :all, value: nil)
+      kw = { interrupt_type: :"INTERRUPT_TYPE_#{type.to_s.upcase}" }
+      kw[:operation_tag] = value if type == :tag
+      kw[:operation_id] = value if type == :operation_id
+      req = Proto::InterruptRequest.new(
+        session_id: @session_id, user_context: @user_context, client_type: @client_type, **kw
+      )
+      with_retries { @stub.interrupt(req, metadata: @metadata) }
+    end
+
+    # Release this client's server-side session.
+    # @return [void]
+    def release_session
+      req = Proto::ReleaseSessionRequest.new(
+        session_id: @session_id, user_context: @user_context, client_type: @client_type
+      )
+      # Best-effort and non-retrying: this runs on teardown, so a dead server
+      # must not block the caller with the retry/backoff loop.
+      @stub.release_session(req, metadata: @metadata)
+      nil
+    rescue StandardError
+      nil
+    end
+
+    private
+
+    def execute(plan)
+      operation_id = SecureRandom.uuid
+      req = Proto::ExecutePlanRequest.new(
+        session_id: @session_id,
+        user_context: @user_context,
+        operation_id: operation_id,
+        plan: plan,
+        client_type: @client_type,
+        tags: @tags
+      )
+
+      result = ExecuteResult.new([], nil, nil, [], nil, 0)
+      result.pipeline_events = []
+      with_retries do
+        responses = @stub.execute_plan(req, metadata: @metadata)
+        responses.each do |resp|
+          @server_side_session_id = resp.server_side_session_id unless resp.server_side_session_id.empty?
+          accumulate(result, resp)
+        end
+      end
+      result
+    end
+
+    def accumulate(result, resp)
+      result.schema = resp.schema if resp.schema
+      result.metrics = resp.metrics if resp.metrics
+      result.observed_metrics += resp.observed_metrics.to_a unless resp.observed_metrics.empty?
+
+      case resp.response_type
+      when :arrow_batch
+        batch = resp.arrow_batch
+        result.arrow_batches << batch.data unless batch.data.empty?
+        result.row_count += batch.row_count
+      when :sql_command_result
+        result.sql_command_result = resp.sql_command_result.relation
+      when :write_stream_operation_start_result
+        result.write_stream_result = resp.write_stream_operation_start_result
+      when :streaming_query_command_result
+        result.streaming_query_result = resp.streaming_query_command_result
+      when :streaming_query_manager_command_result
+        result.streaming_manager_result = resp.streaming_query_manager_command_result
+      when :checkpoint_command_result
+        result.checkpoint_relation = resp.checkpoint_command_result.relation
+      when :pipeline_command_result
+        result.pipeline_command_result = resp.pipeline_command_result
+      when :pipeline_event_result
+        result.pipeline_events << resp.pipeline_event_result.event
+      end
+    end
+
+    def with_retries
+      attempt = 0
+      begin
+        yield
+      rescue GRPC::BadStatus => e
+        if retryable?(e) && attempt < @max_retries
+          delay = backoff(attempt)
+          attempt += 1
+          sleep(delay)
+          retry
+        end
+        raise translate_error(e)
+      end
+    end
+
+    RETRYABLE_CODES = [
+      GRPC::Core::StatusCodes::UNAVAILABLE,
+      GRPC::Core::StatusCodes::DEADLINE_EXCEEDED,
+      GRPC::Core::StatusCodes::ABORTED,
+      GRPC::Core::StatusCodes::RESOURCE_EXHAUSTED,
+    ].freeze
+
+    def retryable?(error)
+      RETRYABLE_CODES.include?(error.code)
+    end
+
+    def backoff(attempt)
+      delay = @retry_base_delay * (2**attempt)
+      delay = [delay, @max_retry_delay].min
+      delay + (rand * delay * 0.5)
+    end
+
+    def translate_error(error)
+      message = error.respond_to?(:details) ? error.details : error.message
+      code = grpc_code_name(error)
+      klass =
+        if /\[(PARSE_SYNTAX_ERROR|PARSE_)/.match?(message.to_s) || /ParseException/.match?(message.to_s)
+          ParseError
+        elsif /AnalysisException|UNRESOLVED_|TABLE_OR_VIEW_NOT_FOUND|\[.*\] /.match?(message.to_s)
+          AnalysisError
+        else
+          SparkConnectError
+        end
+      error_class = message.to_s[/\[([A-Z0-9_.]+)\]/, 1]
+      klass.new(message, error_class: error_class, grpc_code: code)
+    end
+
+    def grpc_code_name(error)
+      GRPC::Core::StatusCodes.constants.find { |c| GRPC::Core::StatusCodes.const_get(c) == error.code }&.to_s
+    end
+  end
+end