5b1aac88-2fcd-4067-b077-48a1fbc5d97d

package openapi

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "net/http"
        "os"
        "path/filepath"
        "strings"
        "time"

        "github.com/astronomer/astro-cli/config"
        "gopkg.in/yaml.v3"
)

const (
        // SpecURL is the URL to fetch the Astro Cloud API OpenAPI specification.
        SpecURL = "https://api.astronomer.io/spec/v1.0"
        // CloudCacheFileName is the name of the cache file for the Astro Cloud API.
        CloudCacheFileName = "openapi-cache.json"
        // AirflowCacheFileNameTemplate is the template for version-specific Airflow cache files.
        AirflowCacheFileNameTemplate = "openapi-airflow-%s-cache.json"
        // CacheTTL is how long the cache is valid.
        CacheTTL = 24 * time.Hour
        // FetchTimeout is the timeout for fetching the spec.
        FetchTimeout = 30 * time.Second
        // dirPermissions is the permission mode for created directories.
        dirPermissions = 0o755
        // filePermissions is the permission mode for the cache file.
        filePermissions = 0o600
)

// AirflowCacheFileNameForVersion returns the cache file name for a specific Airflow version.
// The version is normalized to strip build metadata (e.g., "3.1.7+astro.1" -> "3.1.7").
func AirflowCacheFileNameForVersion(version string) string {
        normalizedVersion := NormalizeAirflowVersion(version)
        return fmt.Sprintf(AirflowCacheFileNameTemplate, normalizedVersion)
}

// CachedSpec wraps the OpenAPI spec with metadata for caching.
type CachedSpec struct {
        Spec      *OpenAPISpec `json:"spec"`
        FetchedAt time.Time    `json:"fetchedAt"`
}

// Cache manages fetching and caching of an OpenAPI specification.
type Cache struct {
        specURL     string
        cachePath   string
        stripPrefix string // optional prefix to strip from endpoint paths (e.g. "/api/v2")
        httpClient  *http.Client
        spec        *OpenAPISpec
        fetchedAt   time.Time
}

// getHTTPClient returns the configured HTTP client, defaulting to http.DefaultClient.
func (c *Cache) getHTTPClient() *http.Client {
        if c.httpClient != nil {
                return c.httpClient
        }
        return http.DefaultClient
}

// SetHTTPClient configures the HTTP client used to fetch specs.
func (c *Cache) SetHTTPClient(client *http.Client) {
        c.httpClient = client
}

// NewCache creates a new OpenAPI cache with default settings.
func NewCache() *Cache {
        return &Cache{
                specURL:   SpecURL,
                cachePath: filepath.Join(config.HomeConfigPath, CloudCacheFileName),
        }
}

// NewCacheWithOptions creates a new OpenAPI cache with custom settings.
func NewCacheWithOptions(specURL, cachePath string) *Cache {
        return &Cache{
                specURL:   specURL,
                cachePath: cachePath,
        }
}

// NewAirflowCacheForVersion creates a new OpenAPI cache configured for a specific Airflow version.
// It automatically determines the correct spec URL and cache file path based on the version.
// Endpoint paths are stripped of their API version prefix (/api/v1 or /api/v2) so that
// callers work with version-agnostic paths like /dags instead of /api/v2/dags.
func NewAirflowCacheForVersion(version string) (*Cache, error) {
        specURL, err := BuildAirflowSpecURL(version)
        if err != nil {
                return nil, fmt.Errorf("building spec URL for version %s: %w", version, err)
        }

        cachePath := filepath.Join(config.HomeConfigPath, AirflowCacheFileNameForVersion(version))

        // Determine the API prefix to strip based on major version
        prefix := "/api/v2"
        normalized := NormalizeAirflowVersion(version)
        if strings.HasPrefix(normalized, "2.") {
                prefix = "/api/v1"
        }

        return &Cache{
                specURL:     specURL,
                cachePath:   cachePath,
                stripPrefix: prefix,
        }, nil
}

// Load loads the OpenAPI spec, using cache if valid or fetching if needed.
// If forceRefresh is true, the cache is ignored and a fresh spec is fetched.
func (c *Cache) Load(forceRefresh bool) error {
        if !forceRefresh {
                // Try to load from cache first
                if err := c.readCache(); err == nil && !c.isExpired() {
                        return nil
                }
        }

        // Fetch fresh spec
        if err := c.fetchSpec(); err != nil {
                // If fetch fails and we have a stale cache, use it
                if c.spec != nil {
                        return nil
                }
                return err
        }

        // Save to cache
        return c.saveCache()
}

// GetSpec returns the loaded OpenAPI spec.
func (c *Cache) GetSpec() *OpenAPISpec {
        return c.spec
}

// GetEndpoints extracts all endpoints from the loaded spec.
// If a stripPrefix is configured, it is removed from each endpoint path.
func (c *Cache) GetEndpoints() []Endpoint {
        if c.spec == nil {
                return nil
        }
        endpoints := ExtractEndpoints(c.spec)
        if c.stripPrefix != "" {
                for i := range endpoints {
                        endpoints[i].Path = strings.TrimPrefix(endpoints[i].Path, c.stripPrefix)
                }
        }
        return endpoints
}

// IsLoaded returns true if a spec has been loaded.
func (c *Cache) IsLoaded() bool {
        return c.spec != nil
}

// readCache attempts to read the cached spec from disk.
func (c *Cache) readCache() error {
        data, err := os.ReadFile(c.cachePath)
        if err != nil {
                return err
        }

        var cached CachedSpec
        if err := json.Unmarshal(data, &cached); err != nil {
                return err
        }

        c.spec = cached.Spec
        c.fetchedAt = cached.FetchedAt
        return nil
}

// saveCache saves the current spec to disk.
func (c *Cache) saveCache() error {
        cached := CachedSpec{
                Spec:      c.spec,
                FetchedAt: c.fetchedAt,
        }

        data, err := json.Marshal(cached)
        if err != nil {
                return err
        }

        // Ensure directory exists
        dir := filepath.Dir(c.cachePath)
        if err := os.MkdirAll(dir, dirPermissions); err != nil {
                return err
        }

        return os.WriteFile(c.cachePath, data, filePermissions)
}

// isExpired returns true if the cached spec has expired.
func (c *Cache) isExpired() bool {
        return time.Since(c.fetchedAt) > CacheTTL
}

// fetchSpec fetches the OpenAPI spec from the remote URL.
func (c *Cache) fetchSpec() error {
        ctx, cancel := context.WithTimeout(context.Background(), FetchTimeout)
        defer cancel()

        req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.specURL, http.NoBody)
        if err != nil {
                return fmt.Errorf("creating request: %w", err)
        }

        // Accept both JSON and YAML
        req.Header.Set("Accept", "application/json, application/yaml, text/yaml, */*")

        resp, err := c.getHTTPClient().Do(req)
        if err != nil {
                return fmt.Errorf("fetching OpenAPI spec: %w", err)
        }
        defer resp.Body.Close()

        if resp.StatusCode != http.StatusOK {
                return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
        }

        body, err := io.ReadAll(resp.Body)
        if err != nil {
                return fmt.Errorf("reading response body: %w", err)
        }

        var spec OpenAPISpec
        contentType := resp.Header.Get("Content-Type")

        // Try to parse based on content type, fallback to trying both formats
        switch {
        case strings.Contains(contentType, "yaml") || strings.Contains(contentType, "yml"):
                if err := yaml.Unmarshal(body, &spec); err != nil {
                        return fmt.Errorf("parsing OpenAPI spec as YAML: %w", err)
                }
        case strings.Contains(contentType, "json"):
                if err := json.Unmarshal(body, &spec); err != nil {
                        return fmt.Errorf("parsing OpenAPI spec as JSON: %w", err)
                }
        default:
                // Content-Type not helpful, try YAML first (more common for OpenAPI), then JSON
                if err := yaml.Unmarshal(body, &spec); err != nil {
                        if err := json.Unmarshal(body, &spec); err != nil {
                                return fmt.Errorf("parsing OpenAPI spec (tried YAML and JSON): %w", err)
                        }
                }
        }

        c.spec = &spec
        c.fetchedAt = time.Now()
        return nil
}

// ClearCache removes the cached spec file.
func (c *Cache) ClearCache() error {
        return os.Remove(c.cachePath)
}

1	package openapi
2
3	import (
4	"context"
5	"encoding/json"
6	"fmt"
7	"io"
8	"net/http"
9	"os"
10	"path/filepath"
11	"strings"
12	"time"
13
14	"github.com/astronomer/astro-cli/config"
15	"gopkg.in/yaml.v3"
16	)
17
18	const (
19	// SpecURL is the URL to fetch the Astro Cloud API OpenAPI specification.
20	SpecURL = "https://api.astronomer.io/spec/v1.0"
21	// CloudCacheFileName is the name of the cache file for the Astro Cloud API.
22	CloudCacheFileName = "openapi-cache.json"
23	// AirflowCacheFileNameTemplate is the template for version-specific Airflow cache files.
24	AirflowCacheFileNameTemplate = "openapi-airflow-%s-cache.json"
25	// CacheTTL is how long the cache is valid.
26	CacheTTL = 24 * time.Hour
27	// FetchTimeout is the timeout for fetching the spec.
28	FetchTimeout = 30 * time.Second
29	// dirPermissions is the permission mode for created directories.
30	dirPermissions = 0o755
31	// filePermissions is the permission mode for the cache file.
32	filePermissions = 0o600
33	)
34
35	// AirflowCacheFileNameForVersion returns the cache file name for a specific Airflow version.
36	// The version is normalized to strip build metadata (e.g., "3.1.7+astro.1" -> "3.1.7").
37	func AirflowCacheFileNameForVersion(version string) string {	4✔
38	normalizedVersion := NormalizeAirflowVersion(version)	4✔
39	return fmt.Sprintf(AirflowCacheFileNameTemplate, normalizedVersion)	4✔
40	}	4✔
41
42	// CachedSpec wraps the OpenAPI spec with metadata for caching.
43	type CachedSpec struct {
44	Spec *OpenAPISpec `json:"spec"`
45	FetchedAt time.Time `json:"fetchedAt"`
46	}
47
48	// Cache manages fetching and caching of an OpenAPI specification.
49	type Cache struct {
50	specURL string
51	cachePath string
52	stripPrefix string // optional prefix to strip from endpoint paths (e.g. "/api/v2")
53	httpClient *http.Client
54	spec *OpenAPISpec
55	fetchedAt time.Time
56	}
57
58	// getHTTPClient returns the configured HTTP client, defaulting to http.DefaultClient.
NEW 59	func (c Cache) getHTTPClient() http.Client {	×
NEW 60	if c.httpClient != nil {	×
NEW 61	return c.httpClient	×
NEW 62	}	×
NEW 63	return http.DefaultClient	×
64	}
65
66	// SetHTTPClient configures the HTTP client used to fetch specs.
NEW 67	func (c Cache) SetHTTPClient(client http.Client) {	×
NEW 68	c.httpClient = client	×
NEW 69	}	×
70
71	// NewCache creates a new OpenAPI cache with default settings.
NEW 72	func NewCache() *Cache {	×
NEW 73	return &Cache{	×
NEW 74	specURL: SpecURL,	×
NEW 75	cachePath: filepath.Join(config.HomeConfigPath, CloudCacheFileName),	×
NEW 76	}	×
NEW 77	}	×
78
79	// NewCacheWithOptions creates a new OpenAPI cache with custom settings.
NEW 80	func NewCacheWithOptions(specURL, cachePath string) *Cache {	×
NEW 81	return &Cache{	×
NEW 82	specURL: specURL,	×
NEW 83	cachePath: cachePath,	×
NEW 84	}	×
NEW 85	}	×
86
87	// NewAirflowCacheForVersion creates a new OpenAPI cache configured for a specific Airflow version.
88	// It automatically determines the correct spec URL and cache file path based on the version.
89	// Endpoint paths are stripped of their API version prefix (/api/v1 or /api/v2) so that
90	// callers work with version-agnostic paths like /dags instead of /api/v2/dags.
NEW 91	func NewAirflowCacheForVersion(version string) (*Cache, error) {	×
NEW 92	specURL, err := BuildAirflowSpecURL(version)	×
NEW 93	if err != nil {	×
NEW 94	return nil, fmt.Errorf("building spec URL for version %s: %w", version, err)	×
NEW 95	}	×
96
NEW 97	cachePath := filepath.Join(config.HomeConfigPath, AirflowCacheFileNameForVersion(version))	×
NEW 98		×
NEW 99	// Determine the API prefix to strip based on major version	×
NEW 100	prefix := "/api/v2"	×
NEW 101	normalized := NormalizeAirflowVersion(version)	×
NEW 102	if strings.HasPrefix(normalized, "2.") {	×
NEW 103	prefix = "/api/v1"	×
NEW 104	}	×
105
NEW 106	return &Cache{	×
NEW 107	specURL: specURL,	×
NEW 108	cachePath: cachePath,	×
NEW 109	stripPrefix: prefix,	×
NEW 110	}, nil	×
111	}
112
113	// Load loads the OpenAPI spec, using cache if valid or fetching if needed.
114	// If forceRefresh is true, the cache is ignored and a fresh spec is fetched.
NEW 115	func (c *Cache) Load(forceRefresh bool) error {	×
NEW 116	if !forceRefresh {	×
NEW 117	// Try to load from cache first	×
NEW 118	if err := c.readCache(); err == nil && !c.isExpired() {	×
NEW 119	return nil	×
NEW 120	}	×
121	}
122
123	// Fetch fresh spec
NEW 124	if err := c.fetchSpec(); err != nil {	×
NEW 125	// If fetch fails and we have a stale cache, use it	×
NEW 126	if c.spec != nil {	×
NEW 127	return nil	×
NEW 128	}	×
NEW 129	return err	×
130	}
131
132	// Save to cache
NEW 133	return c.saveCache()	×
134	}
135
136	// GetSpec returns the loaded OpenAPI spec.
NEW 137	func (c Cache) GetSpec() OpenAPISpec {	×
NEW 138	return c.spec	×
NEW 139	}	×
140
141	// GetEndpoints extracts all endpoints from the loaded spec.
142	// If a stripPrefix is configured, it is removed from each endpoint path.
NEW 143	func (c *Cache) GetEndpoints() []Endpoint {	×
NEW 144	if c.spec == nil {	×
NEW 145	return nil	×
NEW 146	}	×
NEW 147	endpoints := ExtractEndpoints(c.spec)	×
NEW 148	if c.stripPrefix != "" {	×
NEW 149	for i := range endpoints {	×
NEW 150	endpoints[i].Path = strings.TrimPrefix(endpoints[i].Path, c.stripPrefix)	×
NEW 151	}	×
152	}
NEW 153	return endpoints	×
154	}
155
156	// IsLoaded returns true if a spec has been loaded.
NEW 157	func (c *Cache) IsLoaded() bool {	×
NEW 158	return c.spec != nil	×
NEW 159	}	×
160
161	// readCache attempts to read the cached spec from disk.
NEW 162	func (c *Cache) readCache() error {	×
NEW 163	data, err := os.ReadFile(c.cachePath)	×
NEW 164	if err != nil {	×
NEW 165	return err	×
NEW 166	}	×
167
NEW 168	var cached CachedSpec	×
NEW 169	if err := json.Unmarshal(data, &cached); err != nil {	×
NEW 170	return err	×
NEW 171	}	×
172
NEW 173	c.spec = cached.Spec	×
NEW 174	c.fetchedAt = cached.FetchedAt	×
NEW 175	return nil	×
176	}
177
178	// saveCache saves the current spec to disk.
NEW 179	func (c *Cache) saveCache() error {	×
NEW 180	cached := CachedSpec{	×
NEW 181	Spec: c.spec,	×
NEW 182	FetchedAt: c.fetchedAt,	×
NEW 183	}	×
NEW 184		×
NEW 185	data, err := json.Marshal(cached)	×
NEW 186	if err != nil {	×
NEW 187	return err	×
NEW 188	}	×
189
190	// Ensure directory exists
NEW 191	dir := filepath.Dir(c.cachePath)	×
NEW 192	if err := os.MkdirAll(dir, dirPermissions); err != nil {	×
NEW 193	return err	×
NEW 194	}	×
195
NEW 196	return os.WriteFile(c.cachePath, data, filePermissions)	×
197	}
198
199	// isExpired returns true if the cached spec has expired.
NEW 200	func (c *Cache) isExpired() bool {	×
NEW 201	return time.Since(c.fetchedAt) > CacheTTL	×
NEW 202	}	×
203
204	// fetchSpec fetches the OpenAPI spec from the remote URL.
NEW 205	func (c *Cache) fetchSpec() error {	×
NEW 206	ctx, cancel := context.WithTimeout(context.Background(), FetchTimeout)	×
NEW 207	defer cancel()	×
NEW 208		×
NEW 209	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.specURL, http.NoBody)	×
NEW 210	if err != nil {	×
NEW 211	return fmt.Errorf("creating request: %w", err)	×
NEW 212	}	×
213
214	// Accept both JSON and YAML
NEW 215	req.Header.Set("Accept", "application/json, application/yaml, text/yaml, /")	×
NEW 216		×
NEW 217	resp, err := c.getHTTPClient().Do(req)	×
NEW 218	if err != nil {	×
NEW 219	return fmt.Errorf("fetching OpenAPI spec: %w", err)	×
NEW 220	}	×
NEW 221	defer resp.Body.Close()	×
NEW 222		×
NEW 223	if resp.StatusCode != http.StatusOK {	×
NEW 224	return fmt.Errorf("unexpected status code: %d", resp.StatusCode)	×
NEW 225	}	×
226
NEW 227	body, err := io.ReadAll(resp.Body)	×
NEW 228	if err != nil {	×
NEW 229	return fmt.Errorf("reading response body: %w", err)	×
NEW 230	}	×
231
NEW 232	var spec OpenAPISpec	×
NEW 233	contentType := resp.Header.Get("Content-Type")	×
NEW 234		×
NEW 235	// Try to parse based on content type, fallback to trying both formats	×
NEW 236	switch {	×
NEW 237	case strings.Contains(contentType, "yaml") \|\| strings.Contains(contentType, "yml"):	×
NEW 238	if err := yaml.Unmarshal(body, &spec); err != nil {	×
NEW 239	return fmt.Errorf("parsing OpenAPI spec as YAML: %w", err)	×
NEW 240	}	×
NEW 241	case strings.Contains(contentType, "json"):	×
NEW 242	if err := json.Unmarshal(body, &spec); err != nil {	×
NEW 243	return fmt.Errorf("parsing OpenAPI spec as JSON: %w", err)	×
NEW 244	}	×
NEW 245	default:	×
NEW 246	// Content-Type not helpful, try YAML first (more common for OpenAPI), then JSON	×
NEW 247	if err := yaml.Unmarshal(body, &spec); err != nil {	×
NEW 248	if err := json.Unmarshal(body, &spec); err != nil {	×
NEW 249	return fmt.Errorf("parsing OpenAPI spec (tried YAML and JSON): %w", err)	×
NEW 250	}	×
251	}
252	}
253
NEW 254	c.spec = &spec	×
NEW 255	c.fetchedAt = time.Now()	×
NEW 256	return nil	×
257	}
258
259	// ClearCache removes the cached spec file.
NEW 260	func (c *Cache) ClearCache() error {	×
NEW 261	return os.Remove(c.cachePath)	×
NEW 262	}	×

astronomer / astro-cli / 5b1aac88-2fcd-4067-b077-48a1fbc5d97d

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous