Check for invalid api key while fetching typedefs

Signed-off-by: Karanjot Singh <[email protected]>
atlanhq · May 24, 2024 · 373f7fc · 373f7fc
1 parent 0218829
commit 373f7fc
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 32 deletions.
diff --git a/atlan/client/client.go b/atlan/client/client.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"net/url"
 	"os"
 	"strings"
 )
@@ -165,14 +166,30 @@ func (ac *AtlanClient) DisableLogging() {
 }
 
 // CallAPI makes a generic API call.
-func (ac *AtlanClient) CallAPI(api *API, queryParams map[string]string, requestObj interface{}) ([]byte, error) {
+func (ac *AtlanClient) CallAPI(api *API, queryParams interface{}, requestObj interface{}) ([]byte, error) {
 	params := deepCopy(ac.requestParams)
 	path := ac.host + api.Endpoint.Atlas + api.Path
 
-	if queryParams != nil {
+	query := url.Values{}
+	switch v := queryParams.(type) {
+	case map[string]string:
+		for key, value := range v {
+			query.Add(key, value)
+		}
+	case map[string][]string:
+		for key, values := range v {
+			for _, value := range values {
+				query.Add(key, value)
+			}
+		}
+	default:
 		params["params"] = queryParams
 	}
 
+	if len(query) > 0 {
+		path += "?" + query.Encode()
+	}
+
 	if requestObj != nil {
 		//fmt.Println("Request Object:", requestObj)
 		requestJSON, err := json.Marshal(requestObj)
@@ -186,6 +203,7 @@ func (ac *AtlanClient) CallAPI(api *API, queryParams map[string]string, requestO
 
 	ac.logAPICall(api.Method, path)
 
+	//logger.Log.Debugf("Params: %v", params)
 	response, err := ac.makeRequest(api.Method, path, params)
 	if err != nil {
 		if response != nil && response.Body != nil {

diff --git a/atlan/client/typedef_client.go b/atlan/client/typedef_client.go
@@ -57,11 +57,11 @@ func buildTypeDefRequest(typedef model.TypeDef) (model.TypeDefResponse, error) {
 func NewTypeDefResponse(rawJSON []byte) (*model.TypeDefResponse, error) {
 	var response model.TypeDefResponse
 	if err := json.Unmarshal(rawJSON, &response); err != nil {
-		//fmt.Println(response)
 		return nil, err
 	}
 	return &response, nil
 }
+
 func RefreshCaches(typedef model.TypeDef) error {
 
 	switch t := typedef.(type) {
@@ -89,13 +89,55 @@ func GetAll() (*model.TypeDefResponse, error) {
 }
 
 // Get retrieves a TypeDefResponse object that contains a list of the specified category type definitions in Atlan.
-func Get(typeCategory atlan.AtlanTypeCategory) (*model.TypeDefResponse, error) {
-	queryParams := map[string]string{"type": typeCategory.String()}
+func Get(typeCategory interface{}) (*model.TypeDefResponse, error) {
+	var categories []string
+	hasStruct := false
+
+	switch v := typeCategory.(type) {
+	case atlan.AtlanTypeCategory:
+		if v == atlan.AtlanTypeCategoryStruct {
+			hasStruct = true
+		}
+		categories = append(categories, v.String())
+	case []atlan.AtlanTypeCategory:
+		for _, tc := range v {
+			if tc == atlan.AtlanTypeCategoryStruct {
+				hasStruct = true
+			}
+			categories = append(categories, tc.String())
+		}
+	default:
+		return nil, fmt.Errorf("invalid type category")
+	}
+
+	if !hasStruct {
+		categories = append(categories, atlan.AtlanTypeCategoryStruct.String())
+	}
+
+	queryParams := map[string][]string{
+		"type": categories,
+	}
+
 	rawJSON, err := DefaultAtlanClient.CallAPI(&GET_ALL_TYPE_DEFS, queryParams, nil)
+	if err != nil {
+		return nil, AtlanError{
+			ErrorCode: errorCodes[CONNECTION_ERROR],
+			Args:      []interface{}{"IOException"},
+		}
+	}
+
+	response, err := NewTypeDefResponse(rawJSON)
 	if err != nil {
 		return nil, err
 	}
-	return NewTypeDefResponse(rawJSON)
+
+	if atlan.Contains(categories, atlan.AtlanTypeCategoryStruct.String()) {
+		if response == nil || response.StructDefs == nil || len(response.StructDefs) == 0 {
+			return nil, AtlanError{ErrorCode: errorCodes[EXPIRED_API_TOKEN]}
+		}
+	}
+
+	return response, nil
 }
 
 func (c *TypeDefClient) Create(typedef model.TypeDef) (*model.TypeDefResponse, error) {

diff --git a/atlan/utils.go b/atlan/utils.go
@@ -27,3 +27,13 @@ func NextID() string {
 	sNextID++
 	return fmt.Sprintf("-%d", sNextID)
 }
+
+// Helper function to check if a slice Contains a string
+func Contains(slice []string, item string) bool {
+	for _, v := range slice {
+		if v == item {
+			return true
+		}
+	}
+	return false
+}
diff --git a/main.go b/main.go
@@ -1,4 +1,3 @@
-// main.go
 package main
 
 import (
@@ -11,42 +10,50 @@ func main() {
 
 	ctx := client.NewContext()
 
+	//ctx, _ := client.Context("API_KEY", "BASE_URL")
 	ctx.SetLogger(true, "debug")
 
-	// Fetch columns of table from Atlan using Qualified Name
-	assetQualifiedName := "default/snowflake/1715371897/RAW/WIDEWORLDIMPORTERS_SALESFORCE/FIVETRAN_API_CALL"
+	response, _ := client.Get(atlan.AtlanTypeCategoryClassification)
+	fmt.Println(response.AtlanTagDefs[0].GUID)
 
-	columnSearchResponse, _ := client.NewFluentSearch().
-		PageSizes(1000).
-		ActiveAssets().
-		Where(ctx.Column.TYPENAME.Eq("Column")).
-		Where(ctx.Column.TABLE_QUALIFIED_NAME.Eq(assetQualifiedName)).
-		IncludeOnResults("userDescription", "dataType", "isPrimary").
-		SetUtmTags(atlan.PROJECT_SDK_CLI).
-		Execute()
+	/*
+		client.GetAll()
+
+		// Fetch columns of table from Atlan using Qualified Name
+		assetQualifiedName := "default/snowflake/1715371897/RAW/WIDEWORLDIMPORTERS_SALESFORCE/FIVETRAN_API_CALL"
+
+		columnSearchResponse, _ := client.NewFluentSearch().
+			PageSizes(1000).
+			ActiveAssets().
+			Where(ctx.Column.TYPENAME.Eq("Column")).
+			Where(ctx.Column.TABLE_QUALIFIED_NAME.Eq(assetQualifiedName)).
+			IncludeOnResults("userDescription", "dataType", "isPrimary").
+			SetUtmTags(atlan.PROJECT_SDK_CLI).
+			Execute()
 
-	fmt.Println(*columnSearchResponse[0].Entities[0].Name)
-	//fmt.Println(*columnSearchResponse[0].Entities[0].DataType)
+		fmt.Println(*columnSearchResponse[0].Entities[0].Name)
+		//fmt.Println(*columnSearchResponse[0].Entities[0].DataType)
 
-	qualifiedname := "default/snowflake/1715371897/RAW/WIDEWORLDIMPORTERS_SALESFORCE/FIVETRAN_API_CALL"
+		qualifiedname := "default/snowflake/1715371897/RAW/WIDEWORLDIMPORTERS_SALESFORCE/FIVETRAN_API_CALL"
 
-	response, atlanErr := client.NewFluentSearch().
-		PageSizes(10).
-		ActiveAssets().
-		Where(ctx.Table.QUALIFIED_NAME.Eq(qualifiedname)).
-		IncludeOnResults("userDescription", "ownerUsers", "ownerGroups", "certificateStatus", "tags").
-		Execute()
+		response, atlanErr := client.NewFluentSearch().
+			PageSizes(10).
+			ActiveAssets().
+			Where(ctx.Table.QUALIFIED_NAME.Eq(qualifiedname)).
+			IncludeOnResults("userDescription", "ownerUsers", "ownerGroups", "certificateStatus", "tags").
+			Execute()
 
-	if atlanErr != nil {
-		fmt.Println(atlanErr)
-	}
+		if atlanErr != nil {
+			fmt.Println(atlanErr)
+		}
 
-	fmt.Println(response[0].Entities[0].SearchMeanings[0].Guid)
+		fmt.Println(response[0].Entities[0].SearchMeanings[0].Guid)
 
-	/*
+		/*
 
 
-	 */
+	*/
+
 	/*
 		// Fetch columns of table from Atlan using Qualified Name
 		qualifiedname := "default/snowflake/1714501359/RAW/WIDEWORLDIMPORTERS_SALESFORCE/WAITLIST_WORK_TYPE_HISTORY/ID"