21121: Get keep byte usage from Prometheus
[arvados.git] / lib / boot / helpers.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package boot
6
7 import (
8         "context"
9         "net/url"
10
11         "git.arvados.org/arvados.git/lib/controller/rpc"
12         "git.arvados.org/arvados.git/sdk/go/arvados"
13         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
14         "git.arvados.org/arvados.git/sdk/go/auth"
15         "git.arvados.org/arvados.git/sdk/go/keepclient"
16         "gopkg.in/check.v1"
17 )
18
19 // ClientsWithToken returns Context, Arvados.Client and keepclient structs
20 // initialized to connect to the cluster with the supplied Arvados token.
21 func (super *Supervisor) ClientsWithToken(clusterID, token string) (context.Context, *arvados.Client, *keepclient.KeepClient) {
22         cl := super.cluster
23         if super.children != nil {
24                 cl = super.children[clusterID].cluster
25         } else if clusterID != cl.ClusterID {
26                 panic("bad clusterID " + clusterID)
27         }
28         ctx := auth.NewContext(super.ctx, auth.NewCredentials(token))
29         ac, err := arvados.NewClientFromConfig(cl)
30         if err != nil {
31                 panic(err)
32         }
33         ac.AuthToken = token
34         arv, err := arvadosclient.New(ac)
35         if err != nil {
36                 panic(err)
37         }
38         kc := keepclient.New(arv)
39         return ctx, ac, kc
40 }
41
42 // UserClients logs in as a user called "example", get the user's API token,
43 // initialize clients with the API token, set up the user and
44 // optionally activate the user.  Return client structs for
45 // communicating with the cluster on behalf of the 'example' user.
46 func (super *Supervisor) UserClients(clusterID string, rootctx context.Context, c *check.C, conn *rpc.Conn, authEmail string, activate bool) (context.Context, *arvados.Client, *keepclient.KeepClient, arvados.User) {
47         login, err := conn.UserSessionCreate(rootctx, rpc.UserSessionCreateOptions{
48                 ReturnTo: ",https://example.com",
49                 AuthInfo: rpc.UserSessionAuthInfo{
50                         Email:     authEmail,
51                         FirstName: "Example",
52                         LastName:  "User",
53                         Username:  "example",
54                 },
55         })
56         c.Assert(err, check.IsNil)
57         redirURL, err := url.Parse(login.RedirectLocation)
58         c.Assert(err, check.IsNil)
59         userToken := redirURL.Query().Get("api_token")
60         c.Logf("user token: %q", userToken)
61         ctx, ac, kc := super.ClientsWithToken(clusterID, userToken)
62         user, err := conn.UserGetCurrent(ctx, arvados.GetOptions{})
63         c.Assert(err, check.IsNil)
64         _, err = conn.UserSetup(rootctx, arvados.UserSetupOptions{UUID: user.UUID})
65         c.Assert(err, check.IsNil)
66         if activate {
67                 _, err = conn.UserActivate(rootctx, arvados.UserActivateOptions{UUID: user.UUID})
68                 c.Assert(err, check.IsNil)
69                 user, err = conn.UserGetCurrent(ctx, arvados.GetOptions{})
70                 c.Assert(err, check.IsNil)
71                 c.Logf("user UUID: %q", user.UUID)
72                 if !user.IsActive {
73                         c.Fatalf("failed to activate user -- %#v", user)
74                 }
75         }
76         return ctx, ac, kc, user
77 }
78
79 // RootClients returns Context, arvados.Client and keepclient structs initialized
80 // to communicate with the cluster as the system root user.
81 func (super *Supervisor) RootClients(clusterID string) (context.Context, *arvados.Client, *keepclient.KeepClient) {
82         return super.ClientsWithToken(clusterID, super.Cluster(clusterID).SystemRootToken)
83 }
84
85 // AnonymousClients returns Context, arvados.Client and keepclient structs initialized
86 // to communicate with the cluster as the anonymous user.
87 func (super *Supervisor) AnonymousClients(clusterID string) (context.Context, *arvados.Client, *keepclient.KeepClient) {
88         return super.ClientsWithToken(clusterID, super.Cluster(clusterID).Users.AnonymousUserToken)
89 }
90
91 // Conn gets rpc connection struct initialized to communicate with the
92 // specified cluster.
93 func (super *Supervisor) Conn(clusterID string) *rpc.Conn {
94         controllerURL := url.URL(super.Cluster(clusterID).Services.Controller.ExternalURL)
95         return rpc.NewConn(clusterID, &controllerURL, true, rpc.PassthroughTokenProvider)
96 }