Cache access token on disk (#85)

* cache access token on disk

* cleanup

* fix tests

* Update src/rest.jl

Co-authored-by: Jacob Quinn <quinn.jacobd@gmail.com>

* addressing PR comments

* Update src/rest.jl

Co-authored-by: Jacob Quinn <quinn.jacobd@gmail.com>

* fix read access token issue && add token caching test

* init rai config folder

* missing shell env

Co-authored-by: Jacob Quinn <quinn.jacobd@gmail.com>

Author: NRHelmi

.github/actions/test/action.yml

@@ -38,6 +38,11 @@ runs:
     - uses: julia-actions/cache@v1
     - uses: julia-actions/julia-buildpkg@v1
 
+    # this folder is required to test access token caching
+    - name: Init rai config folder
+      run: mkdir -p ~/.rai
+      shell: bash
+
     - name: Test
       uses: julia-actions/julia-runtest@v1
       env:

src/creds.jl

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-using Dates: DateTime, Second
+using Dates: datetime2unix
 
 """
     AccessToken
@@ -23,21 +23,21 @@ but we can still print the secret if needed:
 
 Example:
 ```
-access_token.token
+t.access_token
 ````
 """
 struct AccessToken
-    token::String
+    access_token::String
     scope::String
     expires_in::Int  # seconds
-    created_on::DateTime
+    created_on::Float64
 end
 
 function Base.show(io::IO, t::AccessToken)
     print(
         io,
         "(",
-        isempty(t.token) ? "" : "$(t.token[1:3])...",
+        isempty(t.access_token) ? "" : "$(t.access_token[1:3])...",
         ", ", t.scope,
         ", ", t.expires_in,
         ", ", t.created_on,
@@ -46,8 +46,8 @@ function Base.show(io::IO, t::AccessToken)
 end
 
 function isexpired(access_token::AccessToken)::Bool
-    expires_on = access_token.created_on + Second(access_token.expires_in)
-    return expires_on - Second(5) < now() # anticipate token expiration by 5 seconds
+    expires_on = access_token.created_on + access_token.expires_in
+    return expires_on - 5 < datetime2unix(now()) # anticipate token expiration by 5 seconds
 end
 
 abstract type Credentials end
@@ -78,8 +78,8 @@ function Base.show(io::IO, c::ClientCredentials)
         io,
         "(",
         c.client_id,
-        c.client_secret == nothing ? "" : ", $(c.client_secret[1:3])...",
-        c.access_token == nothing ? "" : ", $(c.access_token)",
+        c.client_secret === nothing ? "" : ", $(c.client_secret[1:3])...",
+        c.access_token === nothing ? "" : ", $(c.access_token)",
         ", ",
         c.client_credentials_url,
         ")"

src/rest.jl

@@ -15,7 +15,7 @@
 # Low level HTTP interface to the RAI REST API. Handles authentication of
 # requests and other protocol level details.
 
-using Dates: now
+using Dates: now, datetime2unix
 import HTTP
 import JSON3
 
@@ -81,7 +81,64 @@ function get_access_token(ctx::Context, creds::ClientCredentials)::AccessToken
     opts = (redirect = false, retry_non_idempotent = true, connect_timeout = 30, readtimeout = 30, keepalive = true)
     rsp = HTTP.request("POST", url, h, body; opts...)
     data = JSON3.read(rsp.body)
-    return AccessToken(data.access_token, data.scope, data.expires_in, now())
+    return AccessToken(data.access_token, data.scope, data.expires_in, datetime2unix(now()))
+end
+
+# cache name
+function _cache_file()
+    return joinpath(homedir(), ".rai", "tokens.json")
+end
+
+# read oauth cache
+function _read_cache()
+    try
+        if isfile(_cache_file())
+            return copy(JSON3.read(read(_cache_file())))
+        else
+            return nothing
+        end
+    catch e
+        @warn e
+        return nothing
+    end
+end
+
+# Read access token from cache
+function _read_token_cache(creds::ClientCredentials)
+    try
+        cache = _read_cache()
+        cache === nothing && return nothing
+
+        if haskey(cache, Symbol(creds.client_id))
+            access_token = cache[Symbol(creds.client_id)]
+            return AccessToken(
+                access_token[:access_token],
+                access_token[:scope],
+                access_token[:expires_in],
+                access_token[:created_on],
+            )
+        else
+            return nothing
+        end
+    catch e
+        @warn e
+        return nothing
+    end
+end
+
+# Write access token to cache
+function _write_token_cache(creds::ClientCredentials)
+    try
+        cache = _read_cache()
+        if cache === nothing
+            cache = Dict(creds.client_id => creds.access_token)
+        else
+            cache[Symbol(creds.client_id)] = creds.access_token
+        end
+        write(_cache_file(), JSON3.write(cache))
+    catch e
+        @warn e
+    end
 end
 
 function _get_client_credentials_url(creds::ClientCredentials)
@@ -102,13 +159,19 @@ function _authenticate!(
     headers,
 )::Nothing
     if isnothing(creds.access_token)
-        creds.access_token = get_access_token(ctx, creds)
+        creds.access_token = _read_token_cache(creds)
+        if isnothing(creds.access_token)
+            creds.access_token = get_access_token(ctx, creds)
+            _write_token_cache(creds)
+        end
     end
 
     if isexpired(creds.access_token)
         creds.access_token = get_access_token(ctx, creds)
+        _write_token_cache(creds)
     end
-    push!(headers, "Authorization" => "Bearer $(creds.access_token.token)")
+
+    push!(headers, "Authorization" => "Bearer $(creds.access_token.access_token)")
     return nothing
 end
 

test/api.jl

@@ -5,7 +5,7 @@ using JSON3
 using Mocking
 using Dates
 using RAI.protocol
-using RAI: _poll_with_specified_overhead
+using RAI: _poll_with_specified_overhead, _write_token_cache, _read_token_cache
 
 using RAI: TransactionResponse
 
@@ -287,11 +287,24 @@ end
 end
 
 @testset "hide client secrets in repl" begin
-    access_token = AccessToken("abc_token", "run:transaction", 3600, DateTime("2022-08-12T17:49:51.365"))
+    access_token = AccessToken("abc_token", "run:transaction", 3600, datetime2unix(DateTime("2022-08-12T17:49:51.365")))
     creds = ClientCredentials("client_id", "xyz_client_secret", "https://login.relationalai.com/oauth/token")
     creds.access_token = access_token
 
     io = IOBuffer()
     show(io, creds)
-    @test String(take!(io)) === "(client_id, xyz..., (abc..., run:transaction, 3600, 2022-08-12T17:49:51.365), https://login.relationalai.com/oauth/token)"
+    @test String(take!(io)) === "(client_id, xyz..., (abc..., run:transaction, 3600, 1.660326591365e9), https://login.relationalai.com/oauth/token)"
+end
+
+@testset "read write access token to cache" begin
+    access_token = AccessToken("abc_token", "run:transaction", 3600, datetime2unix(DateTime("2022-08-12T17:49:51.365")))
+    creds = ClientCredentials("client_id", "xyz_client_secret", "https://login.relationalai.com/oauth/token")
+    creds.access_token = access_token
+
+    # write/read access token to cache
+    _write_token_cache(creds)
+    cached_token = _read_token_cache(creds)
+
+    # check if access token is serialized/de-serialized correctly from the cache
+    @test cached_token === access_token
 end