Skip to content

Commit c9f72d6

Browse files
adnan-alhomssiadnan-alhomssi
authoredOct 17, 2024··
Merge pull request #44 from RelationalAI/snowflake-tss
Add support for Snowflake Stages and SNOWFLAKE_FULL encryption
·
v0.12.3v0.9.1
2 parents c2b2442 + 2a2f430 commit c9f72d6

7 files changed

+1286
-10
lines changed
 

‎Project.toml

+6-2
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
name = "RustyObjectStore"
22
uuid = "1b5eed3d-1f46-4baa-87f3-a4a892b23610"
3-
version = "0.8.2"
3+
version = "0.9.1"
44

55
[deps]
6+
Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
7+
CloudBase = "85eb1798-d7c4-4918-bb13-c944d38e27ed"
68
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
9+
HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
710
JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
11+
Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
812
object_store_ffi_jll = "0e112785-0821-598c-8835-9f07837e8d7b"
913

1014
[compat]
@@ -16,7 +20,7 @@ ReTestItems = "1"
1620
Sockets = "1"
1721
Test = "1"
1822
julia = "1.8"
19-
object_store_ffi_jll = "0.8.2"
23+
object_store_ffi_jll = "0.9.1"
2024

2125
[extras]
2226
CloudBase = "85eb1798-d7c4-4918-bb13-c944d38e27ed"

‎src/RustyObjectStore.jl

+256-3
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
module RustyObjectStore
22

33
export init_object_store, get_object!, put_object, delete_object
4-
export StaticConfig, ClientOptions, Config, AzureConfig, AWSConfig
4+
export StaticConfig, ClientOptions, Config, AzureConfig, AWSConfig, SnowflakeConfig
55
export status_code, is_connection, is_timeout, is_early_eof, is_unknown, is_parse_url
66
export get_object_stream, ReadStream, finish!
77
export put_object_stream, WriteStream, cancel!, shutdown!
8-
export current_metrics
8+
export current_metrics, invalidate_config
99
export max_entries_per_chunk, ListEntry, list_objects, list_objects_stream, next_chunk!
1010

1111
using Base.Libc.Libdl: dlext
@@ -550,6 +550,156 @@ function Base.show(io::IO, conf::AWSConfig)
550550
print(io, ", ", "opts=", repr(conf.opts), ")")
551551
end
552552

553+
"""
554+
$TYPEDEF
555+
556+
Configuration for the Snowflake stage object store backend.
557+
558+
It is recommended to reuse an instance for many operations.
559+
560+
# Keyword Arguments
561+
- `stage::String`: Snowflake stage
562+
- `encryption_scheme::Option{String}`: (Optional) Encryption scheme to enforce (one of AES_128_CBC, AES_256_GCM)
563+
- `account::Option{String}`: (Optional) Snowflake account (read from SNOWFLAKE_ACCOUNT env var if missing)
564+
- `database::Option{String}`: (Optional) Snwoflake database (read from SNOWFLAKE_DATABASE env var if missing)
565+
- `schema::Option{String}`: (Optional) Snowflake schema (read from SNOWFLAKE_SCHEMA env var if missing)
566+
- `endpoint::Option{String}`: (Optional) Snowflake endpoint (read from SNOWFLAKE_ENDPOINT or SNOWFLAKE_HOST env vars if missing)
567+
- `warehouse::Option{String}`: (Optional) Snowflake warehouse
568+
- `username::Option{String}`: (Optional) Snowflake username (required for user/pass flow)
569+
- `password::Option{String}`: (Optional) Snowflake password (required for user/pass flow)
570+
- `role::Option{String}`: (Optional) Snowflake role (required for user/pass flow)
571+
- `master_token_path::Option{String}`: (Optional) Path to Snowflake master token (read from MASTER_TOKEN_PATH or defaults to `/snowflake/session/token` if missing)
572+
- `keyring_capacity::Option{Int}`: (Optional) Maximum number of keys to be kept in the in-memory keyring (key cache)
573+
- `keyring_ttl_secs::Option{Int}`: (Optional) Duration in seconds after which a key is removed from the keyring
574+
- `opts::ClientOptions`: (Optional) Client configuration options.
575+
"""
576+
struct SnowflakeConfig <: AbstractConfig
577+
stage::String
578+
encryption_scheme::Option{String}
579+
account::Option{String}
580+
database::Option{String}
581+
schema::Option{String}
582+
endpoint::Option{String}
583+
warehouse::Option{String}
584+
username::Option{String}
585+
password::Option{String}
586+
role::Option{String}
587+
master_token_path::Option{String}
588+
keyring_capacity::Option{Int}
589+
keyring_ttl_secs::Option{Int}
590+
opts::ClientOptions
591+
cached_config::Config
592+
function SnowflakeConfig(;
593+
stage::String,
594+
encryption_scheme::Option{String} = nothing,
595+
account::Option{String} = nothing,
596+
database::Option{String} = nothing,
597+
schema::Option{String} = nothing,
598+
endpoint::Option{String} = nothing,
599+
warehouse::Option{String} = nothing,
600+
username::Option{String} = nothing,
601+
password::Option{String} = nothing,
602+
role::Option{String} = nothing,
603+
master_token_path::Option{String} = nothing,
604+
keyring_capacity::Option{Int} = nothing,
605+
keyring_ttl_secs::Option{Int} = nothing,
606+
opts::ClientOptions = ClientOptions()
607+
)
608+
params = copy(opts.params)
609+
610+
params["snowflake_stage"] = stage
611+
612+
if !isnothing(encryption_scheme)
613+
params["snowflake_encryption_scheme"] = encryption_scheme
614+
end
615+
616+
if !isnothing(account)
617+
params["snowflake_account"] = account
618+
end
619+
620+
if !isnothing(database)
621+
params["snowflake_database"] = database
622+
end
623+
624+
if !isnothing(schema)
625+
params["snowflake_schema"] = schema
626+
end
627+
628+
if !isnothing(endpoint)
629+
params["snowflake_endpoint"] = endpoint
630+
end
631+
632+
if !isnothing(warehouse)
633+
params["snowflake_warehouse"] = warehouse
634+
end
635+
636+
if !isnothing(username)
637+
params["snowflake_username"] = username
638+
end
639+
640+
if !isnothing(password)
641+
params["snowflake_password"] = password
642+
end
643+
644+
if !isnothing(role)
645+
params["snowflake_role"] = role
646+
end
647+
648+
if !isnothing(master_token_path)
649+
params["snowflake_master_token_path"] = master_token_path
650+
end
651+
652+
if !isnothing(keyring_capacity)
653+
params["snowflake_keyring_capacity"] = string(keyring_capacity)
654+
end
655+
656+
if !isnothing(keyring_ttl_secs)
657+
params["snowflake_keyring_ttl_secs"] = string(keyring_ttl_secs)
658+
end
659+
660+
# All defaults for the optional values are defined on the Rust side.
661+
map!(v -> strip(v), values(params))
662+
cached_config = Config("snowflake://$(strip(stage))/", params)
663+
return new(
664+
stage,
665+
encryption_scheme,
666+
account,
667+
database,
668+
schema,
669+
endpoint,
670+
warehouse,
671+
username,
672+
password,
673+
role,
674+
master_token_path,
675+
keyring_capacity,
676+
keyring_ttl_secs,
677+
opts,
678+
cached_config
679+
)
680+
end
681+
end
682+
683+
into_config(conf::SnowflakeConfig) = conf.cached_config
684+
685+
function Base.show(io::IO, conf::SnowflakeConfig)
686+
print(io, "SnowflakeConfig("),
687+
print(io, "stage=", repr(conf.stage))
688+
@option_print(conf, encryption_scheme)
689+
@option_print(conf, account)
690+
@option_print(conf, database)
691+
@option_print(conf, schema)
692+
@option_print(conf, endpoint)
693+
@option_print(conf, warehouse)
694+
@option_print(conf, username)
695+
@option_print(conf, password, true)
696+
@option_print(conf, role)
697+
@option_print(conf, master_token_path)
698+
@option_print(conf, keyring_capacity)
699+
@option_print(conf, keyring_ttl_secs)
700+
print(io, ", ", "opts=", repr(conf.opts), ")")
701+
end
702+
553703
mutable struct Response
554704
result::Cint
555705
length::Culonglong
@@ -1755,6 +1905,105 @@ function finish!(stream::ListStream)
17551905
return true
17561906
end
17571907

1908+
mutable struct StageInfoResponseFFI
1909+
result::Cint
1910+
stage_info::Ptr{Cchar}
1911+
error_message::Ptr{Cchar}
1912+
context::Ptr{Cvoid}
1913+
1914+
StageInfoResponseFFI() = new(-1, C_NULL, C_NULL, C_NULL)
1915+
end
1916+
1917+
function current_stage_info(conf::AbstractConfig)
1918+
response = StageInfoResponseFFI()
1919+
ct = current_task()
1920+
event = Base.Event()
1921+
handle = pointer_from_objref(event)
1922+
config = into_config(conf)
1923+
while true
1924+
preserve_task(ct)
1925+
result = GC.@preserve config response event try
1926+
result = @ccall rust_lib.current_stage_info(
1927+
config::Ref{Config},
1928+
response::Ref{StageInfoResponseFFI},
1929+
handle::Ptr{Cvoid}
1930+
)::Cint
1931+
1932+
wait_or_cancel(event, response)
1933+
1934+
result
1935+
finally
1936+
unpreserve_task(ct)
1937+
end
1938+
1939+
if result == 2
1940+
# backoff
1941+
sleep(0.01)
1942+
continue
1943+
end
1944+
1945+
# No need to destroy_cstring(response.stage_info) in case of errors here
1946+
@throw_on_error(response, "current_stage_info", GetException)
1947+
1948+
info_string = unsafe_string(response.stage_info)
1949+
@ccall rust_lib.destroy_cstring(response.stage_info::Ptr{Cchar})::Cint
1950+
1951+
stage_info = JSON3.read(info_string, Dict{String, String})
1952+
return stage_info
1953+
end
1954+
end
1955+
1956+
"""
1957+
invalidate_config(conf::Option{AbstractConfig}) -> Bool
1958+
1959+
Invalidates the specified config (or all if no config is provided) in the Rust
1960+
config cache. This is useful to mitigate test interference.
1961+
1962+
# Arguments
1963+
- `conf::AbstractConfig`: (Optional) The config to be invalidated.
1964+
"""
1965+
function invalidate_config(conf::Option{AbstractConfig}=nothing)
1966+
response = Response()
1967+
ct = current_task()
1968+
event = Base.Event()
1969+
handle = pointer_from_objref(event)
1970+
while true
1971+
preserve_task(ct)
1972+
result = GC.@preserve conf response event try
1973+
result = if !isnothing(conf)
1974+
config = into_config(conf)
1975+
@ccall rust_lib.invalidate_config(
1976+
config::Ref{Config},
1977+
response::Ref{Response},
1978+
handle::Ptr{Cvoid}
1979+
)::Cint
1980+
else
1981+
@ccall rust_lib.invalidate_config(
1982+
C_NULL::Ptr{Cvoid},
1983+
response::Ref{Response},
1984+
handle::Ptr{Cvoid}
1985+
)::Cint
1986+
end
1987+
1988+
wait_or_cancel(event, response)
1989+
1990+
result
1991+
finally
1992+
unpreserve_task(ct)
1993+
end
1994+
1995+
if result == 2
1996+
# backoff
1997+
sleep(0.01)
1998+
continue
1999+
end
2000+
2001+
@throw_on_error(response, "invalidate_config", PutException)
2002+
2003+
return true
2004+
end
2005+
end
2006+
17582007
struct Metrics
17592008
live_bytes::Int64
17602009
end
@@ -1763,4 +2012,8 @@ function current_metrics()
17632012
return @ccall rust_lib.current_metrics()::Metrics
17642013
end
17652014

1766-
end # module
2015+
module Test
2016+
include("mock_server.jl")
2017+
end # Test module
2018+
2019+
end # RustyObjectStore module

‎src/mock_server.jl

+317
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,317 @@
1+
using CloudBase: CloudCredentials, AWSCredentials, AbstractStore, AWS
2+
using JSON3, HTTP, Sockets, Base64
3+
using RustyObjectStore: SnowflakeConfig, ClientOptions
4+
using Base: UUID
5+
6+
export SFGatewayMock, start, with
7+
8+
struct SFConfig
9+
account::String
10+
database::String
11+
default_schema::String
12+
end
13+
14+
struct Stage
15+
database::String
16+
schema::String
17+
name::String
18+
end
19+
20+
mutable struct SFGatewayMock
21+
credentials::CloudCredentials
22+
store::AbstractStore
23+
opts::ClientOptions
24+
config::SFConfig
25+
allowed_stages::Vector{Stage}
26+
encrypted::Bool
27+
keys_lock::ReentrantLock
28+
next_key_id::Int
29+
keys::Dict{String, String}
30+
end
31+
32+
33+
function to_stage(stage::AbstractString, config::SFConfig)
34+
parts = split(stage, ".")
35+
if length(parts) == 1
36+
return Stage(uppercase(config.database), uppercase(config.default_schema), uppercase(parts[1]))
37+
elseif length(parts) == 2
38+
return Stage(uppercase(config.database), uppercase(parts[1]), uppercase(parts[2]))
39+
elseif length(parts) == 3
40+
return Stage(uppercase(parts[1]), uppercase(parts[2]), uppercase(parts[3]))
41+
else
42+
error("Invalid stage spec")
43+
end
44+
end
45+
46+
fqsn(s::Stage) = "$(s.database).$(s.schema).$(s.name)"
47+
stage_uuid(s::Stage) = UUID((hash(fqsn(s)), hash(fqsn(s))))
48+
stage_path(s::Stage) = "stages/$(stage_uuid(s))/"
49+
50+
function SFGatewayMock(
51+
credentials::CloudCredentials,
52+
store::AbstractStore,
53+
encrypted::Bool;
54+
opts=ClientOptions(),
55+
default_schema::String="testschema",
56+
allowed_stages::Vector{String}=["teststage" * string(rand(UInt64), base=16)]
57+
)
58+
config = SFConfig(
59+
"testaccount",
60+
"testdatabase",
61+
default_schema
62+
)
63+
allowed_stages_parsed = map(s -> to_stage(s, config), allowed_stages)
64+
SFGatewayMock(
65+
credentials,
66+
store,
67+
opts,
68+
config,
69+
allowed_stages_parsed,
70+
encrypted,
71+
ReentrantLock(),
72+
1,
73+
Dict{String, String}()
74+
)
75+
end
76+
77+
function authorized(request::HTTP.Request)
78+
return HTTP.header(request, "Authorization") == "Snowflake Token=\"dummy-token\""
79+
end
80+
81+
function unauthorized_response()
82+
return HTTP.Response(401, "Invalid token")
83+
end
84+
85+
function error_json_response(msg::String; code::String="1234")
86+
response_data = Dict(
87+
"data" => Dict(
88+
"queryId" => "dummy-query-id"
89+
),
90+
"code" => code,
91+
"message" => msg,
92+
"success" => false
93+
)
94+
return HTTP.Response(200, JSON3.write(response_data))
95+
end
96+
97+
function construct_stage_info(credentials::AWSCredentials, store::AWS.Bucket, path::String, encrypted::Bool)
98+
m = match(r"(https?://.*?)/", store.baseurl)
99+
@assert !isnothing(m)
100+
test_endpoint = m.captures[1]
101+
102+
Dict(
103+
"locationType" => "S3",
104+
"location" => joinpath(store.name, "xox50000-s", path),
105+
"path" => path,
106+
"region" => "us-east-1",
107+
"storageAccount" => nothing,
108+
"isClientSideEncrypted" => encrypted,
109+
"ciphers" => encrypted ? "AES_CBC" : nothing,
110+
"creds" => Dict(
111+
"AWS_KEY_ID" => credentials.access_key_id,
112+
"AWS_SECRET_KEY" => credentials.secret_access_key,
113+
"AWS_TOKEN" => credentials.session_token
114+
),
115+
"useS3RegionalUrl" => false,
116+
"endPoint" => "us-east-1.s3.amazonaws.com",
117+
"testEndpoint" => test_endpoint
118+
)
119+
end
120+
121+
function next_id_and_key(gw::SFGatewayMock)
122+
@lock gw.keys_lock begin
123+
key_id = gw.next_key_id
124+
gw.next_key_id += 1
125+
key = base64encode(rand(UInt8, 16))
126+
push!(gw.keys, string(key_id) => key)
127+
return key_id, key
128+
end
129+
end
130+
131+
function find_key_by_id(gw::SFGatewayMock, id::String)
132+
@lock gw.keys_lock begin
133+
return get(gw.keys, id, nothing)
134+
end
135+
end
136+
137+
# Returns a SnowflakeConfig and a server instance.
138+
# The config can be used to perform operations against
139+
# a simulated Snowflake stage backed by a Minio instance.
140+
function start(gw::SFGatewayMock)
141+
(port, tcp_server) = Sockets.listenany(8080)
142+
http_server = HTTP.serve!(tcp_server) do request::HTTP.Request
143+
if request.method == "POST" && startswith(request.target, "/session/heartbeat")
144+
# Heartbeat
145+
authorized(request) || return unauthorized_response()
146+
return HTTP.Response(200, "Pong")
147+
elseif request.method == "POST" && startswith(request.target, "/session/token-request")
148+
# Token Renewal
149+
authorized(request) || return unauthorized_response()
150+
object = JSON3.read(request.body)
151+
if get(object, "oldSessionToken", nothing) != "dummy-token"
152+
return error_json_response("Invalid session token")
153+
end
154+
155+
response_data = Dict(
156+
"data" => Dict(
157+
"sessionToken" => "dummy-token",
158+
"validityInSecondsST" => 3600,
159+
"masterToken" => "dummy-master-token",
160+
"validityInSecondsMT" => 3600
161+
),
162+
"success" => true
163+
)
164+
return HTTP.Response(200, JSON3.write(response_data))
165+
elseif request.method == "POST" && startswith(request.target, "/session/v1/login-request")
166+
# Login
167+
response_data = Dict(
168+
"data" => Dict(
169+
"token" => "dummy-token",
170+
"validityInSeconds" => 3600,
171+
"masterToken" => "dummy-master-token",
172+
"masterValidityInSeconds" => 3600
173+
),
174+
"success" => true
175+
)
176+
return HTTP.Response(200, JSON3.write(response_data))
177+
elseif request.method == "POST" && startswith(request.target, "/queries/v1/query-request")
178+
# Query
179+
authorized(request) || return unauthorized_response()
180+
object = JSON3.read(request.body)
181+
sql = get(object, "sqlText", nothing)
182+
if isnothing(sql) || !isa(sql, String)
183+
return error_json_response("Missing sql query text")
184+
end
185+
186+
sql = strip(sql)
187+
188+
if startswith(sql, "PUT")
189+
m = match(r"PUT\s+?file://.*?\s+?@(.+?)(\s|$)", sql)
190+
if isnothing(m)
191+
return error_json_response("Missing stage name or file path")
192+
end
193+
194+
stage = try
195+
to_stage(m.captures[1], gw.config)
196+
catch e
197+
return error_json_response("$(e)")
198+
end
199+
200+
if !(stage in gw.allowed_stages)
201+
return error_json_response("Stage not found")
202+
end
203+
204+
encryption_material = if gw.encrypted
205+
# generate new key
206+
key_id, key = next_id_and_key(gw)
207+
Dict(
208+
"queryStageMasterKey" => key,
209+
"queryId" => string(key_id),
210+
"smkId" => key_id
211+
)
212+
else
213+
nothing
214+
end
215+
216+
217+
stage_info = if isa(gw.credentials, AWSCredentials) && isa(gw.store, AWS.Bucket)
218+
construct_stage_info(gw.credentials, gw.store, stage_path(stage), gw.encrypted)
219+
else
220+
error("unimplemented")
221+
end
222+
223+
# PUT Query
224+
response_data = Dict(
225+
"data" => Dict(
226+
"queryId" => "dummy-query-id",
227+
"encryptionMaterial" => encryption_material,
228+
"stageInfo" => stage_info
229+
),
230+
"success" => true
231+
)
232+
return HTTP.Response(200, JSON3.write(response_data))
233+
elseif startswith(sql, "GET")
234+
# GET Query
235+
m = match(r"GET\s+?@(.+?)/(.+?)\s", sql)
236+
if isnothing(m)
237+
return error_json_response("Missing stage name or file path")
238+
end
239+
240+
stage = try
241+
to_stage(m.captures[1], gw.config)
242+
catch e
243+
return error_json_response("$(e)")
244+
end
245+
246+
if !(stage in gw.allowed_stages)
247+
return error_json_response("Stage not found")
248+
end
249+
250+
path = m.captures[2]
251+
252+
stage_info = if isa(gw.credentials, AWSCredentials) && isa(gw.store, AWS.Bucket)
253+
construct_stage_info(gw.credentials, gw.store, stage_path(stage), gw.encrypted)
254+
else
255+
error("unimplemented")
256+
end
257+
258+
encryption_material = if gw.encrypted
259+
# fetch key id from s3 meta and return key
260+
response = AWS.head(
261+
stage_info["testEndpoint"] * "/" * stage_info["location"] * path;
262+
service="s3", region="us-east-1", credentials=gw.credentials
263+
)
264+
pos = findfirst(x -> x[1] == "x-amz-meta-x-amz-matdesc", response.headers)
265+
matdesc = JSON3.read(response.headers[pos][2])
266+
key_id = matdesc["queryId"]
267+
key = find_key_by_id(gw, key_id)
268+
Dict(
269+
"queryStageMasterKey" => key,
270+
"queryId" => key_id,
271+
"smkId" => parse(Int, key_id)
272+
)
273+
else
274+
nothing
275+
end
276+
277+
response_data = Dict(
278+
"data" => Dict(
279+
"queryId" => "dummy-query-id",
280+
"src_locations" => [path],
281+
"encryptionMaterial" => [encryption_material],
282+
"stageInfo" => stage_info
283+
),
284+
"success" => true
285+
)
286+
return HTTP.Response(200, JSON3.write(response_data))
287+
else
288+
return error_json_response("Unsupported query")
289+
end
290+
else
291+
return HTTP.Response(404, "Not Found")
292+
end
293+
end
294+
295+
master_path, fileio = mktemp()
296+
write(fileio, "dummy-file-master-token")
297+
sfconfig = SnowflakeConfig(
298+
stage=fqsn(gw.allowed_stages[1]),
299+
account=gw.config.account,
300+
database=gw.config.database,
301+
schema=gw.allowed_stages[1].schema,
302+
endpoint="http://127.0.0.1:$(port)",
303+
master_token_path=master_path,
304+
opts=gw.opts
305+
)
306+
return sfconfig, http_server
307+
end
308+
309+
function with(f::Function, gw::SFGatewayMock)
310+
config, server = start(gw)
311+
try
312+
f(config)
313+
finally
314+
HTTP.forceclose(server)
315+
rm(config.master_token_path)
316+
end
317+
end

‎test/basic_unified_tests.jl

+49-5
Original file line numberDiff line numberDiff line change
@@ -531,7 +531,13 @@ function run_sanity_test_cases(read_config::AbstractConfig, write_config::Abstra
531531
end
532532
end
533533

534-
function run_list_test_cases(config::AbstractConfig)
534+
function within_margin(a, b, margin = 32)
535+
return all(abs.(a .- b) .<= margin)
536+
end
537+
538+
539+
function run_list_test_cases(config::AbstractConfig; strict_entry_size=true)
540+
margin = strict_entry_size ? 0 : 32
535541
@testset "basic listing" begin
536542
for i in range(10; step=10, length=5)
537543
nbytes_written = put_object(codeunits(repeat('=', i)), "list/$(i).csv", config)
@@ -540,7 +546,7 @@ function run_list_test_cases(config::AbstractConfig)
540546

541547
entries = list_objects("list/", config)
542548
@test length(entries) == 5
543-
@test map(x -> x.size, entries) == range(10; step=10, length=5)
549+
@test within_margin(map(x -> x.size, entries), range(10; step=10, length=5), margin)
544550
@test map(x -> x.location, entries) == ["list/10.csv", "list/20.csv", "list/30.csv", "list/40.csv", "list/50.csv"]
545551
end
546552

@@ -560,7 +566,7 @@ function run_list_test_cases(config::AbstractConfig)
560566

561567
entries = list_objects("other/prefix/", config)
562568
@test length(entries) == 5
563-
@test map(x -> x.size, entries) == range(110; step=10, length=5)
569+
@test within_margin(map(x -> x.size, entries), range(110; step=10, length=5), margin)
564570
@test map(x -> x.location, entries) ==
565571
["other/prefix/110.csv", "other/prefix/120.csv", "other/prefix/130.csv", "other/prefix/140.csv", "other/prefix/150.csv"]
566572

@@ -569,6 +575,10 @@ function run_list_test_cases(config::AbstractConfig)
569575

570576
entries = list_objects("other/p/", config)
571577
@test length(entries) == 0
578+
579+
entries = list_objects("other/prefix/150.csv", config)
580+
@test length(entries) == 1
581+
@test map(x -> x.location, entries) == ["other/prefix/150.csv"]
572582
end
573583

574584
@testset "list empty entries" begin
@@ -602,7 +612,7 @@ function run_list_test_cases(config::AbstractConfig)
602612

603613
append!(entries, one_entry)
604614

605-
@test sort(map(x -> x.size, entries)) == data
615+
@test within_margin(sort(map(x -> x.size, entries)), data, margin)
606616
@test sort(map(x -> x.location, entries)) == sort(map(x -> "list/$(x).csv", data))
607617
end
608618

@@ -640,7 +650,7 @@ function run_list_test_cases(config::AbstractConfig)
640650

641651
@test isnothing(next_chunk!(stream))
642652

643-
@test sort(map(x -> x.size, entries)) == data[51:end]
653+
@test within_margin(sort(map(x -> x.size, entries)), data[51:end], margin)
644654
@test sort(map(x -> x.location, entries)) == sort(map(x -> key(x), data[51:end]))
645655
end
646656
end
@@ -765,3 +775,37 @@ Minio.with(; debug=true, public=true) do conf
765775
run_read_write_test_cases(config_no_creds, config)
766776
end # Minio.with
767777
end # @testitem
778+
779+
@testitem "Basic Snowflake Stage usage" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
780+
using CloudBase.CloudTest: Minio
781+
using RustyObjectStore: SnowflakeConfig, ClientOptions
782+
783+
# For interactive testing, use Minio.run() instead of Minio.with()
784+
# conf, p = Minio.run(; debug=true, public=false); atexit(() -> kill(p))
785+
Minio.with(; debug=true, public=false) do conf
786+
credentials, container = conf
787+
with(SFGatewayMock(credentials, container, false)) do config::SnowflakeConfig
788+
run_read_write_test_cases(config)
789+
run_stream_test_cases(config)
790+
run_list_test_cases(config)
791+
run_sanity_test_cases(config)
792+
end
793+
end # Minio.with
794+
end # @testitem
795+
796+
@testitem "Basic Snowflake Stage usage (encrypted)" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
797+
using CloudBase.CloudTest: Minio
798+
using RustyObjectStore: SnowflakeConfig, ClientOptions
799+
800+
# For interactive testing, use Minio.run() instead of Minio.with()
801+
# conf, p = Minio.run(; debug=true, public=false); atexit(() -> kill(p))
802+
Minio.with(; debug=true, public=false) do conf
803+
credentials, container = conf
804+
with(SFGatewayMock(credentials, container, true)) do config::SnowflakeConfig
805+
run_read_write_test_cases(config)
806+
run_stream_test_cases(config)
807+
run_list_test_cases(config; strict_entry_size=false)
808+
run_sanity_test_cases(config)
809+
end
810+
end # Minio.with
811+
end # @testitem

‎test/common_testsetup.jl

+6
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,9 @@
1212
)
1313
init_object_store(test_config)
1414
end
15+
16+
@testsetup module SnowflakeMock
17+
using RustyObjectStore.Test: SFGatewayMock, start, with
18+
19+
export SFGatewayMock, start, with
20+
end

‎test/snowflake_api_tests.jl

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
@testitem "SnowflakeConfig" begin
2+
# basic case
3+
@test repr(SnowflakeConfig(;
4+
stage="a"
5+
)) == "SnowflakeConfig(stage=\"a\", opts=ClientOptions())"
6+
7+
# password is obscured when printing
8+
@test repr(SnowflakeConfig(;
9+
stage="a",
10+
username="b",
11+
password="c",
12+
role="d"
13+
)) == "SnowflakeConfig(stage=\"a\", username=\"b\", password=*****, role=\"d\", opts=ClientOptions())"
14+
15+
# optional params are supported
16+
@test repr(SnowflakeConfig(;
17+
stage="a",
18+
encryption_scheme="b",
19+
account="c",
20+
database="d",
21+
schema="e",
22+
endpoint="f"
23+
)) == "SnowflakeConfig(stage=\"a\", encryption_scheme=\"b\", account=\"c\", database=\"d\", schema=\"e\", endpoint=\"f\", opts=ClientOptions())"
24+
end

‎test/snowflake_stage_exception_tests.jl

+628
Large diffs are not rendered by default.

2 commit comments

Comments
 (2)

adnan-alhomssi commented on Oct 17, 2024

@adnan-alhomssi
adnan-alhomssi
on Oct 17, 2024
MemberAuthor

@JuliaRegistrator register

Upgrades to a newer version of the object_store_ffi library that supports Snowflake Stages as a backing store. Adds the new SnowflakeConfig (for interacting with those stages) including tests and some test related helpers.

JuliaRegistrator commented on Oct 17, 2024

@JuliaRegistrator
JuliaRegistrator
on Oct 17, 2024

Registration pull request created: JuliaRegistries/General/117455

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.9.1 -m "<description of version>" c9f72d6ece3a71c3214282acd61a2217b2817895
git push origin v0.9.1

Also, note the warning: Version 0.9.1 skips over 0.9.0
This can be safely ignored. However, if you want to fix this you can do so. Call register() again after making the fix. This will update the Pull request.

Please sign in to comment.