chalk_ruby 0.2.3 → 0.2.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: de352344a0d00f988be587d33f4e1d181256c0a075832fa4daa804639af83f9d
-  data.tar.gz: d4bde660475858f5eeb7b6bdfb08bdabacacbc36c92962007e3f1acf8f929dff
+  metadata.gz: cd71abeb7b4278848fe875e49df4ebfcc5e944ade6aa19de9e1bf5c02d9c744e
+  data.tar.gz: 2b38fccf00540a8387a0e11f898efe088ca7026edd2a9f95eb4d00f283c68bcc
 SHA512:
-  metadata.gz: a811e70301fadceb6fd1e0822e51dd60f3a297f1ef1afe1e20919824b877a6bcd479dc3233b9a326f08b1b5028569cd42e42b9778ac29fde073349cb7f3d4e7a
-  data.tar.gz: f45d5ba28647765f57a527c1602053f7daa7aad04b72f17f594f68bf1ef0f901453a35d9a87525a66b99d1b10aace71cd14936b8699a4efe7943db4911045d58
+  metadata.gz: 103633662a2281978355ff8f91962732e5bdfe07b567583bbe8d67ca69a8e572ec98ac12f5cccda388f670db258ac39d0fefd086c95e49449c5dd0e79adfcf2d
+  data.tar.gz: 5dc1ccd3f29ae918930eade977c3664a1cd129739694ccd168c3941d2d46703619b08c79ac4f0895cbb3a4e8e4179f1f366bf41c38d1601f0359a3e099397d96

data/chalk_ruby.gemspec

@@ -38,10 +38,11 @@ Gem::Specification.new do |spec|
 
   spec.add_dependency 'faraday', ['>= 0.15', '< 3']
   spec.add_dependency 'faraday-net_http_persistent', ['>= 0.15', '< 3']
-  spec.add_dependency 'grpc', ['>=1.64.3', '< 2']
+  spec.add_dependency 'grpc', ['>=1.68.1', '< 2']
 
   spec.add_dependency 'multi_json', '~> 1.0'
   spec.add_dependency 'net-http-persistent'
+  spec.add_dependency 'red-arrow', '~> 18.0.0'
 
   spec.add_development_dependency 'httpclient'
   spec.add_development_dependency 'm'

data/lib/chalk_ruby/config.rb

@@ -16,6 +16,7 @@ module ChalkRuby
                   :query_timeout,
                   :api_timeout,
                   :connect_timeout,
+                  :query_service_root_ca_path,
                   :additional_headers
 
     # Creates a new ChalkRuby::Config object for use with ChalkRuby::Client.
@@ -45,18 +46,22 @@ module ChalkRuby
     # @option options [Float?] :connect_timeout
     #   The timeout for connect operations (in seconds).
     #
+    # @option options [String?] :query_service_root_ca_path
+    #   The root ca to use for trusting SSL verifying connections to the query server in grpc.
+    #
     # @option options [Hash<String, String>?] :additional_headers
     #  Additional headers to be sent with every request. Typically not required.
     #
     def initialize(opts = {})
-      @client_id          = opts[:client_id] || ENV['CHALK_CLIENT_ID']
-      @client_secret      = opts[:client_secret] || ENV['CHALK_CLIENT_SECRET']
-      @environment        = opts[:environment] || ENV['CHALK_ACTIVE_ENVIRONMENT']
-      @query_server       = opts[:query_server] || ENV['CHALK_QUERY_SERVER'] || Defaults::QUERY_SERVER
-      @api_server         = opts[:api_server] || ENV['CHALK_API_SERVER'] || Defaults::API_SERVER
-      @query_timeout      = opts[:query_timeout] || Defaults::QUERY_TIMEOUT
-      @api_timeout        = opts[:api_timeout] || Defaults::API_TIMEOUT
-      @connect_timeout    = opts[:connect_timeout] || Defaults::CONNECT_TIMEOUT
+      @client_id                        = opts[:client_id] || ENV['CHALK_CLIENT_ID']
+      @client_secret                    = opts[:client_secret] || ENV['CHALK_CLIENT_SECRET']
+      @environment                      = opts[:environment] || ENV['CHALK_ACTIVE_ENVIRONMENT']
+      @query_server                     = opts[:query_server] || ENV['CHALK_QUERY_SERVER'] || Defaults::QUERY_SERVER
+      @api_server                       = opts[:api_server] || ENV['CHALK_API_SERVER'] || Defaults::API_SERVER
+      @query_timeout                    = opts[:query_timeout] || Defaults::QUERY_TIMEOUT
+      @api_timeout                      = opts[:api_timeout] || Defaults::API_TIMEOUT
+      @connect_timeout                  = opts[:connect_timeout] || Defaults::CONNECT_TIMEOUT
+      @query_service_root_ca_path       = opts[:query_service_root_ca_path] || Defaults::ROOT_CA_PATH
       @additional_headers = opts[:additional_headers] || {}
 
       raise ChalkError, 'No Client ID provided, please set :client_id' if @client_id.nil?

data/lib/chalk_ruby/defaults.rb

@@ -32,6 +32,10 @@ module ChalkRuby
     API_TIMEOUT     = 300
     QUERY_TIMEOUT   = 1200
 
+    # GRPC
+    # ----------------------------------------
+    ROOT_CA_PATH = nil
+
     WAIT_TASK_DEFAULT_TIME_BEFORE_RETRY = 100
   end
 end

data/lib/chalk_ruby/grpc/auth_interceptor.rb

@@ -11,12 +11,13 @@ module ChalkRuby
         @client_secret = client_secret
         @environment_id = environment_id
         @token = nil
+        @token_expiry = nil
       end
 
       def request_response(request:, call:, method:, metadata:)
-        # If we haven't fetched a token yet or if you'd like to handle token expiration,
-        # this is where you'd refresh it. For now, let's assume a long-lived token.
-        if @token.nil?
+
+        # Leave 1 minute buffer before refreshing auth token
+        if @token.nil? || @token_expiry.nil? || Time.now >= (@token_expiry - 60)
           response = @auth_stub.get_token(
             Chalk::Server::V1::GetTokenRequest.new(
               client_id: @client_id,
@@ -24,6 +25,9 @@ module ChalkRuby
             )
           )
           @token = response.access_token
+
+          # expires_in assumes to be 'seconds'
+          @token_expiry = Time.now + response.expires_in
         end
 
 

data/lib/chalk_ruby/grpc_client.rb

@@ -146,6 +146,72 @@ module ChalkRuby
       query_service.ping(Chalk::Engine::V1::PingRequest.new(num: 1))
     end
 
+    def query_bulk(
+      input:,
+      output:,
+      now: nil,
+      staleness: nil,
+      context: nil,
+      response_options: nil,
+      body_type: nil,
+      timeout: nil
+    )
+      # Convert input to feather format
+      inputs_feather = to_feather(input)
+
+
+      r = Chalk::Common::V1::OnlineQueryBulkRequest.new(
+        inputs_feather: inputs_feather,
+        outputs: output.map { |o| Chalk::Common::V1::OutputExpr.new(feature_fqn: o) },
+        staleness: staleness || {},
+        context: context || Chalk::Common::V1::OnlineQueryContext.new,
+        response_options: response_options || Chalk::Common::V1::OnlineQueryResponseOptions.new,
+        body_type: body_type || :FEATHER_BODY_TYPE_UNSPECIFIED
+      )
+
+      if timeout.nil?
+        response = query_service.online_query_bulk(r)
+      else
+        response = query_service.online_query_bulk(r, deadline: Time.now + timeout)
+      end
+
+      output_data = nil
+
+      if (!response.scalars_data.nil?) and response.scalars_data.length > 0
+        buffer = Arrow::Buffer.new(response.scalars_data)
+
+        # Create a buffer reader
+        buffer_reader = Arrow::BufferInputStream.new(buffer)
+
+        # Create an IPC reader from the buffer reader
+        reader = Arrow::FeatherFileReader.new(buffer_reader)
+
+        # Read the table
+
+
+        output_data = []
+
+        table = reader.read
+
+
+        field_names = table.schema.fields.map(&:name)
+        table.each_record do |r|
+          row = {}
+          field_names.each do |f|
+            row[f] = r[f]
+          end
+
+          output_data << row
+        end
+      end
+
+      {
+        data: output_data,
+        errors: response.errors,
+        meta: response.response_meta
+      }
+    end
+
     def query(
       input:,
       output:,
@@ -183,18 +249,14 @@ module ChalkRuby
     end
 
     def get_token
-      Token.new(
-        api_server_request(
-          method: :post,
-          path: '/v1/oauth/token',
-          body: {
-            client_id: @config.client_id,
-            client_secret: @config.client_secret,
-            grant_type: 'client_credentials'
-          },
-          headers: get_unauthenticated_headers
+      response = @auth_stub.get_token(
+        Chalk::Server::V1::GetTokenRequest.new(
+          client_id: @config.client_id,
+          client_secret: @config.client_secret
         )
       )
+
+      response.access_token
     end
 
     def get_unauthenticated_headers
@@ -280,7 +342,14 @@ module ChalkRuby
 
     def query_service
       if @query_service.nil?
-        @query_service = Chalk::Engine::V1::QueryService::Stub.new(query_server_host, GRPC::Core::ChannelCredentials.new(), interceptors: [engine_interceptor])
+        channel_creds =
+          if @config.query_service_root_ca_path.nil?
+            GRPC::Core::ChannelCredentials.new
+          else
+            GRPC::Core::ChannelCredentials.new(File.read(@config.query_service_root_ca_path))
+          end
+
+        @query_service = Chalk::Engine::V1::QueryService::Stub.new(query_server_host, channel_creds, interceptors: [engine_interceptor])
       end
 
       @query_service
@@ -342,7 +411,29 @@ module ChalkRuby
         ::Google::Protobuf::Value.new(list_value: list_value)
       else
         raise "Unsupported type: #{value.class}"
+        end
       end
+
+      private
+
+    def to_feather(input_hash)
+      require 'arrow'
+
+      # Ensure all values in the input hash are arrays
+      array_input_hash = input_hash.transform_values { |v| v.is_a?(Array) ? v : [v] }
+
+      # Create a table from the transformed input hash
+      table = Arrow::Table.new(array_input_hash)
+
+      # Write to feather format in memory
+      buffer = Arrow::ResizableBuffer.new(0)
+
+      output = Arrow::BufferOutputStream.new(buffer)
+
+      table.write_as_feather(output)
+
+      # Remove trailing null bytes from the resizable buffer; the buffer is 512 aligned
+      buffer.data.to_s.b.gsub(/ARROW1(.*)ARROW1.*/m) {"ARROW1#{$1}ARROW1"}
+    end
     end
   end
-end

data/lib/chalk_ruby/version.rb

@@ -1,3 +1,3 @@
 module ChalkRuby
-  VERSION = '0.2.3'.freeze
+  VERSION = '0.2.5'.freeze
 end

data/test/chalk_ruby/integration/client_test.rb

@@ -3,7 +3,6 @@ require 'rspec/autorun'
 require 'chalk_ruby/client'
 require 'chalk_ruby/error'
 
-
 CLIENT_ID     = ''
 CLIENT_SECRET = ''
 

data/test/chalk_ruby/integration/grpc_client_test.rb

@@ -0,0 +1,64 @@
+$LOAD_PATH.unshift File.expand_path('../../../../lib', __FILE__)
+# require 'rspec'
+# require 'chalk_ruby'
+# require 'chalk_ruby'
+
+require 'date'
+require 'rspec/autorun'
+require 'chalk_ruby/grpc_client'
+require 'chalk_ruby/grpc/auth_interceptor'
+require 'chalk_ruby/error'
+require 'chalk_ruby/protos/chalk/server/v1/auth_pb'
+require 'chalk_ruby/protos/chalk/server/v1/auth_services_pb'
+require 'chalk_ruby/protos/chalk/engine/v1/query_server_services_pb'
+require 'arrow'
+
+
+
+
+RSpec.describe ChalkRuby::GrpcClient do
+  describe '#query' do
+    let(:client) do
+      ChalkRuby::GrpcClient.new(
+        ChalkRuby::Config.new(
+          query_server: "standard-gke.chalk-develop.gcp.chalk.ai",
+          api_server: "api.staging.chalk.ai:443",
+          client_id: CLIENT_ID,
+          client_secret: CLIENT_SECRET,
+          environment: "tmnmc9beyujew",
+          # api_timeout: 0.6, # seconds
+          # connect_timeout: 0.3, # seconds
+          # query_service_root_ca_path: "/Users/andrew/found_ca.pem" # path to the root ca for chalkai.internal.found.app,
+        )
+      )
+    end
+
+    # it 'can perform queries' do
+    #   response = client.query(
+    #     input: { 'business.id': 1 },
+    #     output: %w(business.id)
+    #   )
+    #
+    #   expect(response).not_to be_nil
+    #
+    #   puts response
+    #   # The response should be a OnlineQueryBulkResponse
+    #   # expect(response).to be_a(Chalk::Common::V1::OnlineQueryBulkResponse)
+    # end
+
+    it 'can perform bulk queries' do
+      response = client.query_bulk(
+        input: { 'user.id': 1 },
+        output: %w(user.id user.socure_score)
+      )
+
+      expect(response).not_to be_nil
+      # The response should be a OnlineQueryBulkResponse
+      # expect(response).to be_a(Chalk::Common::V1::OnlineQueryBulkResponse)
+
+      # puts response.scalars_data
+
+      puts response
+    end
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: chalk_ruby
 version: !ruby/object:Gem::Version
-  version: 0.2.3
+  version: 0.2.5
 platform: ruby
 authors:
 - Chalk AI, Inc.
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2025-02-23 00:00:00.000000000 Z
+date: 2025-03-06 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bundler
@@ -98,7 +98,7 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 1.64.3
+        version: 1.68.1
     - - "<"
       - !ruby/object:Gem::Version
         version: '2'
@@ -108,7 +108,7 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 1.64.3
+        version: 1.68.1
     - - "<"
       - !ruby/object:Gem::Version
         version: '2'
@@ -140,6 +140,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: red-arrow
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 18.0.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 18.0.0
 - !ruby/object:Gem::Dependency
   name: httpclient
   requirement: !ruby/object:Gem::Requirement
@@ -348,6 +362,7 @@ files:
 - sig/chalk_ruby/token.rbs
 - sig/chalk_ruby/versions.rbs
 - test/chalk_ruby/integration/client_test.rb
+- test/chalk_ruby/integration/grpc_client_test.rb
 - test/chalk_ruby/test_helper.rb
 homepage: https://github.com/chalk-ai/chalk-ruby
 licenses:
@@ -377,4 +392,5 @@ specification_version: 4
 summary: A simple Ruby client for Chalk
 test_files:
 - test/chalk_ruby/integration/client_test.rb
+- test/chalk_ruby/integration/grpc_client_test.rb
 - test/chalk_ruby/test_helper.rb