Merge pull request #15 from atlanhq/Unmarshal-search-assets

Custom Unmarshalling for search assets
atlanhq · Apr 26, 2024 · 62008b3 · 62008b3
2 parents 44a15ab + 4912bee
commit 62008b3
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 21 deletions.
diff --git a/atlan/client/atlan_tag_cache_test.go b/atlan/client/atlan_tag_cache_test.go
@@ -59,10 +59,12 @@ func TestIntegrationAtlanTagCache_GetNameForID(t *testing.T) {
 	cache := NewAtlanTagCache(client)
 
 	// Ensure the cache is populated
-	_ = cache.RefreshCache()
+	resp, _ := GetAll()
+	tagName := resp.AtlanTagDefs[0].DisplayName
+	id, err := cache.GetIDForName(tagName)
 
 	// Assuming "BBDjIBZUNHtKPExR1Z3a5I" is a valid GUID
-	name, err := cache.GetNameForID("BBDjIBZUNHtKPExR1Z3a5I")
+	name, err := cache.GetNameForID(id)
 
 	// Verify
 	assert.NoError(t, err)

diff --git a/atlan/model/search.go b/atlan/model/search.go
@@ -3,6 +3,7 @@
 package model
 
 import (
+	"encoding/json"
 	"github.com/atlanhq/atlan-go/atlan"
 	"github.com/atlanhq/atlan-go/atlan/model/assets"
 )
@@ -455,29 +456,12 @@ type SearchParameters struct {
 	Query                 string                 `json:"query"`
 }
 
-// Entity represents an entity in the Atlas search response.
-type Entity struct {
-	TypeName            string                 `json:"typeName"`
-	Attributes          map[string]interface{} `json:"attributes"`
-	Guid                string                 `json:"guid"`
-	Status              string                 `json:"status"`
-	DisplayText         string                 `json:"displayText"`
-	ClassificationNames []string               `json:"classificationNames"`
-	Tags                []interface{}          `json:"classifications"`
-	MeaningNames        []interface{}          `json:"meaningNames"`
-	Meanings            []interface{}          `json:"meanings"`
-	IsIncomplete        bool                   `json:"isIncomplete"`
-	Labels              []interface{}          `json:"labels"`
-	CreatedBy           string                 `json:"createdBy"`
-	UpdatedBy           string                 `json:"updatedBy"`
-	CreateTime          int64                  `json:"createTime"`
-	UpdateTime          int64                  `json:"updateTime"`
-}
-
 type SearchAssets struct {
 	assets.Asset
 	assets.Table
 	assets.Column
+	QualifiedName    *string           `json:"qualifiedName,omitempty"`
+	Name             *string           `json:"name,omitempty"`
 	SearchAttributes *SearchAttributes `json:"Attributes,omitempty"`
 	NotNull          *bool             `json:"notNull,omitempty"`
 }
@@ -486,3 +470,36 @@ type SearchAttributes struct {
 	QualifiedName *string `json:"qualifiedName,omitempty"`
 	Name          *string `json:"name,omitempty"`
 }
+
+func (sa *SearchAssets) UnmarshalJSON(data []byte) error {
+	// Define an auxiliary struct to decode the JSON
+	type AuxSearchAssets struct {
+		assets.Asset
+		assets.Table
+		assets.Column
+		QualifiedName    *string           `json:"qualifiedName,omitempty"`
+		Name             *string           `json:"name,omitempty"`
+		SearchAttributes *SearchAttributes `json:"attributes,omitempty"`
+		NotNull          *bool             `json:"notNull,omitempty"`
+	}
+
+	// Decode into the auxiliary struct
+	var aux AuxSearchAssets
+	if err := json.Unmarshal(data, &aux); err != nil {
+		return err
+	}
+
+	// Copy fields from auxiliary struct to the main struct
+	sa.Asset = aux.Asset
+	sa.Table = aux.Table
+	sa.Column = aux.Column
+	sa.NotNull = aux.NotNull
+
+	// Check if any search attributes are present
+	if aux.SearchAttributes != nil {
+		sa.Name = aux.SearchAttributes.Name
+		sa.QualifiedName = aux.SearchAttributes.QualifiedName
+	}
+
+	return nil
+}