chalk_ruby 0.2.4 → 0.2.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: cfe57e2b3d1f09eaa9dece4b777788941b347319bca50b3242c52070287aa86f
-  data.tar.gz: '039220d201af61aff4e19fc73ea242a797b8681f87dc449e13187996eae653e6'
+  metadata.gz: cd71abeb7b4278848fe875e49df4ebfcc5e944ade6aa19de9e1bf5c02d9c744e
+  data.tar.gz: 2b38fccf00540a8387a0e11f898efe088ca7026edd2a9f95eb4d00f283c68bcc
 SHA512:
-  metadata.gz: cc1d7cbfcdb2175ee61773918800270f963b136fea24573b06b97a8b22225f5488c9fed440755d964f32577822d61021122469c75c65b7e5271b19cccec03489
-  data.tar.gz: 79b8d14bbf904286f5589542dc85a0ff471b7e79ad675860893bd3e7fdaf9ceb3736a57773ed8bc78da9709a41df34dc6d4badcedad6d0c427a851dfb029aad0
+  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/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:,
@@ -345,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.4'.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.4
+  version: 0.2.5
 platform: ruby
 authors:
 - Chalk AI, Inc.
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2025-02-24 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