rubydb 0.1.0

data/lib/rubydb/client/client.rb

@@ -0,0 +1,304 @@
+# frozen_string_literal: true
+
+require "socket"
+require "json"
+require "time"
+require "digest"
+
+require_relative "connection"
+require_relative "connection_url"
+require_relative "result"
+require_relative "statement"
+require_relative "prepared_statement"
+require_relative "transaction"
+require_relative "connection_pool"
+require_relative "../protocol/protocol"
+require_relative "../protocol/message"
+
+module RubyDB
+  module Client
+    # Client - Main database client
+    class Client
+      attr_reader :config, :connection, :pool, :stats, :last_commit_ack
+
+      def initialize(config = {})
+        url = config[:url] || config["url"]
+        config = ConnectionURL.parse(url).merge(config.each_with_object({}) { |(key, value), result| result[key.to_sym] = value }) if url
+        @config = {
+          host: config[:host] || "localhost",
+          port: config[:port] || 7432,
+          username: config[:username] || "rubydb",
+          password: config[:password] || "",
+          database: config[:database] || "rubydb",
+          timeout: config[:timeout] || 30,
+          pool_size: config[:pool_size] || 5,
+          auto_connect: config.fetch(:auto_connect, true),
+          ssl: config[:ssl] || false,
+          compress: config[:compress] || false,
+          format: config[:format] || :json
+        }
+
+        @stats = {
+          queries_executed: 0,
+          queries_failed: 0,
+          transactions_started: 0,
+          transactions_committed: 0,
+          transactions_rolled_back: 0,
+          prepared_statements: 0,
+          bytes_sent: 0,
+          bytes_received: 0,
+          connection_time_ms: 0,
+          total_query_time_ms: 0,
+          avg_query_time_ms: 0
+        }
+
+        @lock = Mutex.new
+        @connected = false
+        @connection = nil
+        @pool = nil
+        @transaction = nil
+
+        # Initialize connection pool if pool_size > 1
+        if @config[:pool_size] > 1
+          @pool = ConnectionPool.new(@config)
+        end
+
+        # Auto-connect if configured
+        connect if @config[:auto_connect]
+      end
+
+      def connect
+        @lock.synchronize do
+          return if @connected
+
+          start_time = Time.now
+
+          begin
+            if @pool
+              @connection = @pool.get_connection
+            else
+              @connection = Connection.new(@config)
+              @connection.connect
+            end
+
+            @connected = true
+            @stats[:connection_time_ms] = (Time.now - start_time) * 1000
+
+            true
+          rescue => e
+            raise ClientError, "Connection failed: #{e.message}"
+          end
+        end
+      end
+
+      def disconnect
+        @lock.synchronize do
+          return unless @connected
+
+          if @pool
+            @pool.release_connection(@connection)
+          else
+            @connection.disconnect
+          end
+
+          @connection = nil
+          @connected = false
+          @transaction = nil
+
+          true
+        end
+      end
+
+      def connected?
+        @connected && @connection && @connection.connected?
+      end
+
+      def query(sql, params = [], timeout: nil)
+        @lock.synchronize do
+          ensure_connected
+
+          start_time = Time.now
+
+          begin
+            # Send query
+            deadline_at = timeout && (Time.now + Float(timeout)).iso8601(6)
+            response = @connection.send_query(sql, params, deadline_at: deadline_at)
+            @stats[:queries_executed] += 1
+
+            # Parse result
+            result = Result.new(response)
+
+            elapsed_ms = (Time.now - start_time) * 1000
+            @stats[:total_query_time_ms] += elapsed_ms
+            @stats[:avg_query_time_ms] = @stats[:total_query_time_ms] / @stats[:queries_executed] if @stats[:queries_executed] > 0
+
+            result
+
+          rescue => e
+            @stats[:queries_failed] += 1
+            raise
+          end
+        end
+      end
+
+      def query_async(sql, params = [], timeout: nil)
+        @lock.synchronize do
+          ensure_connected
+          deadline_at = timeout && (Time.now + Float(timeout)).iso8601(6)
+          handle = @connection.send_query_async(sql, params, deadline_at: deadline_at)
+          AsyncResult.new(handle)
+        end
+      end
+
+      class AsyncResult
+        attr_reader :request_id
+
+        def initialize(handle)
+          @handle = handle
+          @request_id = handle.request_id
+        end
+
+        def cancel
+          @handle.cancel
+        end
+
+        def wait(timeout = nil)
+          Result.new(@handle.wait(timeout))
+        end
+      end
+
+      def prepare(sql)
+        @lock.synchronize do
+          ensure_connected
+
+          begin
+            response = @connection.send_prepare(sql)
+            @stats[:prepared_statements] += 1
+
+            PreparedStatement.new(response[:statement_id], sql, self)
+          rescue => e
+            @stats[:queries_failed] += 1
+            raise
+          end
+        end
+      end
+
+      def execute(statement_id, params = [])
+        @lock.synchronize do
+          ensure_connected
+
+          start_time = Time.now
+
+          begin
+            response = @connection.send_execute(statement_id, params)
+            @stats[:queries_executed] += 1
+
+            result = Result.new(response)
+
+            elapsed_ms = (Time.now - start_time) * 1000
+            @stats[:total_query_time_ms] += elapsed_ms
+            @stats[:avg_query_time_ms] = @stats[:total_query_time_ms] / @stats[:queries_executed] if @stats[:queries_executed] > 0
+
+            result
+
+          rescue => e
+            @stats[:queries_failed] += 1
+            raise
+          end
+        end
+      end
+
+      def close_statement(statement_id)
+        @lock.synchronize do
+          ensure_connected
+          @connection.send_close(statement_id)
+        end
+      end
+
+      def begin_transaction(options = {})
+        @lock.synchronize do
+          ensure_connected
+
+          if @transaction && @transaction.active?
+            raise ClientError, "Transaction already active"
+          end
+
+          @connection.send_begin(options)
+          @transaction = Transaction.new(self, options)
+          @stats[:transactions_started] += 1
+
+          @transaction
+        end
+      end
+
+      def commit
+        @lock.synchronize do
+          ensure_connected
+
+          unless @transaction && @transaction.active?
+            raise ClientError, "No active transaction"
+          end
+
+          response = @connection.send_commit
+          unless response[:success] != false
+            raise ClientError, response[:error] || "Transaction commit failed"
+          end
+          @last_commit_ack = response[:commit_ack]
+          @last_commit_ack[:status] = @last_commit_ack[:status].to_sym if @last_commit_ack && @last_commit_ack[:status]
+          @transaction.mark_committed
+          @stats[:transactions_committed] += 1
+
+          @transaction = nil
+        end
+      end
+
+      def rollback
+        @lock.synchronize do
+          ensure_connected
+
+          unless @transaction && @transaction.active?
+            raise ClientError, "No active transaction"
+          end
+
+          @connection.send_rollback
+          @transaction.mark_rolled_back
+          @stats[:transactions_rolled_back] += 1
+
+          @transaction = nil
+        end
+      end
+
+      def in_transaction?
+        @transaction && @transaction.active?
+      end
+
+      def ping
+        @lock.synchronize do
+          ensure_connected
+          @connection.send_ping
+          true
+        end
+      end
+
+      def stats
+        @lock.synchronize do
+          @stats.merge({
+            connected: @connected,
+            host: @config[:host],
+            port: @config[:port],
+            database: @config[:database],
+            pool_size: @config[:pool_size],
+            pool_active: @pool ? @pool.active_count : 0
+          })
+        end
+      end
+
+      private
+
+      def ensure_connected
+        connect unless @connected
+        raise ClientError, "Not connected" unless @connected
+      end
+    end
+  end
+end

data/lib/rubydb/client/connection.rb

@@ -0,0 +1,415 @@
+# frozen_string_literal: true
+
+require "socket"
+require "json"
+require "time"
+require "monitor"
+require "openssl"
+require "base64"
+require "securerandom"
+
+module RubyDB
+  module Client
+    # Connection - Client connection to database server
+    class Connection
+      class QueryHandle
+        attr_reader :request_id
+
+        def initialize(connection, request_id, thread)
+          @connection = connection
+          @request_id = request_id
+          @thread = thread
+        end
+
+        def cancel
+          @connection.cancel(@request_id)
+        end
+
+        def cancelled?
+          @cancelled == true
+        end
+
+        def wait(timeout = nil)
+          @thread.join(timeout)
+          raise ConnectionError, "Query did not finish within #{timeout} seconds" if @thread.alive?
+
+          @thread.value
+        end
+      end
+
+      attr_reader :config, :socket, :stats
+
+      def initialize(config = {})
+        @config = config
+        @socket = nil
+        @connected = false
+        @protocol = nil
+        @authenticated = false
+        @session_id = nil
+        @stats = {
+          bytes_sent: 0,
+          bytes_received: 0,
+          messages_sent: 0,
+          messages_received: 0,
+          errors: 0,
+          reconnects: 0
+        }
+        @lock = Monitor.new
+        @buffer = ""
+      end
+
+      def connect
+        @lock.synchronize do
+          return if @connected
+
+          begin
+            tcp_socket = TCPSocket.new(@config[:host], @config[:port])
+            ssl = @config[:ssl]
+            if ssl == true || ssl.is_a?(Hash) && (ssl[:enabled] || ssl["enabled"])
+              context = OpenSSL::SSL::SSLContext.new
+              context.verify_mode = if ssl[:verify_peer] || ssl["verify_peer"]
+                OpenSSL::SSL::VERIFY_PEER
+              else
+                OpenSSL::SSL::VERIFY_NONE
+              end
+              context.ca_file = ssl[:ca_file] || ssl["ca_file"] if ssl.is_a?(Hash) && (ssl[:ca_file] || ssl["ca_file"])
+              if ssl.is_a?(Hash) && (ssl[:cert_file] || ssl["cert_file"])
+                context.cert = OpenSSL::X509::Certificate.new(File.binread(ssl[:cert_file] || ssl["cert_file"]))
+              end
+              if ssl.is_a?(Hash) && (ssl[:key_file] || ssl["key_file"])
+                context.key = OpenSSL::PKey.read(File.binread(ssl[:key_file] || ssl["key_file"]))
+              end
+              if ssl.is_a?(Hash) && (ssl[:min_version] || ssl["min_version"])
+                context.min_version = ssl[:min_version] || ssl["min_version"]
+              end
+              @socket = OpenSSL::SSL::SSLSocket.new(tcp_socket, context)
+              @socket.hostname = @config[:host] if @socket.respond_to?(:hostname=)
+              @socket.connect
+            else
+              @socket = tcp_socket
+            end
+
+            # Set timeouts
+            if @config[:timeout]
+              @socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_RCVTIMEO, [@config[:timeout], 0].pack("l_2"))
+              @socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_SNDTIMEO, [@config[:timeout], 0].pack("l_2"))
+            end
+
+            # Initialize protocol
+            @protocol = Protocol::Protocol.new(
+              format: @config[:format] || :json,
+              compression: @config[:compress] || false,
+              encryption: @config[:ssl] || false
+            )
+
+            # Perform handshake
+            perform_handshake
+
+            @connected = true
+
+          rescue => e
+            raise ConnectionError, "Connection failed: #{e.message}"
+          end
+        end
+      end
+
+      def disconnect
+        @lock.synchronize do
+          return unless @connected
+
+          begin
+            # Send terminate message
+            send_message(Protocol::Message.new(Protocol::Message::TYPE_TERMINATE))
+          rescue
+            # Ignore errors on disconnect
+          end
+
+          @socket.close if @socket
+          @socket = nil
+          @connected = false
+          @authenticated = false
+        end
+      end
+
+      def connected?
+        @connected && @socket && !@socket.closed?
+      end
+
+      def send_query(sql, params = [], deadline_at: nil)
+        @lock.synchronize do
+          ensure_connected
+
+          payload = { sql: sql, params: params }
+          payload[:deadline_at] = deadline_at if deadline_at
+          message = Protocol::Message.new(
+            Protocol::Message::TYPE_QUERY,
+            payload
+          )
+
+          send_message(message)
+          unwrap_result(receive_message.payload)
+        end
+      end
+
+      # Starts a query and returns a handle that can issue a wire-level
+      # cancellation while the server is executing it. The same connection
+      # must not be used for another application request until wait returns.
+      def send_query_async(sql, params = [], deadline_at: nil)
+        ensure_connected
+        payload = { sql: sql, params: params }
+        payload[:deadline_at] = deadline_at if deadline_at
+        message = Protocol::Message.new(Protocol::Message::TYPE_QUERY, payload)
+        @lock.synchronize { send_message(message) }
+        thread = Thread.new do
+          loop do
+            response = receive_message
+            next if response.type == :cancel_response
+            next if response.type == :notice && response.payload[:target_request_id].to_s != message.id.to_s
+
+            break unwrap_result(response.payload.merge(request_id: message.id))
+          end
+        end
+        QueryHandle.new(self, message.id, thread)
+      end
+
+      def cancel(request_id)
+        ensure_connected
+        message = Protocol::Message.new(
+          Protocol::Message::TYPE_CANCEL,
+          { target_request_id: request_id.to_s }
+        )
+        @lock.synchronize { send_message(message) }
+        true
+      end
+
+      def send_prepare(sql)
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(
+            Protocol::Message::TYPE_PREPARE,
+            { sql: sql }
+          )
+
+          send_message(message)
+          receive_message.payload
+        end
+      end
+
+      def send_execute(statement_id, params = [])
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(
+            Protocol::Message::TYPE_EXECUTE,
+            { statement_id: statement_id, params: params }
+          )
+
+          send_message(message)
+          unwrap_result(receive_message.payload)
+        end
+      end
+
+      def send_close(statement_id)
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(
+            Protocol::Message::TYPE_CLOSE,
+            { statement_id: statement_id }
+          )
+
+          send_message(message)
+          receive_message.payload
+        end
+      end
+
+      def send_begin(options = {})
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(
+            Protocol::Message::TYPE_BEGIN,
+            options
+          )
+
+          send_message(message)
+          receive_message.payload
+        end
+      end
+
+      def send_commit
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(Protocol::Message::TYPE_COMMIT)
+          send_message(message)
+          receive_message.payload
+        end
+      end
+
+      def send_rollback
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(Protocol::Message::TYPE_ROLLBACK)
+          send_message(message)
+          receive_message.payload
+        end
+      end
+
+      def send_ping
+        @lock.synchronize do
+          ensure_connected
+
+          message = Protocol::Message.new(Protocol::Message::TYPE_PING)
+          send_message(message)
+          receive_message
+        end
+      end
+
+      def stats
+        @lock.synchronize do
+          @stats.merge({
+            connected: @connected,
+            authenticated: @authenticated,
+            host: @config[:host],
+            port: @config[:port]
+          })
+        end
+      end
+
+      private
+
+      def ensure_connected
+        unless @connected
+          connect
+        end
+      end
+
+      def send_message(message)
+        data = @protocol.encoder.encode(message)
+        @socket.write(data)
+        @socket.flush
+
+        @stats[:bytes_sent] += data.bytesize
+        @stats[:messages_sent] += 1
+      end
+
+      def receive_message
+        data = @socket.gets
+        raise EOFError, "server closed connection" unless data
+        @stats[:bytes_received] += data.bytesize
+        @stats[:messages_received] += 1
+
+        @protocol.decoder.decode(data, @protocol.encoder.format)
+      rescue => e
+        @stats[:errors] += 1
+        raise ConnectionError, "Receive failed: #{e.message}"
+      end
+
+      def unwrap_result(payload)
+        result = payload[:result]
+        return payload unless result.is_a?(Hash)
+
+        result.merge(success: payload[:success] != false)
+      end
+
+      def perform_handshake
+        # Send handshake request
+        handshake_msg = Protocol::Message.new(
+          Protocol::Message::TYPE_HANDSHAKE,
+          {
+            protocol_version: Protocol::ProtocolVersion.current,
+            client_name: "RubyDB Client",
+            client_version: RubyDB::VERSION,
+            username: @config[:username],
+            database: @config[:database]
+          }
+        )
+
+        send_message(handshake_msg)
+        response = receive_message
+
+        if response.type == Protocol::Message::TYPE_HANDSHAKE_RESPONSE
+          if response.payload[:success]
+            # Send authentication, including a real SCRAM proof when required.
+            auth_payload = {
+              username: @config[:username],
+              password: @config[:password],
+              auth_method: @config[:auth_method] || "none"
+            }
+            if response.payload[:default_auth].to_s == "scram-sha-256"
+              auth_payload.delete(:password)
+              scram_payload = scram_data(
+                @config[:username].to_s,
+                @config[:password].to_s,
+                response.payload[:challenge].to_s,
+                response.payload[:scram_salt].to_s,
+                response.payload[:scram_iterations].to_i
+              )
+              expected_server_signature = scram_payload.delete(:expected_server_signature)
+              auth_payload[:scram_data] = scram_payload
+            end
+            auth_msg = Protocol::Message.new(
+              Protocol::Message::TYPE_AUTHENTICATION_RESPONSE,
+              auth_payload
+            )
+
+            send_message(auth_msg)
+            auth_response = receive_message
+
+            if auth_response.type == Protocol::Message::TYPE_AUTHENTICATION && auth_response.payload[:success]
+              if expected_server_signature && auth_response.payload[:server_signature] != expected_server_signature
+                raise ConnectionError, "SCRAM server signature verification failed"
+              end
+              @authenticated = true
+              @session_id = auth_response.payload[:session_id]
+
+              # Complete capability negotiation before exposing the connection
+              # to callers. The server will not process application requests
+              # until it receives this synchronization frame.
+              sync_msg = Protocol::Message.new(
+                Protocol::Message::TYPE_SYNCHRONIZE,
+                { capabilities: Protocol::Capabilities.default_client.to_hash }
+              )
+              send_message(sync_msg)
+              ready_response = receive_message
+              unless ready_response.type == Protocol::Message::TYPE_READY_FOR_QUERY
+                raise ConnectionError, "Protocol synchronization failed"
+              end
+            else
+              raise ConnectionError, "Authentication failed"
+            end
+          else
+            raise ConnectionError, "Handshake failed: #{response.payload[:error]}"
+          end
+        else
+          raise ConnectionError, "Unexpected handshake response"
+        end
+      end
+
+      def scram_data(username, password, challenge, encoded_salt, iterations)
+        client_nonce = SecureRandom.hex(18)
+        client_first_bare = "n=#{username},r=#{client_nonce}"
+        server_nonce = client_nonce + challenge
+        client_final = "c=biws,r=#{server_nonce}"
+        salt = Base64.decode64(encoded_salt)
+        salted = OpenSSL::PKCS5.pbkdf2_hmac(password, salt, iterations, 32, OpenSSL::Digest::SHA256.new)
+        auth_message = [client_first_bare, "r=#{server_nonce},s=#{encoded_salt},i=#{iterations}", client_final].join(",")
+        client_key = OpenSSL::HMAC.digest("SHA256", salted, "Client Key")
+        stored_key = OpenSSL::Digest::SHA256.digest(client_key)
+        signature = OpenSSL::HMAC.digest("SHA256", stored_key, auth_message)
+        proof = client_key.bytes.each_with_index.map { |byte, index| byte ^ signature.getbyte(index) }.pack("C*")
+        server_key = OpenSSL::HMAC.digest("SHA256", salted, "Server Key")
+        {
+          client_first_bare: client_first_bare,
+          client_final_without_proof: client_final,
+          client_nonce: client_nonce,
+          proof: Base64.strict_encode64(proof),
+          expected_server_signature: Base64.strict_encode64(OpenSSL::HMAC.digest("SHA256", server_key, auth_message))
+        }
+      end
+    end
+  end
+end