New command for exporting organization audit logs (#749)

* PoC of using core API generated bindings in the cLI

* add a script to fetch and generate an openAPI client for core API

* remove generated client

* * Just issue and HTTP query to retrieve audit logs
* Use proper `net/http` methods
* Update usage of client.DoPublic across the code

* handle arguments

* accept organization name and pass the short nameé

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update test

* update mock

* remove obsolete .gitignore entry

* remove unused field

* go fmt

* undo format change

* update comment and go fmt

* use pointer and rename variable

* fix formatting

* fix tests

* replace magic number by a variable

* store regex in variable

* new http client for the cloud APIs

* break down command

* go fmt

* add TestGetOrganizationAuditLogs

* go fmt

* add TestOrganizationExportAuditLogs

* use LocalCore rather than a new config entry

* add TestContextGetPublicRESTAPIURL

* test the --output-file flag

* re-organize flags

* Update cmd/cloud/organization.go

Co-authored-by: Neel Dalsania <neel.dalsania@astronomer.io>

* add missing import

* remove redundant part of usage text

* update unit test

* Remove redundant information from usage and correctly set Local flags (Flags() instead of LocalFlags())

* stream the logs to the desired output rather than staging a copy

* update test and mocks

* remove print

* remove unused variable

* close buffer on error too

* add unit tests

* fix casing

* fix test

* add client name and version  in headers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* never stream results to stdout + indicate the output are GZIP

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Neel Dalsania <neel.dalsania@astronomer.io>

airflow_versions/http_client.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net/http"
 	"strings"
 
 	"github.com/astronomer/astro-cli/pkg/httputil"
@@ -33,7 +34,7 @@ type Request struct{}
 
 // DoWithClient (request) is a wrapper to more easily pass variables to a client.Do request
 func (r *Request) DoWithClient(api *Client) (*Response, error) {
-	doOpts := httputil.DoOptions{
+	doOpts := &httputil.DoOptions{
 		Headers: map[string]string{
 			"Accept": "application/json",
 		},
@@ -48,13 +49,14 @@ func (r *Request) Do() (*Response, error) {
 }
 
 // Do executes a query against the updates astronomer API, logging out any errors contained in the response object
-func (c *Client) Do(doOpts httputil.DoOptions) (*Response, error) {
+func (c *Client) Do(doOpts *httputil.DoOptions) (*Response, error) {
 	var response httputil.HTTPResponse
-	url := RuntimeReleaseURL
+	doOpts.Path = RuntimeReleaseURL
 	if c.useAstronomerCertified {
-		url = AirflowReleaseURL
+		doOpts.Path = AirflowReleaseURL
 	}
-	httpResponse, err := c.HTTPClient.Do("GET", url, &doOpts)
+	doOpts.Method = http.MethodGet
+	httpResponse, err := c.HTTPClient.Do(doOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -72,7 +74,7 @@ func (c *Client) Do(doOpts httputil.DoOptions) (*Response, error) {
 	decode := Response{}
 	err = json.NewDecoder(strings.NewReader(response.Body)).Decode(&decode)
 	if err != nil {
-		return nil, fmt.Errorf("failed to JSON decode %s response: %w", url, err)
+		return nil, fmt.Errorf("failed to JSON decode %s response: %w", doOpts.Path, err)
 	}
 
 	return &decode, nil

airflow_versions/http_client_test.go

@@ -27,7 +27,7 @@ func TestClientDo(t *testing.T) {
 		})
 
 		client := &Client{HTTPClient: httpClient, useAstronomerCertified: false}
-		resp, err := client.Do(httputil.DoOptions{})
+		resp, err := client.Do(&httputil.DoOptions{})
 		assert.NoError(t, err)
 		assert.Equal(t, mockResp, *resp)
 	})

astro-client/astro.go

@@ -3,8 +3,16 @@ package astro
 import (
 	"errors"
 	"fmt"
+	"io"
+	"net/http"
+	"regexp"
+	"strings"
+
+	"github.com/astronomer/astro-cli/pkg/httputil"
 )
 
+var organizationShortNameRegex = regexp.MustCompile("[^a-z0-9-]")
+
 type Client interface {
 	GetUserInfo() (*Self, error)
 	// Workspace
@@ -32,6 +40,7 @@ type Client interface {
 	GetWorkerQueueOptions() (WorkerQueueDefaultOptions, error)
 	// Organizations
 	GetOrganizations() ([]Organization, error)
+	GetOrganizationAuditLogs(orgName string, earliest int) (io.ReadCloser, error)
 }
 
 func (c *HTTPClient) GetUserInfo() (*Self, error) {
@@ -286,3 +295,19 @@ func (c *HTTPClient) GetOrganizations() ([]Organization, error) {
 	}
 	return resp.Data.GetOrganizations, nil
 }
+
+func (c *HTTPClient) GetOrganizationAuditLogs(orgName string, earliest int) (io.ReadCloser, error) {
+	// An organization short name has only lower case characters and is stripped of non-alphanumeric characters (expect for hyphens)
+	orgShortName := strings.ToLower(orgName)
+	orgShortName = organizationShortNameRegex.ReplaceAllString(orgShortName, "")
+	doOpts := &httputil.DoOptions{
+		Method:  http.MethodGet,
+		Headers: make(map[string]string),
+		Path:    fmt.Sprintf("/organizations/%s/audit-logs?earliest=%d", orgShortName, earliest),
+	}
+	streamBuffer, err := c.DoPublicRESTStreamQuery(doOpts)
+	if err != nil {
+		return nil, err
+	}
+	return streamBuffer, nil
+}

astro-client/astro_test.go

@@ -952,3 +952,59 @@ func TestGetOrganizations(t *testing.T) {
 		assert.Contains(t, err.Error(), "Internal Server Error")
 	})
 }
+
+func TestGetOrganizationAuditLogs(t *testing.T) {
+	testUtil.InitTestConfig(testUtil.CloudPlatform)
+
+	t.Run("Can export organization audit logs", func(t *testing.T) {
+		mockResponse := "A lot of audit logs entries"
+		client := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+			return &http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(bytes.NewBuffer([]byte(mockResponse))),
+				Header:     make(http.Header),
+			}
+		})
+		astroClient := NewAstroClient(client)
+
+		resp, err := astroClient.GetOrganizationAuditLogs("test-org-id", 50)
+		assert.NoError(t, err)
+		output := new(bytes.Buffer)
+		io.Copy(output, resp)
+		assert.Equal(t, mockResponse, output.String())
+	})
+
+	t.Run("Permission denied", func(t *testing.T) {
+		errorMessage := "Invalid authorization token."
+		errorResponse, err := json.Marshal(map[string]string{"message": errorMessage})
+		assert.NoError(t, err)
+		client := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+			return &http.Response{
+				StatusCode: 401,
+				Body:       io.NopCloser(bytes.NewBuffer(errorResponse)),
+				Header:     make(http.Header),
+			}
+		})
+		astroClient := NewAstroClient(client)
+
+		_, err = astroClient.GetOrganizationAuditLogs("test-org-id", 50)
+		assert.Contains(t, err.Error(), errorMessage)
+	})
+
+	t.Run("Invalid earliest parameter", func(t *testing.T) {
+		errorMessage := "Invalid query parameter values: field earliest failed check on min;"
+		errorResponse, err := json.Marshal(map[string]string{"message": errorMessage})
+		assert.NoError(t, err)
+		client := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+			return &http.Response{
+				StatusCode: 400,
+				Body:       io.NopCloser(bytes.NewBuffer(errorResponse)),
+				Header:     make(http.Header),
+			}
+		})
+		astroClient := NewAstroClient(client)
+
+		_, err = astroClient.GetOrganizationAuditLogs("test-org-id", 50)
+		assert.Contains(t, err.Error(), errorMessage)
+	})
+}

astro-client/client.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net/http"
 	"strings"
 
 	"github.com/astronomer/astro-cli/context"
@@ -45,50 +46,67 @@ func (r *Request) DoWithPublicClient(api *HTTPClient) (*Response, error) {
 	if err != nil {
 		return nil, err
 	}
-	doOpts := httputil.DoOptions{
+	doOpts := &httputil.DoOptions{
 		Data: data,
 		Headers: map[string]string{
 			"Accept": "application/json",
 		},
 	}
 
-	return api.DoPublic(doOpts)
+	return api.doPublicGraphQLQuery(doOpts)
 }
 
-// Do (request) is a wrapper to more easily pass variables to a client.Do request
-func (r *Request) DoPublic() (*Response, error) {
-	return r.DoWithPublicClient(NewAstroClient(httputil.NewHTTPClient()))
-}
-
-// Do executes a query against the Astrohub API, logging out any errors contained in the response object
-func (c *HTTPClient) DoPublic(doOpts httputil.DoOptions) (*Response, error) {
+// Set the path and Authorization header for a REST request to core API
+func (c *HTTPClient) prepareRESTRequest(doOpts *httputil.DoOptions) error {
 	cl, err := context.GetCurrentContext()
 	if err != nil {
-		return nil, err
+		return err
 	}
-
-	// set headers
 	if cl.Token != "" {
 		doOpts.Headers["authorization"] = cl.Token
 	}
-	doOpts.Headers["apollographql-client-name"] = "cli" // nolint: goconst
-	doOpts.Headers["apollographql-client-version"] = version.CurrVersion
+	doOpts.Path = cl.GetPublicRESTAPIURL() + doOpts.Path
+	doOpts.Headers["x-astro-client-identifier"] = "cli" // nolint: goconst
+	doOpts.Headers["x-astro-client-version"] = version.CurrVersion
+	return nil
+}
 
-	var response httputil.HTTPResponse
-	httpResponse, err := c.HTTPClient.Do("POST", cl.GetPublicAPIURL(), &doOpts)
+// DoPublicRESTQuery executes a query against core API
+func (c *HTTPClient) DoPublicRESTQuery(doOpts *httputil.DoOptions) (*httputil.HTTPResponse, error) {
+	err := c.prepareRESTRequest(doOpts)
 	if err != nil {
 		return nil, err
 	}
-	defer httpResponse.Body.Close()
+	return c.DoPublic(doOpts)
+}
 
-	body, err := io.ReadAll(httpResponse.Body)
+// DoPublicRESTQuery executes a query against core API and returns a raw buffer to stream
+func (c *HTTPClient) DoPublicRESTStreamQuery(doOpts *httputil.DoOptions) (io.ReadCloser, error) {
+	err := c.prepareRESTRequest(doOpts)
+	if err != nil {
+		return nil, err
+	}
+	return c.DoPublicStream(doOpts)
+}
+
+// DoPublicGraphQLQuery executes a query against Astrohub GraphQL API, logging out any errors contained in the response object
+func (c *HTTPClient) doPublicGraphQLQuery(doOpts *httputil.DoOptions) (*Response, error) {
+	cl, err := context.GetCurrentContext()
 	if err != nil {
 		return nil, err
 	}
 
-	response = httputil.HTTPResponse{
-		Raw:  httpResponse,
-		Body: string(body),
+	if cl.Token != "" {
+		doOpts.Headers["authorization"] = cl.Token
+	}
+	doOpts.Headers["apollographql-client-name"] = "cli" // nolint: goconst
+	doOpts.Headers["apollographql-client-version"] = version.CurrVersion
+	doOpts.Method = http.MethodPost
+	doOpts.Path = cl.GetPublicGraphQLAPIURL()
+
+	response, err := c.DoPublic(doOpts)
+	if err != nil {
+		return nil, fmt.Errorf("Error processing GraphQL request: %w", err)
 	}
 	decode := Response{}
 	err = json.NewDecoder(strings.NewReader(response.Body)).Decode(&decode)
@@ -106,3 +124,31 @@ func (c *HTTPClient) DoPublic(doOpts httputil.DoOptions) (*Response, error) {
 
 	return &decode, nil
 }
+
+func (c *HTTPClient) DoPublic(doOpts *httputil.DoOptions) (*httputil.HTTPResponse, error) {
+	httpResponse, err := c.HTTPClient.Do(doOpts)
+	if err != nil {
+		return nil, err
+	}
+	defer httpResponse.Body.Close()
+
+	body, err := io.ReadAll(httpResponse.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	response := &httputil.HTTPResponse{
+		Raw:  httpResponse,
+		Body: string(body),
+	}
+
+	return response, nil
+}
+
+func (c *HTTPClient) DoPublicStream(doOpts *httputil.DoOptions) (io.ReadCloser, error) {
+	httpResponse, err := c.HTTPClient.Do(doOpts)
+	if err != nil {
+		return nil, err
+	}
+	return httpResponse.Body, nil
+}

astro-client/client_test.go

@@ -1,13 +1,53 @@
 package astro
 
 import (
+	"bytes"
+	"io"
+	"net/http"
 	"testing"
 
 	"github.com/astronomer/astro-cli/pkg/httputil"
+	testUtil "github.com/astronomer/astro-cli/pkg/testing"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestNewAstroClient(t *testing.T) {
 	client := NewAstroClient(httputil.NewHTTPClient())
 	assert.NotNil(t, client, "Can't create new Astro client")
 }
+
+func TestPrepareRESTRequest(t *testing.T) {
+	client := NewAstroClient(httputil.NewHTTPClient())
+	doOpts := &httputil.DoOptions{
+		Path: "/test",
+		Headers: map[string]string{
+			"test": "test",
+		},
+	}
+	err := client.prepareRESTRequest(doOpts)
+	assert.NoError(t, err)
+	assert.Equal(t, "test", doOpts.Headers["test"])
+	// Test context has no token
+	assert.Equal(t, "", doOpts.Headers["Authorization"])
+}
+
+func TestDoPublicRESTQuery(t *testing.T) {
+	mockResponse := "A REST query response"
+	client := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+		return &http.Response{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBuffer([]byte(mockResponse))),
+			Header:     make(http.Header),
+		}
+	})
+	astroClient := NewAstroClient(client)
+	doOpts := &httputil.DoOptions{
+		Path: "/test",
+		Headers: map[string]string{
+			"test": "test",
+		},
+	}
+	resp, err := astroClient.DoPublicRESTQuery(doOpts)
+	assert.NoError(t, err)
+	assert.Equal(t, mockResponse, resp.Body)
+}