statelydb 0.12.2 → 0.14.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 16704510255d3248c18f2f26a63e2e94debe5207e6bdb820b8be999a44aeb486
-  data.tar.gz: 981f9a3030231adfcf253dbde5c84bb54232d64c4cec169a630d8c085194cc0e
+  metadata.gz: a723737fdce6e8ad2bec8937cfd79f6d968df0fc3e9e747c64e18af9ce769601
+  data.tar.gz: 1f0279a8b9507f4cffacc360ab82e2289471a40de39bd3891f0c7db5c63877c6
 SHA512:
-  metadata.gz: 647935dfc070e7b871933d9f75d24c7b6e651cd95005b374a65c5af7e1d0f3271259c5bf585350b3c9065fef8e33c10e86212bb3500d87ef0647ed77cf9bb79d
-  data.tar.gz: a5857dbb91d7f4f91e2d3518d1e0fcdaf8104e5e26b792440207c18f4c18f3d7e0a639c2d090c9ef95cf25e623b00ee8d708956d8989993a54e4a44b8f5f0947
+  metadata.gz: 815835070a91366b18948afe145eb548ae4f6ecbc6559ed3b8193288cd1084b7ac1a416264b0769afea93f4a8c01e17b0474facb46fe90460b0b3db803466428
+  data.tar.gz: e0c75f31911432f2544c4097fbc0c7a1615624257e80495aef9e5b3a0dddc92dd9341683d24416c4a32ab19b8bd34a2d79abb577781d75ee41b55a02f2e2e11c

data/lib/api/auth/get_auth_token_pb.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: auth/get_auth_token.proto
+
+require 'google/protobuf'
+
+
+descriptor_data = "\n\x19\x61uth/get_auth_token.proto\x12\x0cstately.auth\"B\n\x13GetAuthTokenRequest\x12\x1f\n\naccess_key\x18\x01 \x01(\tH\x00R\taccessKeyB\n\n\x08identity\"W\n\x14GetAuthTokenResponse\x12\x1d\n\nauth_token\x18\x01 \x01(\tR\tauthToken\x12 \n\x0c\x65xpires_in_s\x18\x02 \x01(\x04R\nexpiresInSBv\n\x10\x63om.stately.authB\x11GetAuthTokenProtoP\x01\xa2\x02\x03SAX\xaa\x02\x0cStately.Auth\xca\x02\x0cStately\\Auth\xe2\x02\x18Stately\\Auth\\GPBMetadata\xea\x02\rStately::Authb\x06proto3"
+
+pool = Google::Protobuf::DescriptorPool.generated_pool
+pool.add_serialized_file(descriptor_data)
+
+module Stately
+  module Auth
+    GetAuthTokenRequest = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("stately.auth.GetAuthTokenRequest").msgclass
+    GetAuthTokenResponse = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("stately.auth.GetAuthTokenResponse").msgclass
+  end
+end

data/lib/api/auth/service_pb.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: auth/service.proto
+
+require 'google/protobuf'
+
+require 'api/auth/get_auth_token_pb'
+
+
+descriptor_data = "\n\x12\x61uth/service.proto\x12\x0cstately.auth\x1a\x19\x61uth/get_auth_token.proto2i\n\x0b\x41uthService\x12Z\n\x0cGetAuthToken\x12!.stately.auth.GetAuthTokenRequest\x1a\".stately.auth.GetAuthTokenResponse\"\x03\x90\x02\x01\x42q\n\x10\x63om.stately.authB\x0cServiceProtoP\x01\xa2\x02\x03SAX\xaa\x02\x0cStately.Auth\xca\x02\x0cStately\\Auth\xe2\x02\x18Stately\\Auth\\GPBMetadata\xea\x02\rStately::Authb\x06proto3"
+
+pool = Google::Protobuf::DescriptorPool.generated_pool
+pool.add_serialized_file(descriptor_data)
+
+module Stately
+  module Auth
+  end
+end

data/lib/api/auth/service_services_pb.rb

@@ -0,0 +1,29 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# Source: auth/service.proto for package 'Stately.Auth'
+
+require 'grpc'
+require 'api/auth/service_pb'
+
+module Stately
+  module Auth
+    module AuthService
+      # AuthService is the service for vending access tokens used to connect to
+      # StatelyDB. This API is meant to be used from SDKs. Access Keys are created
+      # and managed from the stately.dbmanagement.UserService.
+      class Service
+
+        include ::GRPC::GenericService
+
+        self.marshal_class_method = :encode
+        self.unmarshal_class_method = :decode
+        self.service_name = 'stately.auth.AuthService'
+
+        # GetAuthToken returns a short-lived access token from some proof of
+        # identity. This operation will fail if the identity cannot be verified.
+        rpc :GetAuthToken, ::Stately::Auth::GetAuthTokenRequest, ::Stately::Auth::GetAuthTokenResponse
+      end
+
+      Stub = Service.rpc_stub_class
+    end
+  end
+end

data/lib/api/db/get_pb.rb

@@ -4,7 +4,7 @@
 
 require 'google/protobuf'
 
-require 'db/item_pb'
+require 'api/db/item_pb'
 
 
 descriptor_data = "\n\x0c\x64\x62/get.proto\x12\nstately.db\x1a\rdb/item.proto\"\xa3\x01\n\nGetRequest\x12\x19\n\x08store_id\x18\x01 \x01(\x04R\x07storeId\x12\'\n\x04gets\x18\x02 \x03(\x0b\x32\x13.stately.db.GetItemR\x04gets\x12\x1f\n\x0b\x61llow_stale\x18\x03 \x01(\x08R\nallowStale\x12*\n\x11schema_version_id\x18\x05 \x01(\rR\x0fschemaVersionIdJ\x04\x08\x04\x10\x05\"$\n\x07GetItem\x12\x19\n\x08key_path\x18\x01 \x01(\tR\x07keyPath\"5\n\x0bGetResponse\x12&\n\x05items\x18\x01 \x03(\x0b\x32\x10.stately.db.ItemR\x05itemsBc\n\x0e\x63om.stately.dbB\x08GetProtoP\x01\xa2\x02\x03SDX\xaa\x02\nStately.Db\xca\x02\nStately\\Db\xe2\x02\x16Stately\\Db\\GPBMetadata\xea\x02\x0bStately::Dbb\x06proto3"

data/lib/api/db/list_pb.rb

@@ -4,9 +4,9 @@
 
 require 'google/protobuf'
 
-require 'db/item_pb'
-require 'db/item_property_pb'
-require 'db/list_token_pb'
+require 'api/db/item_pb'
+require 'api/db/item_property_pb'
+require 'api/db/list_token_pb'
 
 
 descriptor_data = "\n\rdb/list.proto\x12\nstately.db\x1a\rdb/item.proto\x1a\x16\x64\x62/item_property.proto\x1a\x13\x64\x62/list_token.proto\"\xbd\x02\n\x10\x42\x65ginListRequest\x12\x19\n\x08store_id\x18\x01 \x01(\x04R\x07storeId\x12&\n\x0fkey_path_prefix\x18\x02 \x01(\tR\rkeyPathPrefix\x12\x14\n\x05limit\x18\x03 \x01(\rR\x05limit\x12\x1f\n\x0b\x61llow_stale\x18\x04 \x01(\x08R\nallowStale\x12\x41\n\rsort_property\x18\x05 \x01(\x0e\x32\x1c.stately.db.SortablePropertyR\x0csortProperty\x12@\n\x0esort_direction\x18\x06 \x01(\x0e\x32\x19.stately.db.SortDirectionR\rsortDirection\x12*\n\x11schema_version_id\x18\x07 \x01(\rR\x0fschemaVersionId\"\x8b\x01\n\x0cListResponse\x12\x37\n\x06result\x18\x01 \x01(\x0b\x32\x1d.stately.db.ListPartialResultH\x00R\x06result\x12\x36\n\x08\x66inished\x18\x02 \x01(\x0b\x32\x18.stately.db.ListFinishedH\x00R\x08\x66inishedB\n\n\x08response\";\n\x11ListPartialResult\x12&\n\x05items\x18\x01 \x03(\x0b\x32\x10.stately.db.ItemR\x05items\";\n\x0cListFinished\x12+\n\x05token\x18\x01 \x01(\x0b\x32\x15.stately.db.ListTokenR\x05token*8\n\rSortDirection\x12\x12\n\x0eSORT_ASCENDING\x10\x00\x12\x13\n\x0fSORT_DESCENDING\x10\x01\x42\x64\n\x0e\x63om.stately.dbB\tListProtoP\x01\xa2\x02\x03SDX\xaa\x02\nStately.Db\xca\x02\nStately\\Db\xe2\x02\x16Stately\\Db\\GPBMetadata\xea\x02\x0bStately::Dbb\x06proto3"

data/lib/api/db/put_pb.rb

@@ -4,7 +4,7 @@
 
 require 'google/protobuf'
 
-require 'db/item_pb'
+require 'api/db/item_pb'
 
 
 descriptor_data = "\n\x0c\x64\x62/put.proto\x12\nstately.db\x1a\rdb/item.proto\"|\n\nPutRequest\x12\x19\n\x08store_id\x18\x01 \x01(\x04R\x07storeId\x12\'\n\x04puts\x18\x02 \x03(\x0b\x32\x13.stately.db.PutItemR\x04puts\x12*\n\x11schema_version_id\x18\x03 \x01(\rR\x0fschemaVersionId\"U\n\x07PutItem\x12$\n\x04item\x18\x01 \x01(\x0b\x32\x10.stately.db.ItemR\x04item\x12$\n\x0emust_not_exist\x18\x03 \x01(\x08R\x0cmustNotExist\"5\n\x0bPutResponse\x12&\n\x05items\x18\x01 \x03(\x0b\x32\x10.stately.db.ItemR\x05itemsBc\n\x0e\x63om.stately.dbB\x08PutProtoP\x01\xa2\x02\x03SDX\xaa\x02\nStately.Db\xca\x02\nStately\\Db\xe2\x02\x16Stately\\Db\\GPBMetadata\xea\x02\x0bStately::Dbb\x06proto3"

data/lib/api/db/service_pb.rb

@@ -4,14 +4,14 @@
 
 require 'google/protobuf'
 
-require 'db/continue_list_pb'
-require 'db/delete_pb'
-require 'db/get_pb'
-require 'db/list_pb'
-require 'db/put_pb'
-require 'db/scan_root_paths_pb'
-require 'db/sync_list_pb'
-require 'db/transaction_pb'
+require 'api/db/continue_list_pb'
+require 'api/db/delete_pb'
+require 'api/db/get_pb'
+require 'api/db/list_pb'
+require 'api/db/put_pb'
+require 'api/db/scan_root_paths_pb'
+require 'api/db/sync_list_pb'
+require 'api/db/transaction_pb'
 
 
 descriptor_data = "\n\x10\x64\x62/service.proto\x12\nstately.db\x1a\x16\x64\x62/continue_list.proto\x1a\x0f\x64\x62/delete.proto\x1a\x0c\x64\x62/get.proto\x1a\rdb/list.proto\x1a\x0c\x64\x62/put.proto\x1a\x18\x64\x62/scan_root_paths.proto\x1a\x12\x64\x62/sync_list.proto\x1a\x14\x64\x62/transaction.proto2\xee\x04\n\x0f\x44\x61tabaseService\x12;\n\x03Put\x12\x16.stately.db.PutRequest\x1a\x17.stately.db.PutResponse\"\x03\x90\x02\x02\x12;\n\x03Get\x12\x16.stately.db.GetRequest\x1a\x17.stately.db.GetResponse\"\x03\x90\x02\x01\x12\x44\n\x06\x44\x65lete\x12\x19.stately.db.DeleteRequest\x1a\x1a.stately.db.DeleteResponse\"\x03\x90\x02\x02\x12J\n\tBeginList\x12\x1c.stately.db.BeginListRequest\x1a\x18.stately.db.ListResponse\"\x03\x90\x02\x01\x30\x01\x12P\n\x0c\x43ontinueList\x12\x1f.stately.db.ContinueListRequest\x1a\x18.stately.db.ListResponse\"\x03\x90\x02\x01\x30\x01\x12L\n\x08SyncList\x12\x1b.stately.db.SyncListRequest\x1a\x1c.stately.db.SyncListResponse\"\x03\x90\x02\x01\x30\x01\x12T\n\x0bTransaction\x12\x1e.stately.db.TransactionRequest\x1a\x1f.stately.db.TransactionResponse\"\x00(\x01\x30\x01\x12Y\n\rScanRootPaths\x12 .stately.db.ScanRootPathsRequest\x1a!.stately.db.ScanRootPathsResponse\"\x03\x90\x02\x01\x42g\n\x0e\x63om.stately.dbB\x0cServiceProtoP\x01\xa2\x02\x03SDX\xaa\x02\nStately.Db\xca\x02\nStately\\Db\xe2\x02\x16Stately\\Db\\GPBMetadata\xea\x02\x0bStately::Dbb\x06proto3"

data/lib/api/db/service_services_pb.rb

@@ -2,7 +2,7 @@
 # Source: db/service.proto for package 'Stately.Db'
 
 require 'grpc'
-require 'db/service_pb'
+require 'api/db/service_pb'
 
 module Stately
   module Db

data/lib/api/db/sync_list_pb.rb

@@ -4,8 +4,8 @@
 
 require 'google/protobuf'
 
-require 'db/item_pb'
-require 'db/list_pb'
+require 'api/db/item_pb'
+require 'api/db/list_pb'
 
 
 descriptor_data = "\n\x12\x64\x62/sync_list.proto\x12\nstately.db\x1a\rdb/item.proto\x1a\rdb/list.proto\"\\\n\x0fSyncListRequest\x12\x1d\n\ntoken_data\x18\x01 \x01(\x0cR\ttokenData\x12*\n\x11schema_version_id\x18\x05 \x01(\rR\x0fschemaVersionId\"\xc8\x01\n\x10SyncListResponse\x12\x31\n\x05reset\x18\x01 \x01(\x0b\x32\x19.stately.db.SyncListResetH\x00R\x05reset\x12=\n\x06result\x18\x02 \x01(\x0b\x32#.stately.db.SyncListPartialResponseH\x00R\x06result\x12\x36\n\x08\x66inished\x18\x03 \x01(\x0b\x32\x18.stately.db.ListFinishedH\x00R\x08\x66inishedB\n\n\x08response\"\x0f\n\rSyncListReset\"\xdf\x01\n\x17SyncListPartialResponse\x12\x35\n\rchanged_items\x18\x01 \x03(\x0b\x32\x10.stately.db.ItemR\x0c\x63hangedItems\x12<\n\rdeleted_items\x18\x02 \x03(\x0b\x32\x17.stately.db.DeletedItemR\x0c\x64\x65letedItems\x12O\n%updated_item_keys_outside_list_window\x18\x03 \x03(\tR updatedItemKeysOutsideListWindow\"(\n\x0b\x44\x65letedItem\x12\x19\n\x08key_path\x18\x01 \x01(\tR\x07keyPathBh\n\x0e\x63om.stately.dbB\rSyncListProtoP\x01\xa2\x02\x03SDX\xaa\x02\nStately.Db\xca\x02\nStately\\Db\xe2\x02\x16Stately\\Db\\GPBMetadata\xea\x02\x0bStately::Dbb\x06proto3"

data/lib/api/db/transaction_pb.rb

@@ -4,13 +4,13 @@
 
 require 'google/protobuf'
 
-require 'db/continue_list_pb'
-require 'db/delete_pb'
-require 'db/get_pb'
-require 'db/item_pb'
-require 'db/item_property_pb'
-require 'db/list_pb'
-require 'db/put_pb'
+require 'api/db/continue_list_pb'
+require 'api/db/delete_pb'
+require 'api/db/get_pb'
+require 'api/db/item_pb'
+require 'api/db/item_property_pb'
+require 'api/db/list_pb'
+require 'api/db/put_pb'
 require 'google/protobuf/empty_pb'
 
 

data/lib/common/auth/{auth0_token_provider.rb → auth_token_provider.rb}

@@ -6,8 +6,10 @@ require "async/http/internet"
 require "async/semaphore"
 require "json"
 require "logger"
-require "weakref"
+require "grpc"
 require_relative "token_provider"
+require_relative "token_fetcher"
+require_relative "../../error"
 
 LOGGER = Logger.new($stdout)
 LOGGER.level = Logger::WARN
@@ -22,20 +24,22 @@ module StatelyDB
       # which vends tokens from auth0 with the given client_id and client_secret.
       # It will default to using the values of `STATELY_CLIENT_ID` and `STATELY_CLIENT_SECRET` if
       # no credentials are explicitly passed and will throw an error if none are found.
-      class Auth0TokenProvider < TokenProvider
+      class AuthTokenProvider < TokenProvider
         # @param [String] origin The origin of the OAuth server
         # @param [String] audience The OAuth Audience for the token
         # @param [String] client_secret The StatelyDB client secret credential
         # @param [String] client_id The StatelyDB client ID credential
+        # @param [String] access_key The StatelyDB access key credential
         def initialize(
           origin: "https://oauth.stately.cloud",
           audience: "api.stately.cloud",
-          client_secret: ENV.fetch("STATELY_CLIENT_SECRET"),
-          client_id: ENV.fetch("STATELY_CLIENT_ID")
+          client_secret: ENV.fetch("STATELY_CLIENT_SECRET", nil),
+          client_id: ENV.fetch("STATELY_CLIENT_ID", nil),
+          access_key: ENV.fetch("STATELY_ACCESS_KEY", nil)
         )
           super()
           @actor = Async::Actor.new(Actor.new(origin: origin, audience: audience,
-                                              client_secret: client_secret, client_id: client_id))
+                                              client_secret: client_secret, client_id: client_id, access_key: access_key))
           # this initialization cannot happen in the constructor because it is async and must run on the event loop
           # which is not available in the constructor
           @actor.init
@@ -64,29 +68,41 @@ module StatelyDB
             origin:,
             audience:,
             client_secret:,
-            client_id:
+            client_id:,
+            access_key:
           )
             super()
-            @client = Async::HTTP::Client.new(Async::HTTP::Endpoint.parse(origin))
-            @client_id = client_id
-            @client_secret = client_secret
-            @audience = audience
 
-            @access_token = nil
-            @expires_at_unix_secs = nil
+            @token_fetcher = nil
+            if !access_key.nil?
+              @token_fetcher = StatelyDB::Common::Auth::StatelyAccessTokenFetcher.new(origin: origin, access_key: access_key)
+            elsif !client_secret.nil? && !client_id.nil?
+              @token_fetcher = StatelyDB::Common::Auth::Auth0TokenFetcher.new(origin: origin, audience: audience,
+                                                                              client_secret: client_secret, client_id: client_id)
+            else
+              raise StatelyDB::Error.new("unable to find client credentials in STATELY_ACCESS_KEY or STATELY_CLIENT_ID and " \
+                                         "STATELY_CLIENT_SECRET environment variables. Either pass your credentials in " \
+                                         "explicitly or set these environment variables",
+                                         code: GRPC::Core::StatusCodes::UNAUTHENTICATED,
+                                         stately_code: "Unauthenticated")
+            end
+
+            @token_state = nil
             @pending_refresh = nil
           end
 
           # Initialize the actor. This runs on the actor thread which means
           # we can dispatch async operations here.
           def init
+            # disable the async lib logger. We do our own error handling and propagation
+            Console.logger.disable(Async::Task)
             refresh_token
           end
 
           # Close the token provider and kill any background operations
           def close
             @scheduled&.stop
-            @client&.close
+            @token_fetcher&.close
           end
 
           # Get the current access token
@@ -94,8 +110,7 @@ module StatelyDB
           # @return [String] The current access token
           def get_token(force: false)
             if force
-              @access_token = nil
-              @expires_at_unix_secs = nil
+              @token_state = nil
             else
               token, ok = valid_access_token
               return token if ok
@@ -107,11 +122,10 @@ module StatelyDB
           # Get the current access token and whether it is valid
           # @return [Array] The current access token and whether it is valid
           def valid_access_token
-            return "", false if @access_token.nil?
-            return "", false if @expires_at_unix_secs.nil?
-            return "", false if @expires_at_unix_secs < Time.now.to_i
+            return "", false if @token_state.nil?
+            return "", false if @token_state.expires_at_unix_secs < Time.now.to_i
 
-            [@access_token, true]
+            [@token_state.token, true]
           end
 
           # Refresh the access token
@@ -151,19 +165,16 @@ module StatelyDB
           # @return [String] The new access token
           def refresh_token_impl
             Sync do
-              resp_data = make_auth0_request
-
-              new_access_token = resp_data["access_token"]
-              new_expires_in_secs = resp_data["expires_in"]
+              token_result = @token_fetcher.fetch
+              new_expires_in_secs = token_result.expires_in_secs
               new_expires_at_unix_secs = Time.now.to_i + new_expires_in_secs
-              if @expires_at_unix_secs.nil? || new_expires_at_unix_secs > @expires_at_unix_secs
 
-                @access_token = new_access_token
-                @expires_at_unix_secs = new_expires_at_unix_secs
+              # only update the token state if the new expiry is later than the current one
+              if @token_state.nil? || new_expires_at_unix_secs > @token_state.expires_at_unix_secs
+                @token_state = TokenState.new(token: token_result.token, expires_at_unix_secs: new_expires_at_unix_secs)
               else
-
-                new_access_token = @access_token
-                new_expires_in_secs = @expires_at_unix_secs - Time.now.to_i
+                # otherwise use the existing expiry time for scheduling the refresh
+                new_expires_in_secs = @token_state.expires_at_unix_secs - Time.now.to_i
               end
 
               # Schedule a refresh of the token ahead of the expiry time
@@ -182,24 +193,21 @@ module StatelyDB
                 @scheduled = nil
               end
 
-              new_access_token
+              @token_state.token
             end
           end
+        end
 
-          def make_auth0_request
-            headers = [["content-type", "application/json"]]
-            body = JSON.dump({ "client_id" => @client_id, client_secret: @client_secret, audience: @audience,
-                               grant_type: DEFAULT_GRANT_TYPE })
-            Sync do
-              # TODO: Wrap this in a retry loop and parse errors like we
-              # do in the Go SDK.
-              response = @client.post("/oauth/token", headers, body)
-              raise "Auth request failed" if response.status != 200
-
-              JSON.parse(response.read)
-            ensure
-              response&.close
-            end
+        # Persistent state for the token provider
+        class TokenState
+          attr_reader :token, :expires_at_unix_secs
+
+          # Create a new TokenState
+          # @param [String] token The access token
+          # @param [Integer] expires_at_unix_secs The unix timestamp when the token expires
+          def initialize(token:, expires_at_unix_secs:)
+            @token = token
+            @expires_at_unix_secs = expires_at_unix_secs
           end
         end
       end

data/lib/common/auth/interceptor.rb

@@ -11,7 +11,7 @@ module StatelyDB
       class Interceptor < GRPC::ClientInterceptor
         # @param [TokenProvider] token_provider The token provider to use for authentication
         def initialize(
-          token_provider: Auth0TokenProvider.new
+          token_provider: AuthTokenProvider.new
         )
           super()
           @token_provider = token_provider

data/lib/common/auth/token_fetcher.rb

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require_relative "../net/conn"
+require_relative "../error_interceptor"
+require_relative "../../api/auth/service_services_pb"
+require "grpc"
+
+module StatelyDB
+  module Common
+    # A module for Stately Cloud auth code
+    module Auth
+      # Result from a token fetch operation
+      class TokenResult
+        attr_reader :token, :expires_in_secs
+
+        # Create a new TokenResult
+        # @param [String] token The access token
+        # @param [Integer] expires_in_secs The number of seconds until the token expires
+        def initialize(token:, expires_in_secs:)
+          @token = token
+          @expires_in_secs = expires_in_secs
+        end
+      end
+
+      # TokenFetcher is an abstract base class that should be extended
+      # for individual token fetcher implementations
+      class TokenFetcher
+        # Get the current access token
+        # @return [TokenResult] The fetched TokenResult
+        def fetch
+          raise "Not Implemented"
+        end
+
+        # Close the token provider and kill any background operations
+        def close
+          raise "Not Implemented"
+        end
+      end
+
+      # Auth0TokenFetcher is a TokenFetcher that fetches tokens from an Auth0 server
+      class Auth0TokenFetcher < TokenFetcher
+        # @param [String] origin The origin of the OAuth server
+        # @param [String] audience The OAuth Audience for the token
+        # @param [String] client_secret The StatelyDB client secret credential
+        # @param [String] client_id The StatelyDB client ID credential
+        def initialize(origin:, audience:, client_secret:, client_id:)
+          super()
+          @client = Async::HTTP::Client.new(Async::HTTP::Endpoint.parse(origin))
+          @audience = audience
+          @client_secret = client_secret
+          @client_id = client_id
+        end
+
+        # Fetch a new token from auth0
+        # @return [TokenResult] The fetched TokenResult
+        def fetch
+          headers = [["content-type", "application/json"]]
+          body = JSON.dump({ "client_id" => @client_id, client_secret: @client_secret, audience: @audience,
+                             grant_type: DEFAULT_GRANT_TYPE })
+          Sync do
+            # TODO: Wrap this in a retry loop and parse errors like we
+            # do in the Go SDK.
+            response = @client.post("/oauth/token", headers, body)
+            raise "Auth request failed" if response.status != 200
+
+            resp_data = JSON.parse(response.read)
+            TokenResult.new(token: resp_data["access_token"], expires_in_secs: resp_data["expires_in"])
+          ensure
+            response&.close
+          end
+        end
+
+        def close
+          @client&.close
+        end
+      end
+
+      # StatelyAccessTokenFetcher is a TokenFetcher that fetches tokens from the StatelyDB API
+      class StatelyAccessTokenFetcher < TokenFetcher
+        # @param [String] origin The origin of the OAuth server
+        # @param [String] access_key The StatelyDB access key credential
+        def initialize(origin:, access_key:)
+          super()
+          @access_key = access_key
+          @channel = Common::Net.new_channel(endpoint: origin)
+          error_interceptor = Common::ErrorInterceptor.new
+          @stub = Stately::Auth::AuthService::Stub.new(nil, nil, channel_override: @channel,
+                                                                 interceptors: [error_interceptor])
+        end
+
+        # Fetch a new token from the StatelyDB API
+        # @return [TokenResult] The fetched TokenResult
+        def fetch
+          resp = @stub.get_auth_token(Stately::Auth::GetAuthTokenRequest.new(access_key: @access_key))
+          TokenResult.new(token: resp.auth_token, expires_in_secs: resp.expires_in_s)
+        end
+
+        def close
+          @channel&.close
+        end
+      end
+    end
+  end
+end

data/lib/error.rb

@@ -1,10 +1,5 @@
 # frozen_string_literal: true
 
-# Add the pb dir to the LOAD_PATH because generated proto imports are not relative and
-# we don't want the protos polluting the main namespace.
-# Tracking here: https://github.com/grpc/grpc/issues/6164
-$LOAD_PATH.unshift "#{File.dirname(__FILE__)}/api"
-
 require "api/errors/error_details_pb"
 
 module StatelyDB

data/lib/statelydb.rb

@@ -1,12 +1,7 @@
 # frozen_string_literal: true
 
-# Add the pb dir to the LOAD_PATH because generated proto imports are not relative and
-# we don't want the protos polluting the main namespace.
-# Tracking here: https://github.com/grpc/grpc/issues/6164
-$LOAD_PATH.unshift "#{File.dirname(__FILE__)}/api"
-
 require "api/db/service_services_pb"
-require "common/auth/auth0_token_provider"
+require "common/auth/auth_token_provider"
 require "common/auth/interceptor"
 require "common/net/conn"
 require "common/error_interceptor"
@@ -35,7 +30,7 @@ module StatelyDB
     # @param region [String] the region to connect to.
     def initialize(store_id:,
                    schema:,
-                   token_provider: Common::Auth::Auth0TokenProvider.new,
+                   token_provider: Common::Auth::AuthTokenProvider.new,
                    endpoint: nil,
                    region: nil)
       if store_id.nil?
@@ -63,6 +58,7 @@ module StatelyDB
       @allow_stale = false
     end
 
+    # @return [void] nil
     def close
       @channel&.close
       @token_provider&.close

data/lib/token.rb

@@ -34,8 +34,7 @@ module StatelyDB
     end
 
     # Returns the schema version ID associated with the token.
-    def schema_version_id
-      !!@schema_version_id
-    end
+    # @return [Integer]
+    attr_reader :schema_version_id
   end
 end

data/lib/transaction/transaction.rb

@@ -254,7 +254,7 @@ module StatelyDB
       #
       # @param items [StatelyDB::Item, Array<StatelyDB::Item>] the items to store. Max
       # 50 items.
-      # @return [Array<StatelyDB::UUID, String, Integer, nil>] the ids of the items
+      # @return [Array<StatelyDB::UUID, String, Integer, NilClass>] the ids of the items
       #
       # @example
       #   results = client.data.transaction do |txn|

data/sig/grpc.rbs

@@ -0,0 +1,80 @@
+module GRPC
+  class ActiveCall
+  end
+
+  class Interceptor
+  end
+  class ClientInterceptor < Interceptor
+    # gRPC client unary interceptor
+    #
+    # _@param_ `request` — The request object
+    #
+    # _@param_ `call` — The active call object
+    #
+    # _@param_ `method` — The method being called
+    #
+    # _@param_ `metadata` — The metadata hash
+    #
+    # _@return_ — The response object
+    def request_response: (
+                            request: Object,
+                            call: GRPC::ActiveCall,
+                            method: Symbol,
+                            metadata: ::Hash[untyped, untyped]
+                          ) -> Object
+
+    # gRPC client streaming interceptor
+    #
+    # _@param_ `requests` — The list of requests
+    #
+    # _@param_ `call` — The active call object
+    #
+    # _@param_ `method` — The method being called
+    #
+    # _@param_ `metadata` — The metadata hash
+    #
+    # _@return_ — The response enumerator
+    def client_streamer: (
+                            requests: ::Enumerable[untyped],
+                            call: GRPC::ActiveCall,
+                            method: Symbol,
+                            metadata: ::Hash[untyped, untyped]
+                          ) -> ::Enumerator[untyped]
+
+    # gRPC  server streaming interceptor
+    #
+    # _@param_ `request` — The request object
+    #
+    # _@param_ `call` — The active call object
+    #
+    # _@param_ `method` — The method being called
+    #
+    # _@param_ `metadata` — The metadata hash
+    #
+    # _@return_ — The response enumerator
+    def server_streamer: (
+                            request: Object,
+                            call: GRPC::ActiveCall,
+                            method: Symbol,
+                            metadata: ::Hash[untyped, untyped]
+                          ) -> ::Enumerator[untyped]
+
+    # gRPC bidirectional streaming interceptor
+    #
+    # _@param_ `requests` — The list of requests
+    #
+    # _@param_ `call` — The active call object
+    #
+    # _@param_ `method` — The method being called
+    #
+    # _@param_ `metadata` — The metadata hash
+    #
+    # _@return_ — The response enumerator
+    def bidi_streamer: (
+                          requests: ::Enumerable[untyped],
+                          call: GRPC::ActiveCall,
+                          method: Symbol,
+                          metadata: ::Hash[untyped, untyped]
+                        ) -> ::Enumerator[untyped]
+  end
+end