spark-connect 0.2.0 → 0.4.0

data/lib/spark_connect/channel_builder.rb

@@ -13,13 +13,18 @@ module SparkConnect
   #
   # Recognised parameters:
   # * `token`      - bearer token; implies TLS and adds an `authorization` header
+  #   (defaults to the `SPARK_CONNECT_AUTHENTICATE_TOKEN` environment variable)
   # * `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)
+  # * `grpc_keepalive_enabled` (default `true`), `grpc_keepalive_time_ms`
+  #   (default 60000), `grpc_keepalive_timeout_ms` (default 20000), and
+  #   `grpc_keepalive_without_calls` (default `true`) - gRPC/HTTP2 keepalive
+  #   PINGs that detect a silently dead connection (Spark 4.3 client defaults)
   #
-  # Any parameter whose name begins with `x-` is forwarded verbatim as gRPC
-  # request metadata.
+  # Every other parameter is forwarded verbatim as gRPC request metadata (like
+  # the official clients), e.g. `x-api-key=...`.
   #
   # @example
   #   cb = SparkConnect::ChannelBuilder.new("sc://localhost:15002")
@@ -29,6 +34,34 @@ module SparkConnect
     DEFAULT_PORT = 15_002
     PARAM_PREFIX = "x-"
 
+    PARAM_TOKEN = "token"
+    PARAM_USER_ID = "user_id"
+    PARAM_USER_AGENT = "user_agent"
+    PARAM_USE_SSL = "use_ssl"
+    PARAM_SESSION_ID = "session_id"
+    PARAM_GRPC_KEEPALIVE_ENABLED = "grpc_keepalive_enabled"
+    PARAM_GRPC_KEEPALIVE_TIME_MS = "grpc_keepalive_time_ms"
+    PARAM_GRPC_KEEPALIVE_TIMEOUT_MS = "grpc_keepalive_timeout_ms"
+    PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS = "grpc_keepalive_without_calls"
+
+    # Parameters consumed by the client itself and never sent as metadata.
+    RESERVED_PARAMS = [
+      PARAM_TOKEN, PARAM_USER_ID, PARAM_USER_AGENT, PARAM_USE_SSL, PARAM_SESSION_ID,
+      PARAM_GRPC_KEEPALIVE_ENABLED, PARAM_GRPC_KEEPALIVE_TIME_MS, PARAM_GRPC_KEEPALIVE_TIMEOUT_MS,
+      PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS,
+    ].freeze
+
+    GRPC_DEFAULT_KEEPALIVE_ENABLED = true
+    GRPC_DEFAULT_KEEPALIVE_TIME_MS = 60 * 1000
+    GRPC_DEFAULT_KEEPALIVE_TIMEOUT_MS = 20 * 1000
+    GRPC_DEFAULT_KEEPALIVE_WITHOUT_CALLS = true
+
+    # Like the official clients, allow messages (e.g. large Arrow batches) of up
+    # to 128 MiB instead of gRPC's 4 MiB receive default.
+    GRPC_MAX_MESSAGE_LENGTH_DEFAULT = 128 * 1024 * 1024
+
+    UUID_PATTERN = /\A\h{8}-\h{4}-\h{4}-\h{4}-\h{12}\z/
+
     # @return [String]
     attr_reader :host
     # @return [Integer]
@@ -52,10 +85,66 @@ module SparkConnect
       @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?
+      @token = @params[PARAM_TOKEN] || ENV.fetch("SPARK_CONNECT_AUTHENTICATE_TOKEN", nil)
+      @user_id = @params[PARAM_USER_ID]
+      @session_id = @params[PARAM_SESSION_ID]
+      validate_session_id!(@session_id)
+      @use_ssl = parse_bool(@params[PARAM_USE_SSL]) || !@token.nil?
+    end
+
+    # A copy of this builder with `session_id` replaced (or removed when nil).
+    # Used to create a fresh ({SparkSession#new_session}) or cloned
+    # ({SparkSession#clone_session}) session against the same endpoint.
+    #
+    # @param session_id [String, nil]
+    # @return [ChannelBuilder]
+    def with_session_id(session_id)
+      copy = dup
+      params = @params.dup
+      if session_id.nil?
+        params.delete(PARAM_SESSION_ID)
+      else
+        params[PARAM_SESSION_ID] = session_id
+      end
+      copy.instance_variable_set(:@params, params)
+      copy.instance_variable_set(:@session_id, session_id)
+      copy
+    end
+
+    # @return [Boolean] whether gRPC keepalive PINGs are sent.
+    def keepalive_enabled?
+      bool_param(PARAM_GRPC_KEEPALIVE_ENABLED, GRPC_DEFAULT_KEEPALIVE_ENABLED)
+    end
+
+    # @return [Integer] idle time (ms) before a keepalive PING.
+    def keepalive_time_ms
+      int_param(PARAM_GRPC_KEEPALIVE_TIME_MS, GRPC_DEFAULT_KEEPALIVE_TIME_MS)
+    end
+
+    # @return [Integer] time (ms) to wait for a keepalive PING ack.
+    def keepalive_timeout_ms
+      int_param(PARAM_GRPC_KEEPALIVE_TIMEOUT_MS, GRPC_DEFAULT_KEEPALIVE_TIMEOUT_MS)
+    end
+
+    # @return [Boolean] whether keepalive PINGs are sent without in-flight calls.
+    def keepalive_without_calls?
+      bool_param(PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS, GRPC_DEFAULT_KEEPALIVE_WITHOUT_CALLS)
+    end
+
+    # The gRPC channel arguments derived from the connection string.
+    # @return [Hash{String=>Object}]
+    def channel_args
+      args = {
+        "grpc.primary_user_agent" => user_agent,
+        "grpc.max_send_message_length" => GRPC_MAX_MESSAGE_LENGTH_DEFAULT,
+        "grpc.max_receive_message_length" => GRPC_MAX_MESSAGE_LENGTH_DEFAULT,
+      }
+      if keepalive_enabled?
+        args["grpc.keepalive_time_ms"] = keepalive_time_ms
+        args["grpc.keepalive_timeout_ms"] = keepalive_timeout_ms
+        args["grpc.keepalive_permit_without_calls"] = keepalive_without_calls? ? 1 : 0
+      end
+      args
     end
 
     # @return [Boolean] whether the channel uses TLS.
@@ -73,25 +162,25 @@ module SparkConnect
       @params["user_agent"] || "spark-connect-ruby/#{SparkConnect::VERSION}"
     end
 
-    # Per-request gRPC metadata derived from the connection string (bearer token
-    # plus any `x-*` parameters).
+    # Per-request gRPC metadata derived from the connection string: the bearer
+    # token plus every parameter the client does not consume itself.
     #
     # @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) }
+      @params.each { |k, v| md[k] = v unless RESERVED_PARAMS.include?(k) }
       md
     end
 
     # Build the gRPC stub for the {Spark::Connect::SparkConnectService}.
     #
-    # @param channel_args [Hash] extra gRPC channel arguments.
+    # @param channel_args [Hash] extra gRPC channel arguments (override the
+    #   defaults from {#channel_args}).
     # @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)
+      Proto::SparkConnectService::Stub.new(target, creds, channel_args: self.channel_args.merge(channel_args))
     end
 
     private
@@ -130,5 +219,23 @@ module SparkConnect
 
       %w[true 1 yes].include?(value.to_s.downcase)
     end
+
+    def bool_param(key, default)
+      value = @params[key]
+      value.nil? ? default : parse_bool(value)
+    end
+
+    def int_param(key, default)
+      value = @params[key]
+      return default if value.nil?
+
+      Integer(value, exception: false) || raise(ConnectionError, "Invalid integer for #{key}: #{value.inspect}")
+    end
+
+    def validate_session_id!(session_id)
+      return if session_id.nil? || UUID_PATTERN.match?(session_id)
+
+      raise ConnectionError, "session_id must be a UUID string, got #{session_id.inspect}"
+    end
   end
 end

data/lib/spark_connect/client.rb

@@ -3,6 +3,42 @@
 require "securerandom"
 
 module SparkConnect
+  # Progress of one Spark stage, reported to the handlers registered with
+  # {SparkSession#register_progress_handler} while a query runs.
+  #
+  # @!attribute [r] stage_id
+  #   @return [Integer]
+  # @!attribute [r] num_tasks
+  #   @return [Integer]
+  # @!attribute [r] num_completed_tasks
+  #   @return [Integer]
+  # @!attribute [r] num_bytes_read
+  #   @return [Integer]
+  # @!attribute [r] done
+  #   @return [Boolean]
+  StageInfo = Struct.new(:stage_id, :num_tasks, :num_completed_tasks, :num_bytes_read, :done, keyword_init: true)
+
+  # Per-RPC timeouts (in seconds; `nil` disables the timeout for that RPC),
+  # mirroring the Spark 4.3 clients' defaults. `ExecutePlan` (the query stream
+  # itself) never has a deadline, since a timeout there would kill the query.
+  #
+  # @example
+  #   SparkConnect::SparkConnectClient.new(cb, rpc_deadlines: SparkConnect::RpcDeadlines.new(analyze_plan: 120))
+  RpcDeadlines = Struct.new(
+    :analyze_plan, :config, :interrupt, :release_session, :clone_session, :get_status,
+    keyword_init: true
+  ) do
+    # @return [RpcDeadlines] the Spark 4.3 defaults (1 hour for AnalyzePlan, 10 minutes otherwise).
+    def self.defaults
+      new(analyze_plan: 3600, config: 600, interrupt: 600, release_session: 600, clone_session: 600, get_status: 600)
+    end
+
+    # @return [RpcDeadlines] all per-RPC timeouts disabled.
+    def self.disabled
+      new
+    end
+  end
+
   # 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
@@ -28,7 +64,7 @@ module SparkConnect
     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
+      :pipeline_command_result, :pipeline_events, :operation_id
     )
 
     # @return [String] the client-side session id (UUID v4).
@@ -37,12 +73,16 @@ module SparkConnect
     attr_reader :client_type
     # @return [ChannelBuilder]
     attr_reader :channel_builder
+    # @return [RpcDeadlines]
+    attr_reader :rpc_deadlines
 
     # @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)
+    # @param rpc_deadlines [RpcDeadlines] per-RPC timeouts.
+    def initialize(channel_builder, session_id: nil, max_retries: 10, retry_base_delay: 0.05, max_retry_delay: 10.0,
+                   rpc_deadlines: RpcDeadlines.defaults)
       @channel_builder = channel_builder
       @stub = channel_builder.build_stub
       @metadata = channel_builder.metadata
@@ -52,8 +92,11 @@ module SparkConnect
       @max_retries = max_retries
       @retry_base_delay = retry_base_delay
       @max_retry_delay = max_retry_delay
+      @rpc_deadlines = rpc_deadlines
       @server_side_session_id = nil
       @tags = []
+      @progress_handlers = []
+      @closed = false
     end
 
     # Operation tags attached to every subsequent execution (used with
@@ -84,6 +127,32 @@ module SparkConnect
       nil
     end
 
+    # ---- Progress handlers ---------------------------------------------------
+
+    # Register a callable notified of execution progress (see {StageInfo}).
+    # @param handler [#call]
+    # @return [void]
+    def register_progress_handler(handler)
+      raise IllegalArgumentError, "progress handler must respond to #call" unless handler.respond_to?(:call)
+
+      @progress_handlers << handler unless @progress_handlers.include?(handler)
+      nil
+    end
+
+    # @return [void]
+    def remove_progress_handler(handler)
+      @progress_handlers.delete(handler)
+      nil
+    end
+
+    # @return [void]
+    def clear_progress_handlers
+      @progress_handlers.clear
+      nil
+    end
+
+    # ---- Core RPCs -----------------------------------------------------------
+
     # Execute a relation plan and accumulate the streamed response.
     #
     # @param relation [Spark::Connect::Relation]
@@ -100,6 +169,22 @@ module SparkConnect
       execute(PlanBuilder.command_plan(command))
     end
 
+    # Execute a command and yield each raw `ExecutePlanResponse` as it arrives,
+    # without accumulating (and without retrying, since responses may already
+    # have been consumed). Used for long-running streams such as the streaming
+    # query listener bus.
+    #
+    # @param command [Spark::Connect::Command]
+    # @yieldparam response [Spark::Connect::ExecutePlanResponse]
+    # @return [void]
+    def stream_command(command, &)
+      req = execute_request(PlanBuilder.command_plan(command), SecureRandom.uuid)
+      @stub.execute_plan(req, metadata: @metadata).each(&)
+      nil
+    rescue GRPC::BadStatus => e
+      raise translate_error(e)
+    end
+
     # Run an `AnalyzePlan` request.
     #
     # @param analyze_kw [Hash] exactly one `analyze` oneof keyword, e.g.
@@ -112,7 +197,7 @@ module SparkConnect
         client_type: @client_type,
         **analyze_kw
       )
-      with_retries { @stub.analyze_plan(req, metadata: @metadata) }
+      with_retries { @stub.analyze_plan(req, **call_options(:analyze_plan)) }
     end
 
     # Run a `Config` request.
@@ -126,7 +211,7 @@ module SparkConnect
         client_type: @client_type,
         operation: operation
       )
-      with_retries { @stub.config(req, metadata: @metadata) }
+      with_retries { @stub.config(req, **call_options(:config)) }
     end
 
     # Interrupt running operations.
@@ -141,7 +226,53 @@ module SparkConnect
       req = Proto::InterruptRequest.new(
         session_id: @session_id, user_context: @user_context, client_type: @client_type, **kw
       )
-      with_retries { @stub.interrupt(req, metadata: @metadata) }
+      with_retries { @stub.interrupt(req, **call_options(:interrupt)) }
+    end
+
+    # The status of operations in this session (the `GetStatus` RPC, Spark 4.2+).
+    #
+    # @param operation_ids [Array<String>] operations to report; all operations
+    #   of the session when empty.
+    # @param operation_extensions [Array<Google::Protobuf::Any>] per-operation
+    #   extension requests.
+    # @param request_extensions [Array<Google::Protobuf::Any>] request-level extensions.
+    # @return [Spark::Connect::GetStatusResponse] each `operation_statuses`
+    #   entry has an `operation_id` and a `state` (e.g. `:OPERATION_STATE_RUNNING`).
+    def get_operation_statuses(operation_ids = [], operation_extensions: [], request_extensions: [])
+      req = Proto::GetStatusRequest.new(
+        session_id: @session_id, user_context: @user_context, client_type: @client_type,
+        operation_status: Proto::GetStatusRequest::OperationStatusRequest.new(
+          operation_ids: Array(operation_ids).map(&:to_s), extensions: operation_extensions
+        ),
+        extensions: request_extensions
+      )
+      req.client_observed_server_side_session_id = @server_side_session_id if @server_side_session_id
+      with_retries { @stub.get_status(req, **call_options(:get_status)) }
+    end
+    alias getOperationStatuses get_operation_statuses
+
+    # Clone this client's server-side session (the `CloneSession` RPC, Spark
+    # 4.1+) and return a client bound to the clone.
+    #
+    # @param new_session_id [String, nil] a UUID for the clone (generated when omitted).
+    # @return [SparkConnectClient]
+    def clone_session(new_session_id = nil)
+      req = Proto::CloneSessionRequest.new(session_id: @session_id, user_context: @user_context, client_type: @client_type)
+      req.new_session_id = new_session_id if new_session_id
+      req.client_observed_server_side_session_id = @server_side_session_id if @server_side_session_id
+      resp = with_retries { @stub.clone_session(req, **call_options(:clone_session)) }
+      if new_session_id && resp.new_session_id != new_session_id
+        raise SparkConnectError, "Server returned session id #{resp.new_session_id}, expected #{new_session_id}"
+      end
+
+      derive(resp.new_session_id)
+    end
+
+    # A client against the same endpoint bound to a fresh, empty server-side
+    # session (no state is copied; no server round trip is made).
+    # @return [SparkConnectClient]
+    def new_session_client
+      derive(nil)
     end
 
     # Release this client's server-side session.
@@ -152,17 +283,49 @@ module SparkConnect
       )
       # 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)
+      @stub.release_session(req, **call_options(:release_session))
+      nil
+    rescue StandardError
+      nil
+    end
+
+    # Close the gRPC channel. Idempotent.
+    # @return [void]
+    def close
+      return nil if @closed
+
+      @closed = true
+      channel = @stub.instance_variable_get(:@ch)
+      channel.close if channel.respond_to?(:close)
       nil
     rescue StandardError
       nil
     end
 
+    # @return [Boolean] whether {#close} has been called.
+    def closed?
+      @closed
+    end
+
     private
 
-    def execute(plan)
-      operation_id = SecureRandom.uuid
-      req = Proto::ExecutePlanRequest.new(
+    def derive(session_id)
+      SparkConnectClient.new(
+        @channel_builder.with_session_id(session_id),
+        max_retries: @max_retries, retry_base_delay: @retry_base_delay, max_retry_delay: @max_retry_delay,
+        rpc_deadlines: @rpc_deadlines
+      )
+    end
+
+    def call_options(rpc)
+      opts = { metadata: @metadata }
+      timeout = @rpc_deadlines&.public_send(rpc)
+      opts[:deadline] = Time.now + timeout if timeout
+      opts
+    end
+
+    def execute_request(plan, operation_id)
+      Proto::ExecutePlanRequest.new(
         session_id: @session_id,
         user_context: @user_context,
         operation_id: operation_id,
@@ -170,17 +333,33 @@ module SparkConnect
         client_type: @client_type,
         tags: @tags
       )
+    end
 
-      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)
+    def execute(plan)
+      operation_id = SecureRandom.uuid
+      req = execute_request(plan, operation_id)
+
+      # Build the accumulator *inside* the retry block so that a mid-stream
+      # failure (which restarts the gRPC stream from the beginning) starts from
+      # a clean slate. Accumulating into a result created outside the block
+      # would re-append already-seen batches and duplicate rows on retry.
+      progress = ProgressTracker.new(@progress_handlers, operation_id)
+      begin
+        with_retries do
+          result = ExecuteResult.new([], nil, nil, [], nil, 0)
+          result.pipeline_events = []
+          result.operation_id = operation_id
+          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?
+            progress.update(resp.execution_progress) if resp.execution_progress
+            accumulate(result, resp)
+          end
+          result
         end
+      ensure
+        progress.finish
       end
-      result
     end
 
     def accumulate(result, resp)
@@ -260,5 +439,46 @@ module SparkConnect
     def grpc_code_name(error)
       GRPC::Core::StatusCodes.constants.find { |c| GRPC::Core::StatusCodes.const_get(c) == error.code }&.to_s
     end
+
+    # Tracks `ExecutionProgress` updates of one execution and notifies the
+    # registered progress handlers, like PySpark's `Progress` helper: on every
+    # update that reports tasks, and once more with `done: true` at the end.
+    # @api private
+    class ProgressTracker
+      def initialize(handlers, operation_id)
+        @handlers = handlers.dup
+        @operation_id = operation_id
+        @stages = []
+        @inflight = 0
+      end
+
+      # Record an `ExecutionProgress` update and notify the handlers.
+      def update(progress)
+        stages = progress.stages.map do |s|
+          StageInfo.new(stage_id: s.stage_id, num_tasks: s.num_tasks, num_completed_tasks: s.num_completed_tasks,
+                        num_bytes_read: s.input_bytes_read, done: s.done)
+        end
+        return if stages.sum(&:num_tasks).zero?
+
+        @stages = stages
+        @inflight = progress.num_inflight_tasks
+        notify(false)
+      end
+
+      # Notify the handlers that the execution finished.
+      def finish
+        notify(true)
+      end
+
+      private
+
+      def notify(done)
+        @handlers.each do |h|
+          h.call(stages: @stages, inflight_tasks: @inflight, operation_id: @operation_id, done: done)
+        rescue StandardError => e
+          warn("spark-connect: progress handler raised #{e.class}: #{e.message}")
+        end
+      end
+    end
   end
 end