Merge pull request #44 from JuliaParallel/tan/ugi

allow both realUser and effectiveUser optionally

src/api_hdfs_base.jl

@@ -62,10 +62,13 @@ struct HDFSFile
     path::AbstractString
 end
 
-function HDFSFile(uristr::AbstractString)
+function HDFSFile(uristr::AbstractString; ugi::Union{UserGroupInformation,Nothing}=nothing, proxy::Bool=false)
     uri = URI(uristr)
     (uri.scheme == "hdfs") || throw(HDFSException("Not a HDFS URI: $uristr"))
-    client = HDFSClient(uri.host, uri.port, UserGroupInformation(uri.userinfo))
+    if ugi === nothing
+        ugi = isempty(uri.userinfo) ? UserGroupInformation() : UserGroupInformation(uri.userinfo)
+    end
+    client = HDFSClient(uri.host, uri.port, ugi)
     HDFSFile(client, uri.path)
 end
 

src/cluster_manager.jl

@@ -119,6 +119,7 @@ function launch(manager::YarnManager, params::Dict, instances_arr::Array, c::Con
     cookie = ":cookie_" * cluster_cookie()
     initargs = "using Elly; Elly.setup_worker($(ipaddr.host), $(port), $(cookie))"
     clc = launchcontext(cmd="$cmd -e '$initargs'", env=appenv)
+    # TODO: do we need to add tokens into clc?
 
     @debug("YarnManager launch", initargs, context=clc)
     on_alloc = (cid) -> container_start(manager.am, cid, clc)

src/ugi.jl

@@ -1,14 +1,22 @@
 # Basic UserGroupInformation to store user information and tokens,
+# realUser: what it is at client side
+# effectiveUser: set this to specify what this should be proxied as at the server end
+# ref: https://www.opencore.com/blog/2016/5/user-name-handling-in-hadoop/
 mutable struct UserGroupInformation
     userinfo::UserInformationProto
     tokens::Dict{AbstractString,TokenProto}
-    function UserGroupInformation(username::AbstractString="")
-        isempty(username) && (username = ENV["USER"])
-        userinfo = protobuild(UserInformationProto, Dict(:realUser => username))
+    function UserGroupInformation(username::AbstractString=default_username(); proxy::Bool=false, proxyuser::AbstractString=username)
+        userinfo = proxy ? protobuild(UserInformationProto, Dict(:realUser => username, :effectiveUser => proxyuser)) : protobuild(UserInformationProto, Dict(:realUser => username))
         new(userinfo, Dict{AbstractString,TokenProto}())
     end
 end
 
+function default_username()
+    ("HADOOP_USER_NAME" in keys(ENV)) && (return ENV["HADOOP_USER_NAME"])
+    ("USER" in keys(ENV)) && (return ENV["USER"])
+    error("Can not determine user information. Either HADOOP_USER_NAME or USER must be set.")
+end
+
 add_token(ugi::UserGroupInformation, token::TokenProto) = add_token(ugi, token.service, token)
 add_token(ugi::UserGroupInformation, alias::AbstractString, token::TokenProto) = (ugi.tokens[alias] = token; nothing)
 

test/hdfstests.jl

@@ -2,10 +2,27 @@ using Elly
 using Test
 using Random
 
+function test_ugi()
+    ugi = UserGroupInformation()
+    @test !isempty(ugi.userinfo.realUser)
+
+    ugi = UserGroupInformation(; proxy=true)
+    @test !isempty(ugi.userinfo.realUser)
+    @test !isempty(ugi.userinfo.effectiveUser)
+
+    proxyuser = ugi.userinfo.realUser * "proxy"
+    ugi = UserGroupInformation(; proxy=true, proxyuser=proxyuser)
+    @test !isempty(ugi.userinfo.realUser)
+    @test !isempty(ugi.userinfo.effectiveUser)
+    @test ugi.userinfo.effectiveUser == proxyuser
+    nothing
+end
+
 function test_hdfs(host="localhost", port=9000)
     limitedtestenv = (get(ENV, "CI", "false") == "true")
 
-    hdfsclnt = HDFSClient(host, port)
+    ugi = UserGroupInformation()
+    hdfsclnt = HDFSClient(host, port, ugi)
 
     exists(hdfsclnt, "/tmp") || mkdir(hdfsclnt, "/tmp")
 
@@ -67,7 +84,7 @@ function test_hdfs(host="localhost", port=9000)
     println(st)
     @test st.name == "bar"
     mv(bar_file, "/tmp/foo/bar2")
-    bar2_file = HDFSFile("hdfs://$(host):$(port)/tmp/foo/bar2")
+    bar2_file = HDFSFile("hdfs://$(host):$(port)/tmp/foo/bar2"; ugi=ugi)
     @test isfile(bar2_file)
     rm(bar2_file)
 

test/runtests.jl

@@ -2,5 +2,6 @@
 include("hdfstests.jl")
 include("yarntests.jl")
 
+test_ugi()
 test_hdfs()
 test_yarn()