12 Commits

Author SHA1 Message Date
Gregor Schulte
a04489ddf5 Simplify flow run filtering logic by removing response struct; adjust pagination calculations to reflect updated results handling 2026-04-10 15:52:16 +02:00
Gregor Schulte
020a835a24 Simplify deployment filtering logic by removing unnecessary response struct; adjust pagination calculations to match updated logic 2026-04-10 13:49:35 +02:00
Gregor Schulte
9b70c13556 Refactor error handling with errors.As for type assertions; add error checks in test encoders and body closures 2026-04-08 16:10:09 +02:00
Gregor Schulte
2dcc0b13dd Extrahiere Deployments-, TaskRun- und WorkPool-Services in separate Dateien; entferne nicht benötigte Services und Funktionen. 2026-04-08 16:07:44 +02:00
Gregor Schulte
57531a7d95 Füge Modelle für Flows, FlowRuns, TaskRuns, WorkPools, WorkQueues, Deployments, Variablen, FlowRunStates, Logs, und Blocks samt zugehöriger Unmarshal-Logik und Zeitfeld-Unterstützung hinzu; ergänze Tests für die FlowRunStates-Service-Methoden. 2026-03-27 14:02:32 +01:00
Gregor Schulte
3aff707116 Entferne das models-Paket und füge neue Services für Artefakte, Automationen, Ereignisse, Suchfilter, Concurrency Limits, Task Worker, sowie FlowRun- und TaskRun-Status hinzu. 2026-03-26 22:29:01 +01:00
Gregor Schulte
6814c7fafb Verbessere die URL-Bau-Logik im Client; stelle sicher, dass der Basis-Pfad nicht ersetzt wird, wenn der Dienst-Pfad absolut ist 2026-02-17 14:14:26 +01:00
Gregor Schulte
727eba0e5e Füge benutzerdefinierte JSON-Unmarshaling-Logik für Zeitfelder hinzu; verbessere die Handhabung von optionalen Zeitwerten 2026-02-17 12:53:56 +01:00
Gregor Schulte
79caa168db Füge BlockTypes, BlockSchemas, BlockDocuments und BlockCapabilities Services hinzu; erweitere Modelle für Blocktypen, Blockschemas und Blockdokumente 2026-02-17 11:34:44 +01:00
Gregor Schulte
ae8a90d1e8 Entferne Integrations-Tests und zugehörige Dateien; füge CLAUDE.md zur .gitignore hinzu 2026-02-17 11:16:36 +01:00
Gregor Schulte
df4978cc04 Update import paths to reflect migration to schultesdev domain in all relevant files 2026-02-02 09:15:39 +01:00
Gregor Schulte
d98b8369fd Update module path to reflect migration to schultesdev domain in go.mod 2026-02-02 09:05:48 +01:00
63 changed files with 34264 additions and 1955 deletions

2
.gitignore vendored
View File

@@ -37,3 +37,5 @@ pkg/models/generated*.go
# Temporary files
tmp/
temp/
CLAUDE.md

6
doc.go
View File

@@ -7,7 +7,7 @@ pagination support, and comprehensive error handling.
# Installation
go get github.com/gregor/prefect-go
go get git.schultes.dev/schultesdev/prefect-go
# Quick Start
@@ -19,8 +19,8 @@ Create a client and interact with the Prefect API:
"context"
"log"
"github.com/gregor/prefect-go/pkg/client"
"github.com/gregor/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/client"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
func main() {

View File

@@ -7,8 +7,8 @@ import (
"log"
"time"
"github.com/gregor/prefect-go/pkg/client"
"github.com/gregor/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/client"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
func main() {

View File

@@ -6,8 +6,8 @@ import (
"fmt"
"log"
"github.com/gregor/prefect-go/pkg/client"
"github.com/gregor/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/client"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
func main() {

View File

@@ -7,8 +7,8 @@ import (
"log"
"time"
"github.com/gregor/prefect-go/pkg/client"
"github.com/gregor/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/client"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
func main() {

View File

@@ -6,8 +6,8 @@ import (
"fmt"
"log"
"github.com/gregor/prefect-go/pkg/client"
"github.com/gregor/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/client"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
func main() {

2
go.mod
View File

@@ -1,4 +1,4 @@
module github.com/gregor/prefect-go
module git.schultes.dev/schultesdev/prefect-go
go 1.21

29948
openapi.json Normal file

File diff suppressed because one or more lines are too long

View File

@@ -1,69 +0,0 @@
# Integration Tests
Integration tests require a running Prefect server instance.
## Running Integration Tests
### 1. Start a Prefect Server
```bash
prefect server start
```
This will start a Prefect server on `http://localhost:4200`.
### 2. Run the Integration Tests
```bash
go test -v -tags=integration ./pkg/client/
```
Or using make:
```bash
make test-integration
```
### 3. Custom Server URL
If your Prefect server is running on a different URL, set the `PREFECT_API_URL` environment variable:
```bash
export PREFECT_API_URL="http://your-server:4200/api"
go test -v -tags=integration ./pkg/client/
```
## Test Coverage
The integration tests cover:
- **Flow Lifecycle**: Create, get, update, list, and delete flows
- **Flow Run Lifecycle**: Create, get, set state, and delete flow runs
- **Deployment Lifecycle**: Create, get, pause, resume, and delete deployments
- **Pagination**: Manual pagination and iterator-based pagination
- **Variables**: Create, get, update, and delete variables
- **Admin Endpoints**: Health check and version
- **Flow Run Monitoring**: Wait for flow run completion
## Notes
- Integration tests use the `integration` build tag to separate them from unit tests
- Each test creates and cleans up its own resources
- Tests use unique UUIDs in names to avoid conflicts
- Some tests may take several seconds to complete (e.g., flow run wait tests)
## Troubleshooting
### Connection Refused
If you see "connection refused" errors, make sure:
1. Prefect server is running (`prefect server start`)
2. The server is accessible at the configured URL
3. No firewall is blocking the connection
### Test Timeout
If tests timeout:
1. Increase the test timeout: `go test -timeout 5m -tags=integration ./pkg/client/`
2. Check server logs for errors
3. Ensure the server is not under heavy load

30
pkg/client/admin.go Normal file
View File

@@ -0,0 +1,30 @@
package client
import (
"context"
"fmt"
)
// AdminService handles administrative operations.
type AdminService struct {
client *Client
}
// Health checks the health of the Prefect server.
func (a *AdminService) Health(ctx context.Context) error {
if err := a.client.get(ctx, "/health", nil); err != nil {
return fmt.Errorf("health check failed: %w", err)
}
return nil
}
// Version retrieves the server version.
func (a *AdminService) Version(ctx context.Context) (string, error) {
var result struct {
Version string `json:"version"`
}
if err := a.client.get(ctx, "/version", &result); err != nil {
return "", fmt.Errorf("failed to get version: %w", err)
}
return result.Version, nil
}

150
pkg/client/artifacts.go Normal file
View File

@@ -0,0 +1,150 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
"github.com/google/uuid"
)
// ArtifactsService handles operations related to artifacts.
type ArtifactsService struct {
client *Client
}
// Create creates a new artifact.
func (s *ArtifactsService) Create(ctx context.Context, req *models.ArtifactCreate) (*models.Artifact, error) {
var artifact models.Artifact
if err := s.client.post(ctx, "/artifacts/", req, &artifact); err != nil {
return nil, fmt.Errorf("failed to create artifact: %w", err)
}
return &artifact, nil
}
// Get retrieves an artifact by ID.
func (s *ArtifactsService) Get(ctx context.Context, id uuid.UUID) (*models.Artifact, error) {
var artifact models.Artifact
path := joinPath("/artifacts", id.String())
if err := s.client.get(ctx, path, &artifact); err != nil {
return nil, fmt.Errorf("failed to get artifact: %w", err)
}
return &artifact, nil
}
// GetLatest retrieves the latest artifact for a given key.
func (s *ArtifactsService) GetLatest(ctx context.Context, key string) (*models.Artifact, error) {
var artifact models.Artifact
path := joinPath("/artifacts", key, "latest")
if err := s.client.get(ctx, path, &artifact); err != nil {
return nil, fmt.Errorf("failed to get latest artifact: %w", err)
}
return &artifact, nil
}
// Update updates an artifact.
func (s *ArtifactsService) Update(ctx context.Context, id uuid.UUID, req *models.ArtifactUpdate) error {
path := joinPath("/artifacts", id.String())
if err := s.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update artifact: %w", err)
}
return nil
}
// Delete deletes an artifact by ID.
func (s *ArtifactsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/artifacts", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete artifact: %w", err)
}
return nil
}
// List retrieves a list of artifacts with optional filtering.
func (s *ArtifactsService) List(ctx context.Context, filter *models.ArtifactFilter, offset, limit int) (*pagination.PaginatedResponse[models.Artifact], error) {
if filter == nil {
filter = &models.ArtifactFilter{}
}
filter.Offset = offset
filter.Limit = limit
type response struct {
Results []models.Artifact `json:"results"`
Count int `json:"count"`
}
var resp response
if err := s.client.post(ctx, "/artifacts/filter", filter, &resp); err != nil {
return nil, fmt.Errorf("failed to list artifacts: %w", err)
}
return &pagination.PaginatedResponse[models.Artifact]{
Results: resp.Results,
Count: resp.Count,
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
}, nil
}
// ListAll returns an iterator for all artifacts matching the filter.
func (s *ArtifactsService) ListAll(ctx context.Context, filter *models.ArtifactFilter) *pagination.Iterator[models.Artifact] {
fetchFunc := func(ctx context.Context, offset, limit int) (*pagination.PaginatedResponse[models.Artifact], error) {
return s.List(ctx, filter, offset, limit)
}
return pagination.NewIterator(fetchFunc, 100)
}
// Count returns the number of artifacts matching the filter.
func (s *ArtifactsService) Count(ctx context.Context, filter *models.ArtifactFilter) (int, error) {
if filter == nil {
filter = &models.ArtifactFilter{}
}
var count int
if err := s.client.post(ctx, "/artifacts/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count artifacts: %w", err)
}
return count, nil
}
// ListLatest retrieves the latest artifact collections with optional filtering.
func (s *ArtifactsService) ListLatest(ctx context.Context, filter *models.ArtifactCollectionFilter, offset, limit int) (*pagination.PaginatedResponse[models.ArtifactCollection], error) {
if filter == nil {
filter = &models.ArtifactCollectionFilter{}
}
filter.Offset = offset
filter.Limit = limit
type response struct {
Results []models.ArtifactCollection `json:"results"`
Count int `json:"count"`
}
var resp response
if err := s.client.post(ctx, "/artifacts/latest/filter", filter, &resp); err != nil {
return nil, fmt.Errorf("failed to list latest artifacts: %w", err)
}
return &pagination.PaginatedResponse[models.ArtifactCollection]{
Results: resp.Results,
Count: resp.Count,
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
}, nil
}
// CountLatest returns the number of latest artifact collections matching the filter.
func (s *ArtifactsService) CountLatest(ctx context.Context, filter *models.ArtifactCollectionFilter) (int, error) {
if filter == nil {
filter = &models.ArtifactCollectionFilter{}
}
var count int
if err := s.client.post(ctx, "/artifacts/latest/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count latest artifacts: %w", err)
}
return count, nil
}

View File

@@ -0,0 +1,151 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestArtifactsService_Create(t *testing.T) {
expected := models.Artifact{
ID: uuid.New(),
Key: strPtr("test-key"),
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/artifacts/" {
t.Errorf("path = %v, want /api/artifacts/", r.URL.Path)
}
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
err := json.NewEncoder(w).Encode(expected)
if err != nil {
return
}
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
artifact, err := client.Artifacts.Create(ctx, &models.ArtifactCreate{Key: strPtr("test-key")})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if artifact.ID != expected.ID {
t.Errorf("ID = %v, want %v", artifact.ID, expected.ID)
}
}
func TestArtifactsService_Get(t *testing.T) {
id := uuid.New()
expected := models.Artifact{ID: id, Key: strPtr("test-key")}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
err := json.NewEncoder(w).Encode(expected)
if err != nil {
return
}
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
artifact, err := client.Artifacts.Get(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if artifact.ID != id {
t.Errorf("ID = %v, want %v", artifact.ID, id)
}
}
func TestArtifactsService_Delete(t *testing.T) {
id := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodDelete {
t.Errorf("method = %v, want DELETE", r.Method)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.Artifacts.Delete(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
func TestArtifactsService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
if r.URL.Path != "/api/artifacts/filter" {
t.Errorf("path = %v, want /api/artifacts/filter", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
resp := map[string]interface{}{
"results": []models.Artifact{{ID: uuid.New()}},
"count": 1,
}
err := json.NewEncoder(w).Encode(resp)
if err != nil {
return
}
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
result, err := client.Artifacts.List(ctx, nil, 0, 10)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(result.Results) != 1 {
t.Errorf("results count = %v, want 1", len(result.Results))
}
}
func TestArtifactsService_Count(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, err := w.Write([]byte(`5`))
if err != nil {
return
}
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
count, err := client.Artifacts.Count(ctx, nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if count != 5 {
t.Errorf("count = %v, want 5", count)
}
}
func strPtr(s string) *string {
return &s
}

120
pkg/client/automations.go Normal file
View File

@@ -0,0 +1,120 @@
package client
import (
"context"
"fmt"
"net/http"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// AutomationsService handles operations related to automations.
type AutomationsService struct {
client *Client
}
// Create creates a new automation.
func (s *AutomationsService) Create(ctx context.Context, req *models.AutomationCreate) (*models.Automation, error) {
var automation models.Automation
if err := s.client.do(ctx, http.MethodPost, "/automations/", req, &automation); err != nil {
return nil, fmt.Errorf("failed to create automation: %w", err)
}
return &automation, nil
}
// Get retrieves an automation by ID.
func (s *AutomationsService) Get(ctx context.Context, id uuid.UUID) (*models.Automation, error) {
var automation models.Automation
path := joinPath("/automations", id.String())
if err := s.client.get(ctx, path, &automation); err != nil {
return nil, fmt.Errorf("failed to get automation: %w", err)
}
return &automation, nil
}
// Update fully replaces an automation.
func (s *AutomationsService) Update(ctx context.Context, id uuid.UUID, req *models.AutomationUpdate) (*models.Automation, error) {
var automation models.Automation
path := joinPath("/automations", id.String())
if err := s.client.do(ctx, http.MethodPut, path, req, &automation); err != nil {
return nil, fmt.Errorf("failed to update automation: %w", err)
}
return &automation, nil
}
// Patch partially updates an automation.
func (s *AutomationsService) Patch(ctx context.Context, id uuid.UUID, req *models.AutomationPartialUpdate) error {
path := joinPath("/automations", id.String())
if err := s.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to patch automation: %w", err)
}
return nil
}
// Delete deletes an automation by ID.
func (s *AutomationsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/automations", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete automation: %w", err)
}
return nil
}
// List retrieves automations with optional filtering.
func (s *AutomationsService) List(ctx context.Context, filter *models.AutomationFilter) ([]models.Automation, error) {
if filter == nil {
filter = &models.AutomationFilter{}
}
var automations []models.Automation
if err := s.client.post(ctx, "/automations/filter", filter, &automations); err != nil {
return nil, fmt.Errorf("failed to list automations: %w", err)
}
return automations, nil
}
// Count returns the number of automations matching the filter.
func (s *AutomationsService) Count(ctx context.Context, filter *models.AutomationFilter) (int, error) {
if filter == nil {
filter = &models.AutomationFilter{}
}
var count int
if err := s.client.post(ctx, "/automations/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count automations: %w", err)
}
return count, nil
}
// GetRelatedTo retrieves automations related to a specific resource.
func (s *AutomationsService) GetRelatedTo(ctx context.Context, resourceID string) ([]models.Automation, error) {
var automations []models.Automation
path := joinPath("/automations/related-to", resourceID)
if err := s.client.get(ctx, path, &automations); err != nil {
return nil, fmt.Errorf("failed to get related automations: %w", err)
}
return automations, nil
}
// DeleteOwnedBy deletes automations owned by a specific resource.
func (s *AutomationsService) DeleteOwnedBy(ctx context.Context, resourceID string) error {
path := joinPath("/automations/owned-by", resourceID)
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete owned automations: %w", err)
}
return nil
}
// ValidateTemplate validates an automation template.
func (s *AutomationsService) ValidateTemplate(ctx context.Context, template string) error {
req := struct {
Template string `json:"template"`
}{
Template: template,
}
if err := s.client.post(ctx, "/automations/templates/validate", req, nil); err != nil {
return fmt.Errorf("failed to validate template: %w", err)
}
return nil
}

View File

@@ -0,0 +1,135 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestAutomationsService_Create(t *testing.T) {
expected := models.Automation{
ID: uuid.New(),
Name: "test-automation",
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/automations/" {
t.Errorf("path = %v, want /api/automations/", r.URL.Path)
}
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusCreated)
json.NewEncoder(w).Encode(expected)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
automation, err := client.Automations.Create(ctx, &models.AutomationCreate{
Name: "test-automation",
Trigger: json.RawMessage(`{"type":"event"}`),
Actions: []json.RawMessage{json.RawMessage(`{"type":"do-nothing"}`)},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if automation.Name != expected.Name {
t.Errorf("Name = %v, want %v", automation.Name, expected.Name)
}
}
func TestAutomationsService_Get(t *testing.T) {
id := uuid.New()
expected := models.Automation{ID: id, Name: "test"}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(expected)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
automation, err := client.Automations.Get(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if automation.ID != id {
t.Errorf("ID = %v, want %v", automation.ID, id)
}
}
func TestAutomationsService_Delete(t *testing.T) {
id := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodDelete {
t.Errorf("method = %v, want DELETE", r.Method)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.Automations.Delete(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
func TestAutomationsService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/automations/filter" {
t.Errorf("path = %v, want /api/automations/filter", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.Automation{{ID: uuid.New(), Name: "test"}})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
automations, err := client.Automations.List(ctx, nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(automations) != 1 {
t.Errorf("count = %v, want 1", len(automations))
}
}
func TestAutomationsService_Patch(t *testing.T) {
id := uuid.New()
enabled := true
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPatch {
t.Errorf("method = %v, want PATCH", r.Method)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.Automations.Patch(ctx, id, &models.AutomationPartialUpdate{Enabled: &enabled})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}

406
pkg/client/blocks.go Normal file
View File

@@ -0,0 +1,406 @@
package client
import (
"context"
"fmt"
"net/url"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// BlockTypesService handles operations related to block types.
type BlockTypesService struct {
client *Client
}
// Create creates a new block type.
func (s *BlockTypesService) Create(ctx context.Context, req *models.BlockTypeCreate) (*models.BlockType, error) {
var blockType models.BlockType
if err := s.client.post(ctx, "/block_types/", req, &blockType); err != nil {
return nil, fmt.Errorf("failed to create block type: %w", err)
}
return &blockType, nil
}
// Get retrieves a block type by ID.
func (s *BlockTypesService) Get(ctx context.Context, id uuid.UUID) (*models.BlockType, error) {
var blockType models.BlockType
path := joinPath("/block_types", id.String())
if err := s.client.get(ctx, path, &blockType); err != nil {
return nil, fmt.Errorf("failed to get block type: %w", err)
}
return &blockType, nil
}
// GetBySlug retrieves a block type by slug.
func (s *BlockTypesService) GetBySlug(ctx context.Context, slug string) (*models.BlockType, error) {
var blockType models.BlockType
path := joinPath("/block_types/slug", slug)
if err := s.client.get(ctx, path, &blockType); err != nil {
return nil, fmt.Errorf("failed to get block type by slug: %w", err)
}
return &blockType, nil
}
// List retrieves block types with optional filtering.
func (s *BlockTypesService) List(ctx context.Context, filter *models.BlockTypeFilter, offset, limit int) ([]models.BlockType, error) {
type nameFilter struct {
Like *string `json:"like,omitempty"`
}
type slugFilter struct {
Any []string `json:"any,omitempty"`
}
type blockTypeFilterCriteria struct {
Name *nameFilter `json:"name,omitempty"`
Slug *slugFilter `json:"slug,omitempty"`
}
type capabilitiesFilter struct {
All []string `json:"all,omitempty"`
}
type blockSchemaFilterCriteria struct {
Capabilities *capabilitiesFilter `json:"block_capabilities,omitempty"`
}
type request struct {
BlockTypes *blockTypeFilterCriteria `json:"block_types,omitempty"`
BlockSchemas *blockSchemaFilterCriteria `json:"block_schemas,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
req := request{Offset: offset, Limit: limit}
if filter != nil {
btf := &blockTypeFilterCriteria{}
if filter.Name != nil {
btf.Name = &nameFilter{Like: filter.Name}
}
if len(filter.Slugs) > 0 {
btf.Slug = &slugFilter{Any: filter.Slugs}
}
if btf.Name != nil || btf.Slug != nil {
req.BlockTypes = btf
}
if len(filter.Capabilities) > 0 {
req.BlockSchemas = &blockSchemaFilterCriteria{
Capabilities: &capabilitiesFilter{All: filter.Capabilities},
}
}
}
var blockTypes []models.BlockType
if err := s.client.post(ctx, "/block_types/filter", req, &blockTypes); err != nil {
return nil, fmt.Errorf("failed to list block types: %w", err)
}
return blockTypes, nil
}
// Update updates a block type.
func (s *BlockTypesService) Update(ctx context.Context, id uuid.UUID, req *models.BlockTypeUpdate) error {
path := joinPath("/block_types", id.String())
if err := s.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update block type: %w", err)
}
return nil
}
// Delete deletes a block type by ID.
func (s *BlockTypesService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/block_types", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete block type: %w", err)
}
return nil
}
// InstallSystemBlockTypes installs the system block types.
func (s *BlockTypesService) InstallSystemBlockTypes(ctx context.Context) error {
if err := s.client.post(ctx, "/block_types/install_system_block_types", nil, nil); err != nil {
return fmt.Errorf("failed to install system block types: %w", err)
}
return nil
}
// ListBlockDocumentsBySlug retrieves all block documents for a block type identified by slug.
func (s *BlockTypesService) ListBlockDocumentsBySlug(ctx context.Context, slug string, includeSecrets bool) ([]models.BlockDocument, error) {
path := joinPath("/block_types/slug", slug, "block_documents")
if includeSecrets {
query := url.Values{}
query.Set("include_secrets", "true")
path = buildPathWithValues(path, query)
}
var blockDocuments []models.BlockDocument
if err := s.client.get(ctx, path, &blockDocuments); err != nil {
return nil, fmt.Errorf("failed to list block documents for block type: %w", err)
}
return blockDocuments, nil
}
// GetBlockDocumentByName retrieves a block document by name for a given block type slug.
func (s *BlockTypesService) GetBlockDocumentByName(ctx context.Context, slug, name string, includeSecrets bool) (*models.BlockDocument, error) {
path := joinPath("/block_types/slug", slug, "block_documents/name", name)
if includeSecrets {
query := url.Values{}
query.Set("include_secrets", "true")
path = buildPathWithValues(path, query)
}
var blockDocument models.BlockDocument
if err := s.client.get(ctx, path, &blockDocument); err != nil {
return nil, fmt.Errorf("failed to get block document by name: %w", err)
}
return &blockDocument, nil
}
// BlockSchemasService handles operations related to block schemas.
type BlockSchemasService struct {
client *Client
}
// Create creates a new block schema.
func (s *BlockSchemasService) Create(ctx context.Context, req *models.BlockSchemaCreate) (*models.BlockSchema, error) {
var blockSchema models.BlockSchema
if err := s.client.post(ctx, "/block_schemas/", req, &blockSchema); err != nil {
return nil, fmt.Errorf("failed to create block schema: %w", err)
}
return &blockSchema, nil
}
// Get retrieves a block schema by ID.
func (s *BlockSchemasService) Get(ctx context.Context, id uuid.UUID) (*models.BlockSchema, error) {
var blockSchema models.BlockSchema
path := joinPath("/block_schemas", id.String())
if err := s.client.get(ctx, path, &blockSchema); err != nil {
return nil, fmt.Errorf("failed to get block schema: %w", err)
}
return &blockSchema, nil
}
// GetByChecksum retrieves a block schema by checksum and optional version.
func (s *BlockSchemasService) GetByChecksum(ctx context.Context, checksum string, version string) (*models.BlockSchema, error) {
path := joinPath("/block_schemas/checksum", checksum)
if version != "" {
query := url.Values{}
query.Set("version", version)
path = buildPathWithValues(path, query)
}
var blockSchema models.BlockSchema
if err := s.client.get(ctx, path, &blockSchema); err != nil {
return nil, fmt.Errorf("failed to get block schema by checksum: %w", err)
}
return &blockSchema, nil
}
// List retrieves block schemas with optional filtering.
func (s *BlockSchemasService) List(ctx context.Context, filter *models.BlockSchemaFilter, offset, limit int) ([]models.BlockSchema, error) {
type blockTypeIDFilter struct {
Any []string `json:"any_,omitempty"`
}
type capabilitiesFilter struct {
All []string `json:"all_,omitempty"`
}
type versionFilter struct {
Any []string `json:"any_,omitempty"`
}
type blockSchemaFilterCriteria struct {
BlockTypeID *blockTypeIDFilter `json:"block_type_id,omitempty"`
Capabilities *capabilitiesFilter `json:"block_capabilities,omitempty"`
Version *versionFilter `json:"version,omitempty"`
}
type request struct {
BlockSchemas *blockSchemaFilterCriteria `json:"block_schemas,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
req := request{Offset: offset, Limit: limit}
if filter != nil {
bsf := &blockSchemaFilterCriteria{}
if filter.BlockTypeID != nil {
bsf.BlockTypeID = &blockTypeIDFilter{Any: []string{filter.BlockTypeID.String()}}
}
if len(filter.Capabilities) > 0 {
bsf.Capabilities = &capabilitiesFilter{All: filter.Capabilities}
}
if filter.Version != nil {
bsf.Version = &versionFilter{Any: []string{*filter.Version}}
}
if bsf.BlockTypeID != nil || bsf.Capabilities != nil || bsf.Version != nil {
req.BlockSchemas = bsf
}
}
var blockSchemas []models.BlockSchema
if err := s.client.post(ctx, "/block_schemas/filter", req, &blockSchemas); err != nil {
return nil, fmt.Errorf("failed to list block schemas: %w", err)
}
return blockSchemas, nil
}
// Delete deletes a block schema by ID.
func (s *BlockSchemasService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/block_schemas", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete block schema: %w", err)
}
return nil
}
// BlockDocumentsService handles operations related to block documents.
type BlockDocumentsService struct {
client *Client
}
// Create creates a new block document.
func (s *BlockDocumentsService) Create(ctx context.Context, req *models.BlockDocumentCreate) (*models.BlockDocument, error) {
var blockDocument models.BlockDocument
if err := s.client.post(ctx, "/block_documents/", req, &blockDocument); err != nil {
return nil, fmt.Errorf("failed to create block document: %w", err)
}
return &blockDocument, nil
}
// Get retrieves a block document by ID.
func (s *BlockDocumentsService) Get(ctx context.Context, id uuid.UUID, includeSecrets bool) (*models.BlockDocument, error) {
path := joinPath("/block_documents", id.String())
if includeSecrets {
query := url.Values{}
query.Set("include_secrets", "true")
path = buildPathWithValues(path, query)
}
var blockDocument models.BlockDocument
if err := s.client.get(ctx, path, &blockDocument); err != nil {
return nil, fmt.Errorf("failed to get block document: %w", err)
}
return &blockDocument, nil
}
// List retrieves block documents with optional filtering.
func (s *BlockDocumentsService) List(ctx context.Context, filter *models.BlockDocumentFilter, includeSecrets bool, sort models.BlockDocumentSort, offset, limit int) ([]models.BlockDocument, error) {
type idFilter struct {
Any []string `json:"any_,omitempty"`
}
type nameFilter struct {
Any []string `json:"any_,omitempty"`
}
type boolFilter struct {
Eq bool `json:"eq_"`
}
type blockDocumentFilterCriteria struct {
BlockTypeID *idFilter `json:"block_type_id,omitempty"`
Name *nameFilter `json:"name,omitempty"`
IsAnonymous *boolFilter `json:"is_anonymous,omitempty"`
}
type request struct {
BlockDocuments *blockDocumentFilterCriteria `json:"block_documents,omitempty"`
IncludeSecrets bool `json:"include_secrets,omitempty"`
Sort models.BlockDocumentSort `json:"sort,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
req := request{
IncludeSecrets: includeSecrets,
Sort: sort,
Offset: offset,
Limit: limit,
}
if filter != nil {
bdf := &blockDocumentFilterCriteria{}
if filter.BlockTypeID != nil {
bdf.BlockTypeID = &idFilter{Any: []string{filter.BlockTypeID.String()}}
}
if filter.Name != nil {
bdf.Name = &nameFilter{Any: []string{*filter.Name}}
}
if filter.IsAnonymous != nil {
bdf.IsAnonymous = &boolFilter{Eq: *filter.IsAnonymous}
}
if bdf.BlockTypeID != nil || bdf.Name != nil || bdf.IsAnonymous != nil {
req.BlockDocuments = bdf
}
}
var blockDocuments []models.BlockDocument
if err := s.client.post(ctx, "/block_documents/filter", req, &blockDocuments); err != nil {
return nil, fmt.Errorf("failed to list block documents: %w", err)
}
return blockDocuments, nil
}
// Count returns the number of block documents matching the filter.
func (s *BlockDocumentsService) Count(ctx context.Context, filter *models.BlockDocumentFilter) (int, error) {
type idFilter struct {
Any []string `json:"any_,omitempty"`
}
type nameFilter struct {
Any []string `json:"any_,omitempty"`
}
type boolFilter struct {
Eq bool `json:"eq_"`
}
type blockDocumentFilterCriteria struct {
BlockTypeID *idFilter `json:"block_type_id,omitempty"`
Name *nameFilter `json:"name,omitempty"`
IsAnonymous *boolFilter `json:"is_anonymous,omitempty"`
}
type request struct {
BlockDocuments *blockDocumentFilterCriteria `json:"block_documents,omitempty"`
}
var req request
if filter != nil {
bdf := &blockDocumentFilterCriteria{}
if filter.BlockTypeID != nil {
bdf.BlockTypeID = &idFilter{Any: []string{filter.BlockTypeID.String()}}
}
if filter.Name != nil {
bdf.Name = &nameFilter{Any: []string{*filter.Name}}
}
if filter.IsAnonymous != nil {
bdf.IsAnonymous = &boolFilter{Eq: *filter.IsAnonymous}
}
if bdf.BlockTypeID != nil || bdf.Name != nil || bdf.IsAnonymous != nil {
req.BlockDocuments = bdf
}
}
var count int
if err := s.client.post(ctx, "/block_documents/count", req, &count); err != nil {
return 0, fmt.Errorf("failed to count block documents: %w", err)
}
return count, nil
}
// Update updates a block document's data.
func (s *BlockDocumentsService) Update(ctx context.Context, id uuid.UUID, req *models.BlockDocumentUpdate) error {
path := joinPath("/block_documents", id.String())
if err := s.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update block document: %w", err)
}
return nil
}
// Delete deletes a block document by ID.
func (s *BlockDocumentsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/block_documents", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete block document: %w", err)
}
return nil
}
// BlockCapabilitiesService handles operations related to block capabilities.
type BlockCapabilitiesService struct {
client *Client
}
// List retrieves all available block capabilities.
func (s *BlockCapabilitiesService) List(ctx context.Context) ([]string, error) {
var capabilities []string
if err := s.client.get(ctx, "/block_capabilities/", &capabilities); err != nil {
return nil, fmt.Errorf("failed to list block capabilities: %w", err)
}
return capabilities, nil
}

View File

@@ -12,8 +12,8 @@ import (
"strings"
"time"
"github.com/gregor/prefect-go/pkg/errors"
"github.com/gregor/prefect-go/pkg/retry"
"git.schultes.dev/schultesdev/prefect-go/pkg/errors"
"git.schultes.dev/schultesdev/prefect-go/pkg/retry"
)
const (
@@ -40,6 +40,19 @@ type Client struct {
Variables *VariablesService
Logs *LogsService
Admin *AdminService
BlockTypes *BlockTypesService
BlockSchemas *BlockSchemasService
BlockDocuments *BlockDocumentsService
BlockCapabilities *BlockCapabilitiesService
Artifacts *ArtifactsService
Automations *AutomationsService
Events *EventsService
SavedSearches *SavedSearchesService
ConcurrencyLimits *ConcurrencyLimitsService
ConcurrencyLimitsV2 *ConcurrencyLimitsV2Service
TaskWorkers *TaskWorkersService
FlowRunStates *FlowRunStatesService
TaskRunStates *TaskRunStatesService
}
// Option is a functional option for configuring the client.
@@ -79,6 +92,19 @@ func NewClient(opts ...Option) (*Client, error) {
c.Variables = &VariablesService{client: c}
c.Logs = &LogsService{client: c}
c.Admin = &AdminService{client: c}
c.BlockTypes = &BlockTypesService{client: c}
c.BlockSchemas = &BlockSchemasService{client: c}
c.BlockDocuments = &BlockDocumentsService{client: c}
c.BlockCapabilities = &BlockCapabilitiesService{client: c}
c.Artifacts = &ArtifactsService{client: c}
c.Automations = &AutomationsService{client: c}
c.Events = &EventsService{client: c}
c.SavedSearches = &SavedSearchesService{client: c}
c.ConcurrencyLimits = &ConcurrencyLimitsService{client: c}
c.ConcurrencyLimitsV2 = &ConcurrencyLimitsV2Service{client: c}
c.TaskWorkers = &TaskWorkersService{client: c}
c.FlowRunStates = &FlowRunStatesService{client: c}
c.TaskRunStates = &TaskRunStatesService{client: c}
return c, nil
}
@@ -140,11 +166,16 @@ func WithUserAgent(ua string) Option {
// do executes an HTTP request with retry logic.
func (c *Client) do(ctx context.Context, method, path string, body, result interface{}) error {
// Build full URL
u, err := c.baseURL.Parse(path)
// Build full URL by joining the base URL path with the service path.
// url.URL.Parse() is not used here because it would replace the base path
// when the service path is absolute (starts with "/").
refURL, err := url.Parse(path)
if err != nil {
return fmt.Errorf("failed to parse path: %w", err)
}
u := *c.baseURL
u.Path = strings.TrimRight(c.baseURL.Path, "/") + "/" + strings.TrimLeft(refURL.Path, "/")
u.RawQuery = refURL.RawQuery
// Serialize request body if present
var reqBody io.Reader

View File

@@ -0,0 +1,114 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// ConcurrencyLimitsService handles operations related to v1 concurrency limits.
type ConcurrencyLimitsService struct {
client *Client
}
// Create creates a new concurrency limit.
func (s *ConcurrencyLimitsService) Create(ctx context.Context, req *models.ConcurrencyLimitCreate) (*models.ConcurrencyLimit, error) {
var limit models.ConcurrencyLimit
if err := s.client.post(ctx, "/concurrency_limits/", req, &limit); err != nil {
return nil, fmt.Errorf("failed to create concurrency limit: %w", err)
}
return &limit, nil
}
// Get retrieves a concurrency limit by ID.
func (s *ConcurrencyLimitsService) Get(ctx context.Context, id uuid.UUID) (*models.ConcurrencyLimit, error) {
var limit models.ConcurrencyLimit
path := joinPath("/concurrency_limits", id.String())
if err := s.client.get(ctx, path, &limit); err != nil {
return nil, fmt.Errorf("failed to get concurrency limit: %w", err)
}
return &limit, nil
}
// GetByTag retrieves a concurrency limit by tag.
func (s *ConcurrencyLimitsService) GetByTag(ctx context.Context, tag string) (*models.ConcurrencyLimit, error) {
var limit models.ConcurrencyLimit
path := joinPath("/concurrency_limits/tag", tag)
if err := s.client.get(ctx, path, &limit); err != nil {
return nil, fmt.Errorf("failed to get concurrency limit by tag: %w", err)
}
return &limit, nil
}
// Delete deletes a concurrency limit by ID.
func (s *ConcurrencyLimitsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/concurrency_limits", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete concurrency limit: %w", err)
}
return nil
}
// DeleteByTag deletes a concurrency limit by tag.
func (s *ConcurrencyLimitsService) DeleteByTag(ctx context.Context, tag string) error {
path := joinPath("/concurrency_limits/tag", tag)
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete concurrency limit by tag: %w", err)
}
return nil
}
// ResetByTag resets a concurrency limit by tag.
func (s *ConcurrencyLimitsService) ResetByTag(ctx context.Context, tag string) error {
path := joinPath("/concurrency_limits/tag", tag, "reset")
if err := s.client.post(ctx, path, nil, nil); err != nil {
return fmt.Errorf("failed to reset concurrency limit: %w", err)
}
return nil
}
// List retrieves concurrency limits with optional filtering.
func (s *ConcurrencyLimitsService) List(ctx context.Context) ([]models.ConcurrencyLimit, error) {
var limits []models.ConcurrencyLimit
if err := s.client.post(ctx, "/concurrency_limits/filter", nil, &limits); err != nil {
return nil, fmt.Errorf("failed to list concurrency limits: %w", err)
}
return limits, nil
}
// Increment increments concurrency limit slots.
func (s *ConcurrencyLimitsService) Increment(ctx context.Context, names []string, slots int, mode string) ([]models.MinimalConcurrencyLimitResponse, error) {
req := struct {
Names []string `json:"names"`
Slots int `json:"slots"`
Mode string `json:"mode,omitempty"`
}{
Names: names,
Slots: slots,
Mode: mode,
}
var resp []models.MinimalConcurrencyLimitResponse
if err := s.client.post(ctx, "/concurrency_limits/increment", req, &resp); err != nil {
return nil, fmt.Errorf("failed to increment concurrency limits: %w", err)
}
return resp, nil
}
// Decrement decrements concurrency limit slots.
func (s *ConcurrencyLimitsService) Decrement(ctx context.Context, names []string, slots int) error {
req := struct {
Names []string `json:"names"`
Slots int `json:"slots"`
}{
Names: names,
Slots: slots,
}
if err := s.client.post(ctx, "/concurrency_limits/decrement", req, nil); err != nil {
return fmt.Errorf("failed to decrement concurrency limits: %w", err)
}
return nil
}

View File

@@ -0,0 +1,126 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestConcurrencyLimitsService_Create(t *testing.T) {
expected := models.ConcurrencyLimit{
ID: uuid.New(),
Tag: "test-tag",
ConcurrencyLimit: 10,
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(expected)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limit, err := client.ConcurrencyLimits.Create(ctx, &models.ConcurrencyLimitCreate{
Tag: "test-tag",
ConcurrencyLimit: 10,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if limit.Tag != "test-tag" {
t.Errorf("Tag = %v, want test-tag", limit.Tag)
}
}
func TestConcurrencyLimitsService_Get(t *testing.T) {
id := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(models.ConcurrencyLimit{ID: id, Tag: "test"})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limit, err := client.ConcurrencyLimits.Get(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if limit.ID != id {
t.Errorf("ID = %v, want %v", limit.ID, id)
}
}
func TestConcurrencyLimitsService_GetByTag(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(models.ConcurrencyLimit{ID: uuid.New(), Tag: "my-tag"})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limit, err := client.ConcurrencyLimits.GetByTag(ctx, "my-tag")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if limit.Tag != "my-tag" {
t.Errorf("Tag = %v, want my-tag", limit.Tag)
}
}
func TestConcurrencyLimitsService_Delete(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodDelete {
t.Errorf("method = %v, want DELETE", r.Method)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.ConcurrencyLimits.Delete(ctx, uuid.New())
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
func TestConcurrencyLimitsService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.ConcurrencyLimit{{ID: uuid.New(), Tag: "test"}})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limits, err := client.ConcurrencyLimits.List(ctx)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(limits) != 1 {
t.Errorf("count = %v, want 1", len(limits))
}
}

View File

@@ -0,0 +1,141 @@
package client
import (
"context"
"fmt"
"net/http"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// ConcurrencyLimitsV2Service handles operations related to v2 concurrency limits.
type ConcurrencyLimitsV2Service struct {
client *Client
}
// Create creates a new v2 concurrency limit.
func (s *ConcurrencyLimitsV2Service) Create(ctx context.Context, req *models.ConcurrencyLimitV2Create) (*models.ConcurrencyLimitV2, error) {
var limit models.ConcurrencyLimitV2
if err := s.client.do(ctx, http.MethodPost, "/v2/concurrency_limits/", req, &limit); err != nil {
return nil, fmt.Errorf("failed to create v2 concurrency limit: %w", err)
}
return &limit, nil
}
// Get retrieves a v2 concurrency limit by ID or name.
func (s *ConcurrencyLimitsV2Service) Get(ctx context.Context, idOrName string) (*models.ConcurrencyLimitV2, error) {
var limit models.ConcurrencyLimitV2
path := joinPath("/v2/concurrency_limits", idOrName)
if err := s.client.get(ctx, path, &limit); err != nil {
return nil, fmt.Errorf("failed to get v2 concurrency limit: %w", err)
}
return &limit, nil
}
// Update updates a v2 concurrency limit.
func (s *ConcurrencyLimitsV2Service) Update(ctx context.Context, idOrName string, req *models.ConcurrencyLimitV2Update) error {
path := joinPath("/v2/concurrency_limits", idOrName)
if err := s.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update v2 concurrency limit: %w", err)
}
return nil
}
// Delete deletes a v2 concurrency limit by ID or name.
func (s *ConcurrencyLimitsV2Service) Delete(ctx context.Context, idOrName string) error {
path := joinPath("/v2/concurrency_limits", idOrName)
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete v2 concurrency limit: %w", err)
}
return nil
}
// List retrieves all v2 concurrency limits.
func (s *ConcurrencyLimitsV2Service) List(ctx context.Context) ([]models.GlobalConcurrencyLimitResponse, error) {
var limits []models.GlobalConcurrencyLimitResponse
if err := s.client.post(ctx, "/v2/concurrency_limits/filter", nil, &limits); err != nil {
return nil, fmt.Errorf("failed to list v2 concurrency limits: %w", err)
}
return limits, nil
}
// Increment increments active slots for v2 concurrency limits.
func (s *ConcurrencyLimitsV2Service) Increment(ctx context.Context, names []string, slots int, mode string) error {
req := struct {
Names []string `json:"names"`
Slots int `json:"slots"`
Mode string `json:"mode,omitempty"`
}{
Names: names,
Slots: slots,
Mode: mode,
}
if err := s.client.post(ctx, "/v2/concurrency_limits/increment", req, nil); err != nil {
return fmt.Errorf("failed to increment v2 concurrency limits: %w", err)
}
return nil
}
// IncrementWithLease increments active slots and returns a lease.
func (s *ConcurrencyLimitsV2Service) IncrementWithLease(ctx context.Context, names []string, slots int, mode string) (*models.ConcurrencyLimitWithLeaseResponse, error) {
req := struct {
Names []string `json:"names"`
Slots int `json:"slots"`
Mode string `json:"mode,omitempty"`
}{
Names: names,
Slots: slots,
Mode: mode,
}
var resp models.ConcurrencyLimitWithLeaseResponse
if err := s.client.post(ctx, "/v2/concurrency_limits/increment-with-lease", req, &resp); err != nil {
return nil, fmt.Errorf("failed to increment v2 concurrency limits with lease: %w", err)
}
return &resp, nil
}
// Decrement decrements active slots for v2 concurrency limits.
func (s *ConcurrencyLimitsV2Service) Decrement(ctx context.Context, names []string, slots int) error {
req := struct {
Names []string `json:"names"`
Slots int `json:"slots"`
}{
Names: names,
Slots: slots,
}
if err := s.client.post(ctx, "/v2/concurrency_limits/decrement", req, nil); err != nil {
return fmt.Errorf("failed to decrement v2 concurrency limits: %w", err)
}
return nil
}
// DecrementWithLease decrements active slots using a lease ID.
func (s *ConcurrencyLimitsV2Service) DecrementWithLease(ctx context.Context, names []string, slots int, leaseID uuid.UUID) error {
req := struct {
Names []string `json:"names"`
Slots int `json:"slots"`
LeaseID uuid.UUID `json:"lease_id"`
}{
Names: names,
Slots: slots,
LeaseID: leaseID,
}
if err := s.client.post(ctx, "/v2/concurrency_limits/decrement-with-lease", req, nil); err != nil {
return fmt.Errorf("failed to decrement v2 concurrency limits with lease: %w", err)
}
return nil
}
// RenewLease renews a concurrency lease.
func (s *ConcurrencyLimitsV2Service) RenewLease(ctx context.Context, leaseID uuid.UUID) error {
path := joinPath("/v2/concurrency_limits/leases", leaseID.String(), "renew")
if err := s.client.post(ctx, path, nil, nil); err != nil {
return fmt.Errorf("failed to renew concurrency lease: %w", err)
}
return nil
}

View File

@@ -0,0 +1,128 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestConcurrencyLimitsV2Service_Create(t *testing.T) {
expected := models.ConcurrencyLimitV2{
ID: uuid.New(),
Name: "test-limit",
Limit: 5,
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v2/concurrency_limits/" {
t.Errorf("path = %v, want /api/v2/concurrency_limits/", r.URL.Path)
}
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusCreated)
json.NewEncoder(w).Encode(expected)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limit, err := client.ConcurrencyLimitsV2.Create(ctx, &models.ConcurrencyLimitV2Create{
Name: "test-limit",
Limit: 5,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if limit.Name != "test-limit" {
t.Errorf("Name = %v, want test-limit", limit.Name)
}
}
func TestConcurrencyLimitsV2Service_Get(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(models.ConcurrencyLimitV2{ID: uuid.New(), Name: "test"})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limit, err := client.ConcurrencyLimitsV2.Get(ctx, "test")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if limit.Name != "test" {
t.Errorf("Name = %v, want test", limit.Name)
}
}
func TestConcurrencyLimitsV2Service_Delete(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodDelete {
t.Errorf("method = %v, want DELETE", r.Method)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.ConcurrencyLimitsV2.Delete(ctx, "test")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
func TestConcurrencyLimitsV2Service_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.GlobalConcurrencyLimitResponse{
{ID: uuid.New(), Name: "test", Limit: 5, ActiveSlots: 0},
})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
limits, err := client.ConcurrencyLimitsV2.List(ctx)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(limits) != 1 {
t.Errorf("count = %v, want 1", len(limits))
}
}
func TestConcurrencyLimitsV2Service_RenewLease(t *testing.T) {
leaseID := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.ConcurrencyLimitsV2.RenewLease(ctx, leaseID)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}

179
pkg/client/deployments.go Normal file
View File

@@ -0,0 +1,179 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
"github.com/google/uuid"
)
// DeploymentsService handles operations related to deployments.
type DeploymentsService struct {
client *Client
}
// Create creates a new deployment.
func (s *DeploymentsService) Create(ctx context.Context, req *models.DeploymentCreate) (*models.Deployment, error) {
var deployment models.Deployment
if err := s.client.post(ctx, "/deployments/", req, &deployment); err != nil {
return nil, fmt.Errorf("failed to create deployment: %w", err)
}
return &deployment, nil
}
// Get retrieves a deployment by ID.
func (s *DeploymentsService) Get(ctx context.Context, id uuid.UUID) (*models.Deployment, error) {
var deployment models.Deployment
path := joinPath("/deployments", id.String())
if err := s.client.get(ctx, path, &deployment); err != nil {
return nil, fmt.Errorf("failed to get deployment: %w", err)
}
return &deployment, nil
}
// GetByName retrieves a deployment by flow name and deployment name.
func (s *DeploymentsService) GetByName(ctx context.Context, flowName, deploymentName string) (*models.Deployment, error) {
var deployment models.Deployment
path := fmt.Sprintf("/deployments/name/%s/%s", flowName, deploymentName)
if err := s.client.get(ctx, path, &deployment); err != nil {
return nil, fmt.Errorf("failed to get deployment by name: %w", err)
}
return &deployment, nil
}
// Update updates a deployment.
func (s *DeploymentsService) Update(ctx context.Context, id uuid.UUID, req *models.DeploymentUpdate) (*models.Deployment, error) {
var deployment models.Deployment
path := joinPath("/deployments", id.String())
if err := s.client.patch(ctx, path, req, &deployment); err != nil {
return nil, fmt.Errorf("failed to update deployment: %w", err)
}
return &deployment, nil
}
// Delete deletes a deployment by ID.
func (s *DeploymentsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/deployments", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete deployment: %w", err)
}
return nil
}
// Pause pauses a deployment's schedule.
func (s *DeploymentsService) Pause(ctx context.Context, id uuid.UUID) error {
path := joinPath("/deployments", id.String(), "pause")
if err := s.client.post(ctx, path, nil, nil); err != nil {
return fmt.Errorf("failed to pause deployment: %w", err)
}
return nil
}
// Resume resumes a deployment's schedule.
func (s *DeploymentsService) Resume(ctx context.Context, id uuid.UUID) error {
path := joinPath("/deployments", id.String(), "resume")
if err := s.client.post(ctx, path, nil, nil); err != nil {
return fmt.Errorf("failed to resume deployment: %w", err)
}
return nil
}
// CreateFlowRun creates a flow run from a deployment.
func (s *DeploymentsService) CreateFlowRun(ctx context.Context, id uuid.UUID, params map[string]interface{}) (*models.FlowRun, error) {
var flowRun models.FlowRun
path := joinPath("/deployments", id.String(), "create_flow_run")
req := struct {
Parameters map[string]interface{} `json:"parameters,omitempty"`
}{
Parameters: params,
}
if err := s.client.post(ctx, path, req, &flowRun); err != nil {
return nil, fmt.Errorf("failed to create flow run from deployment: %w", err)
}
return &flowRun, nil
}
// List retrieves a list of deployments with optional filtering.
func (s *DeploymentsService) List(ctx context.Context, filter *models.DeploymentFilter, offset, limit int) (*pagination.PaginatedResponse[models.Deployment], error) {
if filter == nil {
filter = &models.DeploymentFilter{}
}
filter.Offset = offset
filter.Limit = limit
var results []models.Deployment
if err := s.client.post(ctx, "/deployments/filter", filter, &results); err != nil {
return nil, fmt.Errorf("failed to list deployments: %w", err)
}
return &pagination.PaginatedResponse[models.Deployment]{
Results: results,
Count: offset + len(results),
Limit: limit,
Offset: offset,
HasMore: len(results) == limit,
}, nil
}
// ListAll returns an iterator for all deployments matching the filter.
func (s *DeploymentsService) ListAll(ctx context.Context, filter *models.DeploymentFilter) *pagination.Iterator[models.Deployment] {
fetchFunc := func(ctx context.Context, offset, limit int) (*pagination.PaginatedResponse[models.Deployment], error) {
return s.List(ctx, filter, offset, limit)
}
return pagination.NewIterator(fetchFunc, 100)
}
// Count returns the number of deployments matching the filter.
func (s *DeploymentsService) Count(ctx context.Context, filter *models.DeploymentFilter) (int, error) {
if filter == nil {
filter = &models.DeploymentFilter{}
}
var count int
if err := s.client.post(ctx, "/deployments/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count deployments: %w", err)
}
return count, nil
}
// GetSchedules retrieves all schedules for a deployment.
func (s *DeploymentsService) GetSchedules(ctx context.Context, id uuid.UUID) ([]models.DeploymentSchedule, error) {
var schedules []models.DeploymentSchedule
path := joinPath("/deployments", id.String(), "schedules")
if err := s.client.get(ctx, path, &schedules); err != nil {
return nil, fmt.Errorf("failed to get deployment schedules: %w", err)
}
return schedules, nil
}
// CreateSchedules creates schedules for a deployment.
func (s *DeploymentsService) CreateSchedules(ctx context.Context, id uuid.UUID, schedules []models.DeploymentScheduleCreate) ([]models.DeploymentSchedule, error) {
var result []models.DeploymentSchedule
path := joinPath("/deployments", id.String(), "schedules")
if err := s.client.post(ctx, path, schedules, &result); err != nil {
return nil, fmt.Errorf("failed to create deployment schedules: %w", err)
}
return result, nil
}
// UpdateSchedule updates a specific schedule for a deployment.
func (s *DeploymentsService) UpdateSchedule(ctx context.Context, deploymentID, scheduleID uuid.UUID, req *models.DeploymentScheduleUpdate) error {
path := joinPath("/deployments", deploymentID.String(), "schedules", scheduleID.String())
if err := s.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update deployment schedule: %w", err)
}
return nil
}
// DeleteSchedule deletes a specific schedule from a deployment.
func (s *DeploymentsService) DeleteSchedule(ctx context.Context, deploymentID, scheduleID uuid.UUID) error {
path := joinPath("/deployments", deploymentID.String(), "schedules", scheduleID.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete deployment schedule: %w", err)
}
return nil
}

57
pkg/client/events.go Normal file
View File

@@ -0,0 +1,57 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
// EventsService handles operations related to events.
type EventsService struct {
client *Client
}
// Create creates new events.
func (s *EventsService) Create(ctx context.Context, events []models.Event) error {
if err := s.client.post(ctx, "/events", events, nil); err != nil {
return fmt.Errorf("failed to create events: %w", err)
}
return nil
}
// List retrieves events with filtering.
func (s *EventsService) List(ctx context.Context, filter *models.EventFilter) (*models.EventPage, error) {
if filter == nil {
filter = &models.EventFilter{}
}
var page models.EventPage
if err := s.client.post(ctx, "/events/filter", filter, &page); err != nil {
return nil, fmt.Errorf("failed to list events: %w", err)
}
return &page, nil
}
// NextPage retrieves the next page of events using a page token.
func (s *EventsService) NextPage(ctx context.Context) (*models.EventPage, error) {
var page models.EventPage
if err := s.client.get(ctx, "/events/filter/next", &page); err != nil {
return nil, fmt.Errorf("failed to get next events page: %w", err)
}
return &page, nil
}
// CountBy counts events grouped by a countable field.
func (s *EventsService) CountBy(ctx context.Context, countable string, filter *models.EventFilter) ([]models.EventCount, error) {
if filter == nil {
filter = &models.EventFilter{}
}
path := joinPath("/events/count-by", countable)
var counts []models.EventCount
if err := s.client.post(ctx, path, filter, &counts); err != nil {
return nil, fmt.Errorf("failed to count events: %w", err)
}
return counts, nil
}

62
pkg/client/events_test.go Normal file
View File

@@ -0,0 +1,62 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestEventsService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/events/filter" {
t.Errorf("path = %v, want /api/events/filter", r.URL.Path)
}
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(models.EventPage{
Events: []models.Event{{ID: uuid.New(), Event: "test.event"}},
Total: 1,
})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
page, err := client.Events.List(ctx, nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if page.Total != 1 {
t.Errorf("total = %v, want 1", page.Total)
}
}
func TestEventsService_CountBy(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.EventCount{{Value: "test", Count: 5}})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
counts, err := client.Events.CountBy(ctx, "event", nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(counts) != 1 {
t.Errorf("count = %v, want 1", len(counts))
}
}

View File

@@ -0,0 +1,33 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// FlowRunStatesService handles operations related to flow run states.
type FlowRunStatesService struct {
client *Client
}
// List retrieves all flow run states.
func (s *FlowRunStatesService) List(ctx context.Context) ([]models.State, error) {
var states []models.State
if err := s.client.get(ctx, "/flow_run_states/", &states); err != nil {
return nil, fmt.Errorf("failed to list flow run states: %w", err)
}
return states, nil
}
// Get retrieves a flow run state by ID.
func (s *FlowRunStatesService) Get(ctx context.Context, id uuid.UUID) (*models.State, error) {
var state models.State
path := joinPath("/flow_run_states", id.String())
if err := s.client.get(ctx, path, &state); err != nil {
return nil, fmt.Errorf("failed to get flow run state: %w", err)
}
return &state, nil
}

View File

@@ -0,0 +1,63 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestFlowRunStatesService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/flow_run_states/" {
t.Errorf("path = %v, want /api/flow_run_states/", r.URL.Path)
}
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.State{
{ID: uuid.New(), Type: models.StateTypeCompleted},
})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
states, err := client.FlowRunStates.List(ctx)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(states) != 1 {
t.Errorf("count = %v, want 1", len(states))
}
}
func TestFlowRunStatesService_Get(t *testing.T) {
id := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(models.State{ID: id, Type: models.StateTypeRunning})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
state, err := client.FlowRunStates.Get(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if state.ID != id {
t.Errorf("ID = %v, want %v", state.ID, id)
}
}

View File

@@ -5,9 +5,9 @@ import (
"fmt"
"time"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
"github.com/google/uuid"
"github.com/gregor/prefect-go/pkg/models"
"github.com/gregor/prefect-go/pkg/pagination"
)
// FlowRunsService handles operations related to flow runs.
@@ -42,22 +42,17 @@ func (s *FlowRunsService) List(ctx context.Context, filter *models.FlowRunFilter
filter.Offset = offset
filter.Limit = limit
type response struct {
Results []models.FlowRun `json:"results"`
Count int `json:"count"`
}
var resp response
if err := s.client.post(ctx, "/flow_runs/filter", filter, &resp); err != nil {
var results []models.FlowRun
if err := s.client.post(ctx, "/flow_runs/filter", filter, &results); err != nil {
return nil, fmt.Errorf("failed to list flow runs: %w", err)
}
return &pagination.PaginatedResponse[models.FlowRun]{
Results: resp.Results,
Count: resp.Count,
Results: results,
Count: offset + len(results),
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
HasMore: len(results) == limit,
}, nil
}

View File

@@ -4,9 +4,9 @@ import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
"github.com/google/uuid"
"github.com/gregor/prefect-go/pkg/models"
"github.com/gregor/prefect-go/pkg/pagination"
)
// FlowsService handles operations related to flows.

View File

@@ -1,415 +0,0 @@
//go:build integration
// +build integration
package client_test
import (
"context"
"os"
"testing"
"time"
"github.com/google/uuid"
"github.com/gregor/prefect-go/pkg/client"
"github.com/gregor/prefect-go/pkg/errors"
"github.com/gregor/prefect-go/pkg/models"
)
var testClient *client.Client
func TestMain(m *testing.M) {
// Get base URL from environment or use default
baseURL := os.Getenv("PREFECT_API_URL")
if baseURL == "" {
baseURL = "http://localhost:4200/api"
}
// Create test client
var err error
testClient, err = client.NewClient(
client.WithBaseURL(baseURL),
)
if err != nil {
panic(err)
}
// Run tests
os.Exit(m.Run())
}
func TestIntegration_FlowLifecycle(t *testing.T) {
ctx := context.Background()
// Create a flow
flow, err := testClient.Flows.Create(ctx, &models.FlowCreate{
Name: "test-flow-" + uuid.New().String(),
Tags: []string{"integration-test"},
})
if err != nil {
t.Fatalf("Failed to create flow: %v", err)
}
defer testClient.Flows.Delete(ctx, flow.ID)
// Verify flow was created
if flow.ID == uuid.Nil {
t.Error("Flow ID is nil")
}
if flow.Name == "" {
t.Error("Flow name is empty")
}
// Get the flow
retrievedFlow, err := testClient.Flows.Get(ctx, flow.ID)
if err != nil {
t.Fatalf("Failed to get flow: %v", err)
}
if retrievedFlow.ID != flow.ID {
t.Errorf("Flow ID mismatch: got %v, want %v", retrievedFlow.ID, flow.ID)
}
// Update the flow
newTags := []string{"integration-test", "updated"}
updatedFlow, err := testClient.Flows.Update(ctx, flow.ID, &models.FlowUpdate{
Tags: &newTags,
})
if err != nil {
t.Fatalf("Failed to update flow: %v", err)
}
if len(updatedFlow.Tags) != 2 {
t.Errorf("Expected 2 tags, got %d", len(updatedFlow.Tags))
}
// List flows
flowsPage, err := testClient.Flows.List(ctx, &models.FlowFilter{
Tags: []string{"integration-test"},
}, 0, 10)
if err != nil {
t.Fatalf("Failed to list flows: %v", err)
}
if len(flowsPage.Results) == 0 {
t.Error("Expected at least one flow in results")
}
// Delete the flow
if err := testClient.Flows.Delete(ctx, flow.ID); err != nil {
t.Fatalf("Failed to delete flow: %v", err)
}
// Verify deletion
_, err = testClient.Flows.Get(ctx, flow.ID)
if !errors.IsNotFound(err) {
t.Errorf("Expected NotFound error, got: %v", err)
}
}
func TestIntegration_FlowRunLifecycle(t *testing.T) {
ctx := context.Background()
// Create a flow first
flow, err := testClient.Flows.Create(ctx, &models.FlowCreate{
Name: "test-flow-run-" + uuid.New().String(),
Tags: []string{"integration-test"},
})
if err != nil {
t.Fatalf("Failed to create flow: %v", err)
}
defer testClient.Flows.Delete(ctx, flow.ID)
// Create a flow run
flowRun, err := testClient.FlowRuns.Create(ctx, &models.FlowRunCreate{
FlowID: flow.ID,
Name: "test-run",
Parameters: map[string]interface{}{
"test_param": "value",
},
})
if err != nil {
t.Fatalf("Failed to create flow run: %v", err)
}
defer testClient.FlowRuns.Delete(ctx, flowRun.ID)
// Verify flow run was created
if flowRun.ID == uuid.Nil {
t.Error("Flow run ID is nil")
}
if flowRun.FlowID != flow.ID {
t.Errorf("Flow ID mismatch: got %v, want %v", flowRun.FlowID, flow.ID)
}
// Get the flow run
retrievedRun, err := testClient.FlowRuns.Get(ctx, flowRun.ID)
if err != nil {
t.Fatalf("Failed to get flow run: %v", err)
}
if retrievedRun.ID != flowRun.ID {
t.Errorf("Flow run ID mismatch: got %v, want %v", retrievedRun.ID, flowRun.ID)
}
// Set state to RUNNING
runningState := models.StateTypeRunning
updatedRun, err := testClient.FlowRuns.SetState(ctx, flowRun.ID, &models.StateCreate{
Type: runningState,
Message: strPtr("Test running"),
})
if err != nil {
t.Fatalf("Failed to set state: %v", err)
}
if updatedRun.StateType == nil || *updatedRun.StateType != runningState {
t.Errorf("State type mismatch: got %v, want %v", updatedRun.StateType, runningState)
}
// Set state to COMPLETED
completedState := models.StateTypeCompleted
completedRun, err := testClient.FlowRuns.SetState(ctx, flowRun.ID, &models.StateCreate{
Type: completedState,
Message: strPtr("Test completed"),
})
if err != nil {
t.Fatalf("Failed to set state: %v", err)
}
if completedRun.StateType == nil || *completedRun.StateType != completedState {
t.Errorf("State type mismatch: got %v, want %v", completedRun.StateType, completedState)
}
// List flow runs
runsPage, err := testClient.FlowRuns.List(ctx, &models.FlowRunFilter{
FlowID: &flow.ID,
}, 0, 10)
if err != nil {
t.Fatalf("Failed to list flow runs: %v", err)
}
if len(runsPage.Results) == 0 {
t.Error("Expected at least one flow run in results")
}
}
func TestIntegration_DeploymentLifecycle(t *testing.T) {
ctx := context.Background()
// Create a flow first
flow, err := testClient.Flows.Create(ctx, &models.FlowCreate{
Name: "test-deployment-flow-" + uuid.New().String(),
Tags: []string{"integration-test"},
})
if err != nil {
t.Fatalf("Failed to create flow: %v", err)
}
defer testClient.Flows.Delete(ctx, flow.ID)
// Create a deployment
workPoolName := "default-pool"
deployment, err := testClient.Deployments.Create(ctx, &models.DeploymentCreate{
Name: "test-deployment",
FlowID: flow.ID,
WorkPoolName: &workPoolName,
})
if err != nil {
t.Fatalf("Failed to create deployment: %v", err)
}
defer testClient.Deployments.Delete(ctx, deployment.ID)
// Verify deployment was created
if deployment.ID == uuid.Nil {
t.Error("Deployment ID is nil")
}
if deployment.FlowID != flow.ID {
t.Errorf("Flow ID mismatch: got %v, want %v", deployment.FlowID, flow.ID)
}
// Get the deployment
retrievedDeployment, err := testClient.Deployments.Get(ctx, deployment.ID)
if err != nil {
t.Fatalf("Failed to get deployment: %v", err)
}
if retrievedDeployment.ID != deployment.ID {
t.Errorf("Deployment ID mismatch: got %v, want %v", retrievedDeployment.ID, deployment.ID)
}
// Pause the deployment
if err := testClient.Deployments.Pause(ctx, deployment.ID); err != nil {
t.Fatalf("Failed to pause deployment: %v", err)
}
// Verify paused
pausedDeployment, err := testClient.Deployments.Get(ctx, deployment.ID)
if err != nil {
t.Fatalf("Failed to get deployment: %v", err)
}
if !pausedDeployment.Paused {
t.Error("Expected deployment to be paused")
}
// Resume the deployment
if err := testClient.Deployments.Resume(ctx, deployment.ID); err != nil {
t.Fatalf("Failed to resume deployment: %v", err)
}
// Verify resumed
resumedDeployment, err := testClient.Deployments.Get(ctx, deployment.ID)
if err != nil {
t.Fatalf("Failed to get deployment: %v", err)
}
if resumedDeployment.Paused {
t.Error("Expected deployment to be resumed")
}
}
func TestIntegration_Pagination(t *testing.T) {
ctx := context.Background()
// Create multiple flows
createdFlows := make([]uuid.UUID, 0)
for i := 0; i < 15; i++ {
flow, err := testClient.Flows.Create(ctx, &models.FlowCreate{
Name: "pagination-test-" + uuid.New().String(),
Tags: []string{"pagination-integration-test"},
})
if err != nil {
t.Fatalf("Failed to create flow %d: %v", i, err)
}
createdFlows = append(createdFlows, flow.ID)
}
// Clean up
defer func() {
for _, id := range createdFlows {
testClient.Flows.Delete(ctx, id)
}
}()
// Test manual pagination
page1, err := testClient.Flows.List(ctx, &models.FlowFilter{
Tags: []string{"pagination-integration-test"},
}, 0, 5)
if err != nil {
t.Fatalf("Failed to list flows: %v", err)
}
if len(page1.Results) != 5 {
t.Errorf("Expected 5 flows in page 1, got %d", len(page1.Results))
}
if !page1.HasMore {
t.Error("Expected more pages")
}
// Test iterator
iter := testClient.Flows.ListAll(ctx, &models.FlowFilter{
Tags: []string{"pagination-integration-test"},
})
count := 0
for iter.Next(ctx) {
count++
}
if err := iter.Err(); err != nil {
t.Fatalf("Iterator error: %v", err)
}
if count != 15 {
t.Errorf("Expected to iterate over 15 flows, got %d", count)
}
}
func TestIntegration_Variables(t *testing.T) {
ctx := context.Background()
// Create a variable
variable, err := testClient.Variables.Create(ctx, &models.VariableCreate{
Name: "test-var-" + uuid.New().String(),
Value: "test-value",
Tags: []string{"integration-test"},
})
if err != nil {
t.Fatalf("Failed to create variable: %v", err)
}
defer testClient.Variables.Delete(ctx, variable.ID)
// Get the variable
retrievedVar, err := testClient.Variables.Get(ctx, variable.ID)
if err != nil {
t.Fatalf("Failed to get variable: %v", err)
}
if retrievedVar.Value != "test-value" {
t.Errorf("Value mismatch: got %v, want test-value", retrievedVar.Value)
}
// Update the variable
newValue := "updated-value"
updatedVar, err := testClient.Variables.Update(ctx, variable.ID, &models.VariableUpdate{
Value: &newValue,
})
if err != nil {
t.Fatalf("Failed to update variable: %v", err)
}
if updatedVar.Value != newValue {
t.Errorf("Value mismatch: got %v, want %v", updatedVar.Value, newValue)
}
}
func TestIntegration_AdminEndpoints(t *testing.T) {
ctx := context.Background()
// Test health check
if err := testClient.Admin.Health(ctx); err != nil {
t.Fatalf("Health check failed: %v", err)
}
// Test version
version, err := testClient.Admin.Version(ctx)
if err != nil {
t.Fatalf("Failed to get version: %v", err)
}
if version == "" {
t.Error("Version is empty")
}
t.Logf("Server version: %s", version)
}
func TestIntegration_FlowRunWait(t *testing.T) {
ctx := context.Background()
// Create a flow
flow, err := testClient.Flows.Create(ctx, &models.FlowCreate{
Name: "wait-test-" + uuid.New().String(),
Tags: []string{"integration-test"},
})
if err != nil {
t.Fatalf("Failed to create flow: %v", err)
}
defer testClient.Flows.Delete(ctx, flow.ID)
// Create a flow run
flowRun, err := testClient.FlowRuns.Create(ctx, &models.FlowRunCreate{
FlowID: flow.ID,
Name: "wait-test-run",
})
if err != nil {
t.Fatalf("Failed to create flow run: %v", err)
}
defer testClient.FlowRuns.Delete(ctx, flowRun.ID)
// Simulate completion in background
go func() {
time.Sleep(2 * time.Second)
completedState := models.StateTypeCompleted
testClient.FlowRuns.SetState(ctx, flowRun.ID, &models.StateCreate{
Type: completedState,
Message: strPtr("Test completed"),
})
}()
// Wait for completion with timeout
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
finalRun, err := testClient.FlowRuns.Wait(waitCtx, flowRun.ID, time.Second)
if err != nil {
t.Fatalf("Failed to wait for flow run: %v", err)
}
if finalRun.StateType == nil || *finalRun.StateType != models.StateTypeCompleted {
t.Errorf("Expected COMPLETED state, got %v", finalRun.StateType)
}
}
func strPtr(s string) *string {
return &s
}

55
pkg/client/logs.go Normal file
View File

@@ -0,0 +1,55 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
)
// LogsService handles operations related to logs.
type LogsService struct {
client *Client
}
// Create creates new log entries.
func (l *LogsService) Create(ctx context.Context, logs []*models.LogCreate) error {
if err := l.client.post(ctx, "/logs/", logs, nil); err != nil {
return fmt.Errorf("failed to create logs: %w", err)
}
return nil
}
// List retrieves logs with filtering.
func (l *LogsService) List(ctx context.Context, filter interface{}, offset, limit int) (*pagination.PaginatedResponse[models.Log], error) {
type request struct {
Filter interface{} `json:"filter,omitempty"`
Offset int `json:"offset"`
Limit int `json:"limit"`
}
req := request{
Filter: filter,
Offset: offset,
Limit: limit,
}
type response struct {
Results []models.Log `json:"results"`
Count int `json:"count"`
}
var resp response
if err := l.client.post(ctx, "/logs/filter", req, &resp); err != nil {
return nil, fmt.Errorf("failed to list logs: %w", err)
}
return &pagination.PaginatedResponse[models.Log]{
Results: resp.Results,
Count: resp.Count,
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
}, nil
}

View File

@@ -0,0 +1,52 @@
package client
import (
"context"
"fmt"
"net/http"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// SavedSearchesService handles operations related to saved searches.
type SavedSearchesService struct {
client *Client
}
// Create creates a new saved search.
func (s *SavedSearchesService) Create(ctx context.Context, req *models.SavedSearchCreate) (*models.SavedSearch, error) {
var savedSearch models.SavedSearch
if err := s.client.do(ctx, http.MethodPut, "/saved_searches/", req, &savedSearch); err != nil {
return nil, fmt.Errorf("failed to create saved search: %w", err)
}
return &savedSearch, nil
}
// Get retrieves a saved search by ID.
func (s *SavedSearchesService) Get(ctx context.Context, id uuid.UUID) (*models.SavedSearch, error) {
var savedSearch models.SavedSearch
path := joinPath("/saved_searches", id.String())
if err := s.client.get(ctx, path, &savedSearch); err != nil {
return nil, fmt.Errorf("failed to get saved search: %w", err)
}
return &savedSearch, nil
}
// Delete deletes a saved search by ID.
func (s *SavedSearchesService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/saved_searches", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete saved search: %w", err)
}
return nil
}
// List retrieves all saved searches.
func (s *SavedSearchesService) List(ctx context.Context) ([]models.SavedSearch, error) {
var savedSearches []models.SavedSearch
if err := s.client.post(ctx, "/saved_searches/filter", nil, &savedSearches); err != nil {
return nil, fmt.Errorf("failed to list saved searches: %w", err)
}
return savedSearches, nil
}

View File

@@ -0,0 +1,103 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestSavedSearchesService_Create(t *testing.T) {
expected := models.SavedSearch{ID: uuid.New(), Name: "test-search"}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPut {
t.Errorf("method = %v, want PUT", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(expected)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
search, err := client.SavedSearches.Create(ctx, &models.SavedSearchCreate{Name: "test-search"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if search.Name != expected.Name {
t.Errorf("Name = %v, want %v", search.Name, expected.Name)
}
}
func TestSavedSearchesService_Get(t *testing.T) {
id := uuid.New()
expected := models.SavedSearch{ID: id, Name: "test"}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(expected)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
search, err := client.SavedSearches.Get(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if search.ID != id {
t.Errorf("ID = %v, want %v", search.ID, id)
}
}
func TestSavedSearchesService_Delete(t *testing.T) {
id := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodDelete {
t.Errorf("method = %v, want DELETE", r.Method)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
err := client.SavedSearches.Delete(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
func TestSavedSearchesService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/saved_searches/filter" {
t.Errorf("path = %v, want /api/saved_searches/filter", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.SavedSearch{{ID: uuid.New(), Name: "test"}})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
searches, err := client.SavedSearches.List(ctx)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(searches) != 1 {
t.Errorf("count = %v, want 1", len(searches))
}
}

View File

@@ -1,301 +0,0 @@
package client
import (
"context"
"fmt"
"github.com/google/uuid"
"github.com/gregor/prefect-go/pkg/models"
"github.com/gregor/prefect-go/pkg/pagination"
)
// DeploymentsService handles operations related to deployments.
type DeploymentsService struct {
client *Client
}
// Create creates a new deployment.
func (s *DeploymentsService) Create(ctx context.Context, req *models.DeploymentCreate) (*models.Deployment, error) {
var deployment models.Deployment
if err := s.client.post(ctx, "/deployments/", req, &deployment); err != nil {
return nil, fmt.Errorf("failed to create deployment: %w", err)
}
return &deployment, nil
}
// Get retrieves a deployment by ID.
func (s *DeploymentsService) Get(ctx context.Context, id uuid.UUID) (*models.Deployment, error) {
var deployment models.Deployment
path := joinPath("/deployments", id.String())
if err := s.client.get(ctx, path, &deployment); err != nil {
return nil, fmt.Errorf("failed to get deployment: %w", err)
}
return &deployment, nil
}
// GetByName retrieves a deployment by flow name and deployment name.
func (s *DeploymentsService) GetByName(ctx context.Context, flowName, deploymentName string) (*models.Deployment, error) {
var deployment models.Deployment
path := fmt.Sprintf("/deployments/name/%s/%s", flowName, deploymentName)
if err := s.client.get(ctx, path, &deployment); err != nil {
return nil, fmt.Errorf("failed to get deployment by name: %w", err)
}
return &deployment, nil
}
// Update updates a deployment.
func (s *DeploymentsService) Update(ctx context.Context, id uuid.UUID, req *models.DeploymentUpdate) (*models.Deployment, error) {
var deployment models.Deployment
path := joinPath("/deployments", id.String())
if err := s.client.patch(ctx, path, req, &deployment); err != nil {
return nil, fmt.Errorf("failed to update deployment: %w", err)
}
return &deployment, nil
}
// Delete deletes a deployment by ID.
func (s *DeploymentsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/deployments", id.String())
if err := s.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete deployment: %w", err)
}
return nil
}
// Pause pauses a deployment's schedule.
func (s *DeploymentsService) Pause(ctx context.Context, id uuid.UUID) error {
path := joinPath("/deployments", id.String(), "pause")
if err := s.client.post(ctx, path, nil, nil); err != nil {
return fmt.Errorf("failed to pause deployment: %w", err)
}
return nil
}
// Resume resumes a deployment's schedule.
func (s *DeploymentsService) Resume(ctx context.Context, id uuid.UUID) error {
path := joinPath("/deployments", id.String(), "resume")
if err := s.client.post(ctx, path, nil, nil); err != nil {
return fmt.Errorf("failed to resume deployment: %w", err)
}
return nil
}
// CreateFlowRun creates a flow run from a deployment.
func (s *DeploymentsService) CreateFlowRun(ctx context.Context, id uuid.UUID, params map[string]interface{}) (*models.FlowRun, error) {
var flowRun models.FlowRun
path := joinPath("/deployments", id.String(), "create_flow_run")
req := struct {
Parameters map[string]interface{} `json:"parameters,omitempty"`
}{
Parameters: params,
}
if err := s.client.post(ctx, path, req, &flowRun); err != nil {
return nil, fmt.Errorf("failed to create flow run from deployment: %w", err)
}
return &flowRun, nil
}
// TaskRunsService handles operations related to task runs.
type TaskRunsService struct {
client *Client
}
// Create creates a new task run.
func (t *TaskRunsService) Create(ctx context.Context, req *models.TaskRunCreate) (*models.TaskRun, error) {
var taskRun models.TaskRun
if err := t.client.post(ctx, "/task_runs/", req, &taskRun); err != nil {
return nil, fmt.Errorf("failed to create task run: %w", err)
}
return &taskRun, nil
}
// Get retrieves a task run by ID.
func (t *TaskRunsService) Get(ctx context.Context, id uuid.UUID) (*models.TaskRun, error) {
var taskRun models.TaskRun
path := joinPath("/task_runs", id.String())
if err := t.client.get(ctx, path, &taskRun); err != nil {
return nil, fmt.Errorf("failed to get task run: %w", err)
}
return &taskRun, nil
}
// Delete deletes a task run by ID.
func (t *TaskRunsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/task_runs", id.String())
if err := t.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete task run: %w", err)
}
return nil
}
// SetState sets the state of a task run.
func (t *TaskRunsService) SetState(ctx context.Context, id uuid.UUID, state *models.StateCreate) (*models.TaskRun, error) {
var taskRun models.TaskRun
path := joinPath("/task_runs", id.String(), "set_state")
req := struct {
State *models.StateCreate `json:"state"`
}{
State: state,
}
if err := t.client.post(ctx, path, req, &taskRun); err != nil {
return nil, fmt.Errorf("failed to set task run state: %w", err)
}
return &taskRun, nil
}
// WorkPoolsService handles operations related to work pools.
type WorkPoolsService struct {
client *Client
}
// Get retrieves a work pool by name.
func (w *WorkPoolsService) Get(ctx context.Context, name string) (*models.WorkPool, error) {
var workPool models.WorkPool
path := joinPath("/work_pools", name)
if err := w.client.get(ctx, path, &workPool); err != nil {
return nil, fmt.Errorf("failed to get work pool: %w", err)
}
return &workPool, nil
}
// WorkQueuesService handles operations related to work queues.
type WorkQueuesService struct {
client *Client
}
// Get retrieves a work queue by ID.
func (w *WorkQueuesService) Get(ctx context.Context, id uuid.UUID) (*models.WorkQueue, error) {
var workQueue models.WorkQueue
path := joinPath("/work_queues", id.String())
if err := w.client.get(ctx, path, &workQueue); err != nil {
return nil, fmt.Errorf("failed to get work queue: %w", err)
}
return &workQueue, nil
}
// VariablesService handles operations related to variables.
type VariablesService struct {
client *Client
}
// Create creates a new variable.
func (v *VariablesService) Create(ctx context.Context, req *models.VariableCreate) (*models.Variable, error) {
var variable models.Variable
if err := v.client.post(ctx, "/variables/", req, &variable); err != nil {
return nil, fmt.Errorf("failed to create variable: %w", err)
}
return &variable, nil
}
// Get retrieves a variable by ID.
func (v *VariablesService) Get(ctx context.Context, id uuid.UUID) (*models.Variable, error) {
var variable models.Variable
path := joinPath("/variables", id.String())
if err := v.client.get(ctx, path, &variable); err != nil {
return nil, fmt.Errorf("failed to get variable: %w", err)
}
return &variable, nil
}
// GetByName retrieves a variable by name.
func (v *VariablesService) GetByName(ctx context.Context, name string) (*models.Variable, error) {
var variable models.Variable
path := joinPath("/variables/name", name)
if err := v.client.get(ctx, path, &variable); err != nil {
return nil, fmt.Errorf("failed to get variable by name: %w", err)
}
return &variable, nil
}
// Update updates a variable.
func (v *VariablesService) Update(ctx context.Context, id uuid.UUID, req *models.VariableUpdate) (*models.Variable, error) {
var variable models.Variable
path := joinPath("/variables", id.String())
if err := v.client.patch(ctx, path, req, &variable); err != nil {
return nil, fmt.Errorf("failed to update variable: %w", err)
}
return &variable, nil
}
// Delete deletes a variable by ID.
func (v *VariablesService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/variables", id.String())
if err := v.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete variable: %w", err)
}
return nil
}
// LogsService handles operations related to logs.
type LogsService struct {
client *Client
}
// Create creates new log entries.
func (l *LogsService) Create(ctx context.Context, logs []*models.LogCreate) error {
if err := l.client.post(ctx, "/logs/", logs, nil); err != nil {
return fmt.Errorf("failed to create logs: %w", err)
}
return nil
}
// List retrieves logs with filtering.
func (l *LogsService) List(ctx context.Context, filter interface{}, offset, limit int) (*pagination.PaginatedResponse[models.Log], error) {
type request struct {
Filter interface{} `json:"filter,omitempty"`
Offset int `json:"offset"`
Limit int `json:"limit"`
}
req := request{
Filter: filter,
Offset: offset,
Limit: limit,
}
type response struct {
Results []models.Log `json:"results"`
Count int `json:"count"`
}
var resp response
if err := l.client.post(ctx, "/logs/filter", req, &resp); err != nil {
return nil, fmt.Errorf("failed to list logs: %w", err)
}
return &pagination.PaginatedResponse[models.Log]{
Results: resp.Results,
Count: resp.Count,
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
}, nil
}
// AdminService handles administrative operations.
type AdminService struct {
client *Client
}
// Health checks the health of the Prefect server.
func (a *AdminService) Health(ctx context.Context) error {
if err := a.client.get(ctx, "/health", nil); err != nil {
return fmt.Errorf("health check failed: %w", err)
}
return nil
}
// Version retrieves the server version.
func (a *AdminService) Version(ctx context.Context) (string, error) {
var result struct {
Version string `json:"version"`
}
if err := a.client.get(ctx, "/version", &result); err != nil {
return "", fmt.Errorf("failed to get version: %w", err)
}
return result.Version, nil
}

View File

@@ -0,0 +1,33 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// TaskRunStatesService handles operations related to task run states.
type TaskRunStatesService struct {
client *Client
}
// List retrieves all task run states.
func (s *TaskRunStatesService) List(ctx context.Context) ([]models.State, error) {
var states []models.State
if err := s.client.get(ctx, "/task_run_states/", &states); err != nil {
return nil, fmt.Errorf("failed to list task run states: %w", err)
}
return states, nil
}
// Get retrieves a task run state by ID.
func (s *TaskRunStatesService) Get(ctx context.Context, id uuid.UUID) (*models.State, error) {
var state models.State
path := joinPath("/task_run_states", id.String())
if err := s.client.get(ctx, path, &state); err != nil {
return nil, fmt.Errorf("failed to get task run state: %w", err)
}
return &state, nil
}

View File

@@ -0,0 +1,63 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
func TestTaskRunStatesService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/task_run_states/" {
t.Errorf("path = %v, want /api/task_run_states/", r.URL.Path)
}
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.State{
{ID: uuid.New(), Type: models.StateTypeFailed},
})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
states, err := client.TaskRunStates.List(ctx)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(states) != 1 {
t.Errorf("count = %v, want 1", len(states))
}
}
func TestTaskRunStatesService_Get(t *testing.T) {
id := uuid.New()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %v, want GET", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(models.State{ID: id, Type: models.StateTypePending})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
state, err := client.TaskRunStates.Get(ctx, id)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if state.ID != id {
t.Errorf("ID = %v, want %v", state.ID, id)
}
}

118
pkg/client/task_runs.go Normal file
View File

@@ -0,0 +1,118 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
"github.com/google/uuid"
)
// TaskRunsService handles operations related to task runs.
type TaskRunsService struct {
client *Client
}
// Create creates a new task run.
func (t *TaskRunsService) Create(ctx context.Context, req *models.TaskRunCreate) (*models.TaskRun, error) {
var taskRun models.TaskRun
if err := t.client.post(ctx, "/task_runs/", req, &taskRun); err != nil {
return nil, fmt.Errorf("failed to create task run: %w", err)
}
return &taskRun, nil
}
// Get retrieves a task run by ID.
func (t *TaskRunsService) Get(ctx context.Context, id uuid.UUID) (*models.TaskRun, error) {
var taskRun models.TaskRun
path := joinPath("/task_runs", id.String())
if err := t.client.get(ctx, path, &taskRun); err != nil {
return nil, fmt.Errorf("failed to get task run: %w", err)
}
return &taskRun, nil
}
// Delete deletes a task run by ID.
func (t *TaskRunsService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/task_runs", id.String())
if err := t.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete task run: %w", err)
}
return nil
}
// SetState sets the state of a task run.
func (t *TaskRunsService) SetState(ctx context.Context, id uuid.UUID, state *models.StateCreate) (*models.TaskRun, error) {
var taskRun models.TaskRun
path := joinPath("/task_runs", id.String(), "set_state")
req := struct {
State *models.StateCreate `json:"state"`
}{
State: state,
}
if err := t.client.post(ctx, path, req, &taskRun); err != nil {
return nil, fmt.Errorf("failed to set task run state: %w", err)
}
return &taskRun, nil
}
// Update updates a task run.
func (t *TaskRunsService) Update(ctx context.Context, id uuid.UUID, req *models.TaskRunUpdate) (*models.TaskRun, error) {
var taskRun models.TaskRun
path := joinPath("/task_runs", id.String())
if err := t.client.patch(ctx, path, req, &taskRun); err != nil {
return nil, fmt.Errorf("failed to update task run: %w", err)
}
return &taskRun, nil
}
// List retrieves a list of task runs with optional filtering.
func (t *TaskRunsService) List(ctx context.Context, filter *models.TaskRunFilter, offset, limit int) (*pagination.PaginatedResponse[models.TaskRun], error) {
if filter == nil {
filter = &models.TaskRunFilter{}
}
filter.Offset = offset
filter.Limit = limit
type response struct {
Results []models.TaskRun `json:"results"`
Count int `json:"count"`
}
var resp response
if err := t.client.post(ctx, "/task_runs/filter", filter, &resp); err != nil {
return nil, fmt.Errorf("failed to list task runs: %w", err)
}
return &pagination.PaginatedResponse[models.TaskRun]{
Results: resp.Results,
Count: resp.Count,
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
}, nil
}
// ListAll returns an iterator for all task runs matching the filter.
func (t *TaskRunsService) ListAll(ctx context.Context, filter *models.TaskRunFilter) *pagination.Iterator[models.TaskRun] {
fetchFunc := func(ctx context.Context, offset, limit int) (*pagination.PaginatedResponse[models.TaskRun], error) {
return t.List(ctx, filter, offset, limit)
}
return pagination.NewIterator(fetchFunc, 100)
}
// Count returns the number of task runs matching the filter.
func (t *TaskRunsService) Count(ctx context.Context, filter *models.TaskRunFilter) (int, error) {
if filter == nil {
filter = &models.TaskRunFilter{}
}
var count int
if err := t.client.post(ctx, "/task_runs/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count task runs: %w", err)
}
return count, nil
}

View File

@@ -0,0 +1,26 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
// TaskWorkersService handles operations related to task workers.
type TaskWorkersService struct {
client *Client
}
// List retrieves task workers matching the filter.
func (s *TaskWorkersService) List(ctx context.Context, filter *models.TaskWorkerFilter) ([]models.TaskWorkerResponse, error) {
if filter == nil {
filter = &models.TaskWorkerFilter{}
}
var workers []models.TaskWorkerResponse
if err := s.client.post(ctx, "/task_workers/filter", filter, &workers); err != nil {
return nil, fmt.Errorf("failed to list task workers: %w", err)
}
return workers, nil
}

View File

@@ -0,0 +1,41 @@
package client
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
func TestTaskWorkersService_List(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/task_workers/filter" {
t.Errorf("path = %v, want /api/task_workers/filter", r.URL.Path)
}
if r.Method != http.MethodPost {
t.Errorf("method = %v, want POST", r.Method)
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode([]models.TaskWorkerResponse{
{Identifier: "worker-1", TaskKeys: []string{"key1"}},
})
}))
defer server.Close()
client, _ := NewClient(WithBaseURL(server.URL + "/api"))
ctx := context.Background()
workers, err := client.TaskWorkers.List(ctx, &models.TaskWorkerFilter{TaskKeys: []string{"key1"}})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(workers) != 1 {
t.Errorf("count = %v, want 1", len(workers))
}
if workers[0].Identifier != "worker-1" {
t.Errorf("Identifier = %v, want worker-1", workers[0].Identifier)
}
}

111
pkg/client/variables.go Normal file
View File

@@ -0,0 +1,111 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"git.schultes.dev/schultesdev/prefect-go/pkg/pagination"
"github.com/google/uuid"
)
// VariablesService handles operations related to variables.
type VariablesService struct {
client *Client
}
// Create creates a new variable.
func (v *VariablesService) Create(ctx context.Context, req *models.VariableCreate) (*models.Variable, error) {
var variable models.Variable
if err := v.client.post(ctx, "/variables/", req, &variable); err != nil {
return nil, fmt.Errorf("failed to create variable: %w", err)
}
return &variable, nil
}
// Get retrieves a variable by ID.
func (v *VariablesService) Get(ctx context.Context, id uuid.UUID) (*models.Variable, error) {
var variable models.Variable
path := joinPath("/variables", id.String())
if err := v.client.get(ctx, path, &variable); err != nil {
return nil, fmt.Errorf("failed to get variable: %w", err)
}
return &variable, nil
}
// GetByName retrieves a variable by name.
func (v *VariablesService) GetByName(ctx context.Context, name string) (*models.Variable, error) {
var variable models.Variable
path := joinPath("/variables/name", name)
if err := v.client.get(ctx, path, &variable); err != nil {
return nil, fmt.Errorf("failed to get variable by name: %w", err)
}
return &variable, nil
}
// Update updates a variable.
func (v *VariablesService) Update(ctx context.Context, id uuid.UUID, req *models.VariableUpdate) (*models.Variable, error) {
var variable models.Variable
path := joinPath("/variables", id.String())
if err := v.client.patch(ctx, path, req, &variable); err != nil {
return nil, fmt.Errorf("failed to update variable: %w", err)
}
return &variable, nil
}
// Delete deletes a variable by ID.
func (v *VariablesService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/variables", id.String())
if err := v.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete variable: %w", err)
}
return nil
}
// List retrieves a list of variables with optional filtering.
func (v *VariablesService) List(ctx context.Context, filter *models.VariableFilter, offset, limit int) (*pagination.PaginatedResponse[models.Variable], error) {
if filter == nil {
filter = &models.VariableFilter{}
}
filter.Offset = offset
filter.Limit = limit
type response struct {
Results []models.Variable `json:"results"`
Count int `json:"count"`
}
var resp response
if err := v.client.post(ctx, "/variables/filter", filter, &resp); err != nil {
return nil, fmt.Errorf("failed to list variables: %w", err)
}
return &pagination.PaginatedResponse[models.Variable]{
Results: resp.Results,
Count: resp.Count,
Limit: limit,
Offset: offset,
HasMore: offset+len(resp.Results) < resp.Count,
}, nil
}
// ListAll returns an iterator for all variables matching the filter.
func (v *VariablesService) ListAll(ctx context.Context, filter *models.VariableFilter) *pagination.Iterator[models.Variable] {
fetchFunc := func(ctx context.Context, offset, limit int) (*pagination.PaginatedResponse[models.Variable], error) {
return v.List(ctx, filter, offset, limit)
}
return pagination.NewIterator(fetchFunc, 100)
}
// Count returns the number of variables matching the filter.
func (v *VariablesService) Count(ctx context.Context, filter *models.VariableFilter) (int, error) {
if filter == nil {
filter = &models.VariableFilter{}
}
var count int
if err := v.client.post(ctx, "/variables/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count variables: %w", err)
}
return count, nil
}

78
pkg/client/work_pools.go Normal file
View File

@@ -0,0 +1,78 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
)
// WorkPoolsService handles operations related to work pools.
type WorkPoolsService struct {
client *Client
}
// Get retrieves a work pool by name.
func (w *WorkPoolsService) Get(ctx context.Context, name string) (*models.WorkPool, error) {
var workPool models.WorkPool
path := joinPath("/work_pools", name)
if err := w.client.get(ctx, path, &workPool); err != nil {
return nil, fmt.Errorf("failed to get work pool: %w", err)
}
return &workPool, nil
}
// Create creates a new work pool.
func (w *WorkPoolsService) Create(ctx context.Context, req *models.WorkPoolCreate) (*models.WorkPool, error) {
var workPool models.WorkPool
if err := w.client.post(ctx, "/work_pools/", req, &workPool); err != nil {
return nil, fmt.Errorf("failed to create work pool: %w", err)
}
return &workPool, nil
}
// Update updates a work pool.
func (w *WorkPoolsService) Update(ctx context.Context, name string, req *models.WorkPoolUpdate) error {
path := joinPath("/work_pools", name)
if err := w.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update work pool: %w", err)
}
return nil
}
// Delete deletes a work pool by name.
func (w *WorkPoolsService) Delete(ctx context.Context, name string) error {
path := joinPath("/work_pools", name)
if err := w.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete work pool: %w", err)
}
return nil
}
// List retrieves work pools with optional filtering.
func (w *WorkPoolsService) List(ctx context.Context, filter *models.WorkPoolFilter, offset, limit int) ([]models.WorkPool, error) {
if filter == nil {
filter = &models.WorkPoolFilter{}
}
filter.Offset = offset
filter.Limit = limit
var workPools []models.WorkPool
if err := w.client.post(ctx, "/work_pools/filter", filter, &workPools); err != nil {
return nil, fmt.Errorf("failed to list work pools: %w", err)
}
return workPools, nil
}
// Count returns the number of work pools matching the filter.
func (w *WorkPoolsService) Count(ctx context.Context, filter *models.WorkPoolFilter) (int, error) {
if filter == nil {
filter = &models.WorkPoolFilter{}
}
var count int
if err := w.client.post(ctx, "/work_pools/count", filter, &count); err != nil {
return 0, fmt.Errorf("failed to count work pools: %w", err)
}
return count, nil
}

76
pkg/client/work_queues.go Normal file
View File

@@ -0,0 +1,76 @@
package client
import (
"context"
"fmt"
"git.schultes.dev/schultesdev/prefect-go/pkg/models"
"github.com/google/uuid"
)
// WorkQueuesService handles operations related to work queues.
type WorkQueuesService struct {
client *Client
}
// Get retrieves a work queue by ID.
func (w *WorkQueuesService) Get(ctx context.Context, id uuid.UUID) (*models.WorkQueue, error) {
var workQueue models.WorkQueue
path := joinPath("/work_queues", id.String())
if err := w.client.get(ctx, path, &workQueue); err != nil {
return nil, fmt.Errorf("failed to get work queue: %w", err)
}
return &workQueue, nil
}
// Create creates a new work queue.
func (w *WorkQueuesService) Create(ctx context.Context, req *models.WorkQueueCreate) (*models.WorkQueue, error) {
var workQueue models.WorkQueue
if err := w.client.post(ctx, "/work_queues/", req, &workQueue); err != nil {
return nil, fmt.Errorf("failed to create work queue: %w", err)
}
return &workQueue, nil
}
// GetByName retrieves a work queue by name.
func (w *WorkQueuesService) GetByName(ctx context.Context, name string) (*models.WorkQueue, error) {
var workQueue models.WorkQueue
path := joinPath("/work_queues/name", name)
if err := w.client.get(ctx, path, &workQueue); err != nil {
return nil, fmt.Errorf("failed to get work queue by name: %w", err)
}
return &workQueue, nil
}
// Update updates a work queue.
func (w *WorkQueuesService) Update(ctx context.Context, id uuid.UUID, req *models.WorkQueueUpdate) error {
path := joinPath("/work_queues", id.String())
if err := w.client.patch(ctx, path, req, nil); err != nil {
return fmt.Errorf("failed to update work queue: %w", err)
}
return nil
}
// Delete deletes a work queue by ID.
func (w *WorkQueuesService) Delete(ctx context.Context, id uuid.UUID) error {
path := joinPath("/work_queues", id.String())
if err := w.client.delete(ctx, path); err != nil {
return fmt.Errorf("failed to delete work queue: %w", err)
}
return nil
}
// List retrieves work queues with optional filtering.
func (w *WorkQueuesService) List(ctx context.Context, filter *models.WorkQueueFilter, offset, limit int) ([]models.WorkQueue, error) {
if filter == nil {
filter = &models.WorkQueueFilter{}
}
filter.Offset = offset
filter.Limit = limit
var workQueues []models.WorkQueue
if err := w.client.post(ctx, "/work_queues/filter", filter, &workQueues); err != nil {
return nil, fmt.Errorf("failed to list work queues: %w", err)
}
return workQueues, nil
}

View File

@@ -3,6 +3,7 @@ package errors
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
@@ -98,7 +99,12 @@ func NewAPIError(resp *http.Response) error {
// Try to read and parse the response body
if resp.Body != nil {
defer resp.Body.Close()
defer func(Body io.ReadCloser) {
err := Body.Close()
if err != nil {
return
}
}(resp.Body)
body, err := io.ReadAll(resp.Body)
if err == nil && len(body) > 0 {
// Try to parse as JSON
@@ -181,7 +187,8 @@ func parseValidationError(apiErr *APIError, detail interface{}) error {
// IsNotFound checks if an error is a 404 Not Found error.
func IsNotFound(err error) bool {
if apiErr, ok := err.(*APIError); ok {
var apiErr *APIError
if errors.As(err, &apiErr) {
return apiErr.IsNotFound()
}
return false
@@ -189,7 +196,8 @@ func IsNotFound(err error) bool {
// IsUnauthorized checks if an error is a 401 Unauthorized error.
func IsUnauthorized(err error) bool {
if apiErr, ok := err.(*APIError); ok {
var apiErr *APIError
if errors.As(err, &apiErr) {
return apiErr.IsUnauthorized()
}
return false
@@ -197,7 +205,8 @@ func IsUnauthorized(err error) bool {
// IsForbidden checks if an error is a 403 Forbidden error.
func IsForbidden(err error) bool {
if apiErr, ok := err.(*APIError); ok {
var apiErr *APIError
if errors.As(err, &apiErr) {
return apiErr.IsForbidden()
}
return false
@@ -205,7 +214,8 @@ func IsForbidden(err error) bool {
// IsRateLimited checks if an error is a 429 Too Many Requests error.
func IsRateLimited(err error) bool {
if apiErr, ok := err.(*APIError); ok {
var apiErr *APIError
if errors.As(err, &apiErr) {
return apiErr.IsRateLimited()
}
return false
@@ -213,7 +223,8 @@ func IsRateLimited(err error) bool {
// IsServerError checks if an error is a 5xx server error.
func IsServerError(err error) bool {
if apiErr, ok := err.(*APIError); ok {
var apiErr *APIError
if errors.As(err, &apiErr) {
return apiErr.IsServerError()
}
return false
@@ -221,6 +232,7 @@ func IsServerError(err error) bool {
// IsValidationError checks if an error is a validation error.
func IsValidationError(err error) bool {
_, ok := err.(*ValidationError)
var validationError *ValidationError
ok := errors.As(err, &validationError)
return ok
}

View File

@@ -1,240 +0,0 @@
package errors
import (
"io"
"net/http"
"strings"
"testing"
)
func TestAPIError_Error(t *testing.T) {
tests := []struct {
name string
err *APIError
want string
}{
{
name: "with message",
err: &APIError{
StatusCode: 404,
Message: "Flow not found",
},
want: "prefect api error (status 404): Flow not found",
},
{
name: "without message",
err: &APIError{
StatusCode: 500,
},
want: "prefect api error (status 500)",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.err.Error(); got != tt.want {
t.Errorf("Error() = %v, want %v", got, tt.want)
}
})
}
}
func TestAPIError_Checks(t *testing.T) {
tests := []struct {
name string
statusCode int
checkFuncs map[string]func(*APIError) bool
want map[string]bool
}{
{
name: "404 not found",
statusCode: 404,
checkFuncs: map[string]func(*APIError) bool{
"IsNotFound": (*APIError).IsNotFound,
"IsUnauthorized": (*APIError).IsUnauthorized,
"IsForbidden": (*APIError).IsForbidden,
"IsRateLimited": (*APIError).IsRateLimited,
"IsServerError": (*APIError).IsServerError,
},
want: map[string]bool{
"IsNotFound": true,
"IsUnauthorized": false,
"IsForbidden": false,
"IsRateLimited": false,
"IsServerError": false,
},
},
{
name: "401 unauthorized",
statusCode: 401,
checkFuncs: map[string]func(*APIError) bool{
"IsNotFound": (*APIError).IsNotFound,
"IsUnauthorized": (*APIError).IsUnauthorized,
"IsForbidden": (*APIError).IsForbidden,
"IsRateLimited": (*APIError).IsRateLimited,
"IsServerError": (*APIError).IsServerError,
},
want: map[string]bool{
"IsNotFound": false,
"IsUnauthorized": true,
"IsForbidden": false,
"IsRateLimited": false,
"IsServerError": false,
},
},
{
name: "500 server error",
statusCode: 500,
checkFuncs: map[string]func(*APIError) bool{
"IsNotFound": (*APIError).IsNotFound,
"IsUnauthorized": (*APIError).IsUnauthorized,
"IsForbidden": (*APIError).IsForbidden,
"IsRateLimited": (*APIError).IsRateLimited,
"IsServerError": (*APIError).IsServerError,
},
want: map[string]bool{
"IsNotFound": false,
"IsUnauthorized": false,
"IsForbidden": false,
"IsRateLimited": false,
"IsServerError": true,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := &APIError{StatusCode: tt.statusCode}
for checkName, checkFunc := range tt.checkFuncs {
if got := checkFunc(err); got != tt.want[checkName] {
t.Errorf("%s() = %v, want %v", checkName, got, tt.want[checkName])
}
}
})
}
}
func TestNewAPIError(t *testing.T) {
tests := []struct {
name string
resp *http.Response
wantStatus int
wantMsg string
}{
{
name: "nil response",
resp: nil,
wantStatus: 0,
wantMsg: "nil response",
},
{
name: "404 with JSON body",
resp: &http.Response{
StatusCode: 404,
Body: io.NopCloser(strings.NewReader(`{"detail": "Flow not found"}`)),
Header: http.Header{},
},
wantStatus: 404,
wantMsg: "Flow not found",
},
{
name: "500 without body",
resp: &http.Response{
StatusCode: 500,
Body: io.NopCloser(strings.NewReader("")),
Header: http.Header{},
},
wantStatus: 500,
wantMsg: "Internal Server Error",
},
{
name: "with request ID",
resp: &http.Response{
StatusCode: 400,
Body: io.NopCloser(strings.NewReader(`{"detail": "Bad request"}`)),
Header: http.Header{
"X-Request-ID": []string{"req-123"},
},
},
wantStatus: 400,
wantMsg: "Bad request",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := NewAPIError(tt.resp)
apiErr, ok := err.(*APIError)
if !ok {
t.Fatalf("expected *APIError, got %T", err)
}
if apiErr.StatusCode != tt.wantStatus {
t.Errorf("StatusCode = %v, want %v", apiErr.StatusCode, tt.wantStatus)
}
if apiErr.Message != tt.wantMsg {
t.Errorf("Message = %v, want %v", apiErr.Message, tt.wantMsg)
}
})
}
}
func TestNewAPIError_Validation(t *testing.T) {
resp := &http.Response{
StatusCode: 422,
Body: io.NopCloser(strings.NewReader(`{
"detail": [
{
"loc": ["body", "name"],
"msg": "field required",
"type": "value_error.missing"
}
]
}`)),
Header: http.Header{},
}
err := NewAPIError(resp)
valErr, ok := err.(*ValidationError)
if !ok {
t.Fatalf("expected *ValidationError, got %T", err)
}
if valErr.StatusCode != 422 {
t.Errorf("StatusCode = %v, want 422", valErr.StatusCode)
}
if len(valErr.ValidationErrors) != 1 {
t.Fatalf("expected 1 validation error, got %d", len(valErr.ValidationErrors))
}
vd := valErr.ValidationErrors[0]
if len(vd.Loc) != 2 || vd.Loc[0] != "body" || vd.Loc[1] != "name" {
t.Errorf("Loc = %v, want [body name]", vd.Loc)
}
if vd.Msg != "field required" {
t.Errorf("Msg = %v, want 'field required'", vd.Msg)
}
if vd.Type != "value_error.missing" {
t.Errorf("Type = %v, want 'value_error.missing'", vd.Type)
}
}
func TestHelperFunctions(t *testing.T) {
notFoundErr := &APIError{StatusCode: 404}
unauthorizedErr := &APIError{StatusCode: 401}
serverErr := &APIError{StatusCode: 500}
valErr := &ValidationError{APIError: &APIError{StatusCode: 422}}
if !IsNotFound(notFoundErr) {
t.Error("IsNotFound should return true for 404 error")
}
if !IsUnauthorized(unauthorizedErr) {
t.Error("IsUnauthorized should return true for 401 error")
}
if !IsServerError(serverErr) {
t.Error("IsServerError should return true for 500 error")
}
if !IsValidationError(valErr) {
t.Error("IsValidationError should return true for validation error")
}
}

127
pkg/models/artifacts.go Normal file
View File

@@ -0,0 +1,127 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// ArtifactSort represents sort options for artifacts.
type ArtifactSort string
const (
ArtifactSortCreatedDesc ArtifactSort = "CREATED_DESC"
ArtifactSortUpdatedDesc ArtifactSort = "UPDATED_DESC"
)
// ArtifactCollectionSort represents sort options for artifact collections.
type ArtifactCollectionSort string
const (
ArtifactCollectionSortCreatedDesc ArtifactCollectionSort = "CREATED_DESC"
ArtifactCollectionSortUpdatedDesc ArtifactCollectionSort = "UPDATED_DESC"
)
// Artifact represents a Prefect artifact.
type Artifact struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Key *string `json:"key"`
Type *string `json:"type"`
Description *string `json:"description"`
Data interface{} `json:"data,omitempty"`
Metadata map[string]string `json:"metadata_,omitempty"`
FlowRunID *uuid.UUID `json:"flow_run_id"`
TaskRunID *uuid.UUID `json:"task_run_id"`
}
// ArtifactCreate represents the request to create an artifact.
type ArtifactCreate struct {
Key *string `json:"key,omitempty"`
Type *string `json:"type,omitempty"`
Description *string `json:"description,omitempty"`
Data interface{} `json:"data,omitempty"`
Metadata map[string]string `json:"metadata_,omitempty"`
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
TaskRunID *uuid.UUID `json:"task_run_id,omitempty"`
}
// ArtifactUpdate represents the request to update an artifact.
type ArtifactUpdate struct {
Data interface{} `json:"data,omitempty"`
Description *string `json:"description,omitempty"`
Metadata map[string]string `json:"metadata_,omitempty"`
}
// ArtifactFilter represents filter criteria for querying artifacts.
type ArtifactFilter struct {
Key *string `json:"key,omitempty"`
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
TaskRunID *uuid.UUID `json:"task_run_id,omitempty"`
Type *string `json:"type,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// ArtifactCollection represents a collection of artifacts with the same key.
type ArtifactCollection struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Key string `json:"key"`
LatestID uuid.UUID `json:"latest_id"`
Type *string `json:"type"`
Description *string `json:"description"`
Data interface{} `json:"data,omitempty"`
Metadata map[string]string `json:"metadata_,omitempty"`
FlowRunID *uuid.UUID `json:"flow_run_id"`
TaskRunID *uuid.UUID `json:"task_run_id"`
}
// ArtifactCollectionFilter represents filter criteria for querying artifact collections.
type ArtifactCollectionFilter struct {
Key *string `json:"key,omitempty"`
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
TaskRunID *uuid.UUID `json:"task_run_id,omitempty"`
Type *string `json:"type,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (a *Artifact) UnmarshalJSON(data []byte) error {
type Alias Artifact
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(a),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
a.Created = aux.Created.V
a.Updated = aux.Updated.V
return nil
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (ac *ArtifactCollection) UnmarshalJSON(data []byte) error {
type Alias ArtifactCollection
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(ac),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
ac.Created = aux.Created.V
ac.Updated = aux.Updated.V
return nil
}

247
pkg/models/automations.go Normal file
View File

@@ -0,0 +1,247 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// AutomationSort represents sort options for automations.
type AutomationSort string
const (
AutomationSortCreatedDesc AutomationSort = "CREATED_DESC"
AutomationSortUpdatedDesc AutomationSort = "UPDATED_DESC"
AutomationSortNameAsc AutomationSort = "NAME_ASC"
AutomationSortNameDesc AutomationSort = "NAME_DESC"
)
// TriggerPosture represents the posture of an event trigger.
type TriggerPosture string
const (
TriggerPostureReactive TriggerPosture = "Reactive"
TriggerPostureProactive TriggerPosture = "Proactive"
)
// ActionSource represents the source type for automation actions.
type ActionSource string
const (
ActionSourceSelected ActionSource = "selected"
ActionSourceInferred ActionSource = "inferred"
)
// ResourceSpecification represents resource labels to match.
type ResourceSpecification map[string]string
// Automation represents a Prefect automation.
type Automation struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Description string `json:"description"`
Enabled bool `json:"enabled"`
Tags []string `json:"tags,omitempty"`
Trigger json.RawMessage `json:"trigger"`
Actions []json.RawMessage `json:"actions"`
ActionsOnTrigger []json.RawMessage `json:"actions_on_trigger,omitempty"`
ActionsOnResolve []json.RawMessage `json:"actions_on_resolve,omitempty"`
}
// AutomationCreate represents the request to create an automation.
type AutomationCreate struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Enabled *bool `json:"enabled,omitempty"`
Tags []string `json:"tags,omitempty"`
Trigger json.RawMessage `json:"trigger"`
Actions []json.RawMessage `json:"actions"`
ActionsOnTrigger []json.RawMessage `json:"actions_on_trigger,omitempty"`
ActionsOnResolve []json.RawMessage `json:"actions_on_resolve,omitempty"`
}
// AutomationUpdate represents the request to fully update an automation.
type AutomationUpdate struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Enabled *bool `json:"enabled,omitempty"`
Tags []string `json:"tags,omitempty"`
Trigger json.RawMessage `json:"trigger"`
Actions []json.RawMessage `json:"actions"`
ActionsOnTrigger []json.RawMessage `json:"actions_on_trigger,omitempty"`
ActionsOnResolve []json.RawMessage `json:"actions_on_resolve,omitempty"`
}
// AutomationPartialUpdate represents the request to partially update an automation.
type AutomationPartialUpdate struct {
Enabled *bool `json:"enabled,omitempty"`
}
// AutomationFilter represents filter criteria for querying automations.
type AutomationFilter struct {
Name *string `json:"name,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// EventTrigger represents an event-based automation trigger.
type EventTrigger struct {
Type string `json:"type"`
ID uuid.UUID `json:"id"`
Match ResourceSpecification `json:"match,omitempty"`
MatchRelated json.RawMessage `json:"match_related,omitempty"`
After []string `json:"after,omitempty"`
Expect []string `json:"expect,omitempty"`
ForEach []string `json:"for_each,omitempty"`
Posture TriggerPosture `json:"posture"`
Threshold int `json:"threshold,omitempty"`
Within float64 `json:"within,omitempty"`
}
// CompoundTrigger represents a compound automation trigger.
type CompoundTrigger struct {
Type string `json:"type"`
ID uuid.UUID `json:"id"`
Triggers json.RawMessage `json:"triggers"`
Within *float64 `json:"within"`
Require json.RawMessage `json:"require"`
}
// SequenceTrigger represents a sequence automation trigger.
type SequenceTrigger struct {
Type string `json:"type"`
ID uuid.UUID `json:"id"`
Triggers json.RawMessage `json:"triggers"`
Within *float64 `json:"within"`
}
// DoNothing represents a no-op automation action.
type DoNothing struct {
Type string `json:"type"`
}
// RunDeployment represents an action to run a deployment.
type RunDeployment struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
JobVariables map[string]interface{} `json:"job_variables,omitempty"`
}
// PauseDeploymentAction represents an action to pause a deployment.
type PauseDeploymentAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
}
// ResumeDeploymentAction represents an action to resume a deployment.
type ResumeDeploymentAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
}
// CancelFlowRun represents an action to cancel a flow run.
type CancelFlowRun struct {
Type string `json:"type"`
}
// SuspendFlowRun represents an action to suspend a flow run.
type SuspendFlowRun struct {
Type string `json:"type"`
}
// ResumeFlowRunAction represents an action to resume a flow run.
type ResumeFlowRunAction struct {
Type string `json:"type"`
}
// ChangeFlowRunState represents an action to change a flow run's state.
type ChangeFlowRunState struct {
Type string `json:"type"`
Name *string `json:"name,omitempty"`
State StateType `json:"state"`
Message *string `json:"message,omitempty"`
Force bool `json:"force,omitempty"`
}
// PauseWorkQueueAction represents an action to pause a work queue.
type PauseWorkQueueAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
WorkQueueID *uuid.UUID `json:"work_queue_id,omitempty"`
}
// ResumeWorkQueueAction represents an action to resume a work queue.
type ResumeWorkQueueAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
WorkQueueID *uuid.UUID `json:"work_queue_id,omitempty"`
}
// PauseWorkPoolAction represents an action to pause a work pool.
type PauseWorkPoolAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
WorkPoolID *uuid.UUID `json:"work_pool_id,omitempty"`
}
// ResumeWorkPoolAction represents an action to resume a work pool.
type ResumeWorkPoolAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
WorkPoolID *uuid.UUID `json:"work_pool_id,omitempty"`
}
// PauseAutomationAction represents an action to pause an automation.
type PauseAutomationAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
AutomationID *uuid.UUID `json:"automation_id,omitempty"`
}
// ResumeAutomationAction represents an action to resume an automation.
type ResumeAutomationAction struct {
Type string `json:"type"`
Source ActionSource `json:"source,omitempty"`
AutomationID *uuid.UUID `json:"automation_id,omitempty"`
}
// SendNotification represents an action to send a notification.
type SendNotification struct {
Type string `json:"type"`
BlockDocumentID uuid.UUID `json:"block_document_id"`
Subject string `json:"subject,omitempty"`
Body string `json:"body"`
}
// CallWebhook represents an action to call a webhook.
type CallWebhook struct {
Type string `json:"type"`
BlockDocumentID uuid.UUID `json:"block_document_id"`
Payload string `json:"payload,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (a *Automation) UnmarshalJSON(data []byte) error {
type Alias Automation
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(a),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
a.Created = aux.Created.V
a.Updated = aux.Updated.V
return nil
}

176
pkg/models/blocks.go Normal file
View File

@@ -0,0 +1,176 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// BlockDocumentSort represents sort options for block documents.
type BlockDocumentSort string
const (
BlockDocumentSortNameAsc BlockDocumentSort = "NAME_ASC"
BlockDocumentSortNameDesc BlockDocumentSort = "NAME_DESC"
)
// BlockType represents a Prefect block type.
type BlockType struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Slug string `json:"slug"`
LogoURL *string `json:"logo_url"`
DocumentationURL *string `json:"documentation_url"`
Description *string `json:"description"`
CodeExample *string `json:"code_example"`
IsProtected bool `json:"is_protected"`
}
// BlockTypeCreate represents the request to create a block type.
type BlockTypeCreate struct {
Name string `json:"name"`
Slug string `json:"slug"`
LogoURL *string `json:"logo_url,omitempty"`
DocumentationURL *string `json:"documentation_url,omitempty"`
Description *string `json:"description,omitempty"`
CodeExample *string `json:"code_example,omitempty"`
}
// BlockTypeUpdate represents the request to update a block type.
type BlockTypeUpdate struct {
LogoURL *string `json:"logo_url,omitempty"`
DocumentationURL *string `json:"documentation_url,omitempty"`
Description *string `json:"description,omitempty"`
CodeExample *string `json:"code_example,omitempty"`
}
// BlockTypeFilter represents filter criteria for querying block types.
type BlockTypeFilter struct {
Name *string // Filter by name (partial match)
Slugs []string // Filter by slugs
Capabilities []string // Filter by block capabilities
}
// BlockSchema represents a Prefect block schema.
type BlockSchema struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Checksum string `json:"checksum"`
Fields map[string]interface{} `json:"fields"`
BlockTypeID *uuid.UUID `json:"block_type_id"`
BlockType *BlockType `json:"block_type"`
Capabilities []string `json:"capabilities"`
Version string `json:"version"`
}
// BlockSchemaCreate represents the request to create a block schema.
type BlockSchemaCreate struct {
Fields map[string]interface{} `json:"fields,omitempty"`
BlockTypeID uuid.UUID `json:"block_type_id"`
Capabilities []string `json:"capabilities,omitempty"`
Version *string `json:"version,omitempty"`
}
// BlockSchemaFilter represents filter criteria for querying block schemas.
type BlockSchemaFilter struct {
BlockTypeID *uuid.UUID // Filter by block type ID
Capabilities []string // Filter by required capabilities
Version *string // Filter by schema version
}
// BlockDocument represents a Prefect block document.
type BlockDocument struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name *string `json:"name"`
Data map[string]interface{} `json:"data"`
BlockSchemaID uuid.UUID `json:"block_schema_id"`
BlockSchema *BlockSchema `json:"block_schema"`
BlockTypeID uuid.UUID `json:"block_type_id"`
BlockTypeName *string `json:"block_type_name"`
BlockType *BlockType `json:"block_type"`
BlockDocumentReferences map[string]interface{} `json:"block_document_references"`
IsAnonymous bool `json:"is_anonymous"`
}
// BlockDocumentCreate represents the request to create a block document.
type BlockDocumentCreate struct {
Name *string `json:"name,omitempty"`
Data map[string]interface{} `json:"data,omitempty"`
BlockSchemaID uuid.UUID `json:"block_schema_id"`
BlockTypeID uuid.UUID `json:"block_type_id"`
IsAnonymous bool `json:"is_anonymous,omitempty"`
}
// BlockDocumentUpdate represents the request to update a block document.
type BlockDocumentUpdate struct {
BlockSchemaID *uuid.UUID `json:"block_schema_id,omitempty"`
Data map[string]interface{} `json:"data"`
MergeExisting *bool `json:"merge_existing_data,omitempty"`
}
// BlockDocumentFilter represents filter criteria for querying block documents.
type BlockDocumentFilter struct {
BlockTypeID *uuid.UUID // Filter by block type ID
Name *string // Filter by document name
IsAnonymous *bool // Filter by anonymity
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (bt *BlockType) UnmarshalJSON(data []byte) error {
type Alias BlockType
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(bt),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
bt.Created = aux.Created.V
bt.Updated = aux.Updated.V
return nil
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (bs *BlockSchema) UnmarshalJSON(data []byte) error {
type Alias BlockSchema
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(bs),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
bs.Created = aux.Created.V
bs.Updated = aux.Updated.V
return nil
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (bd *BlockDocument) UnmarshalJSON(data []byte) error {
type Alias BlockDocument
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(bd),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
bd.Created = aux.Created.V
bd.Updated = aux.Updated.V
return nil
}

View File

@@ -0,0 +1,137 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// ConcurrencyLimit represents a Prefect v1 concurrency limit.
type ConcurrencyLimit struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Tag string `json:"tag"`
ConcurrencyLimit int `json:"concurrency_limit"`
ActiveSlots []string `json:"active_slots,omitempty"`
}
// ConcurrencyLimitCreate represents the request to create a v1 concurrency limit.
type ConcurrencyLimitCreate struct {
Tag string `json:"tag"`
ConcurrencyLimit int `json:"concurrency_limit"`
}
// ConcurrencyLimitV2 represents a Prefect v2 concurrency limit.
type ConcurrencyLimitV2 struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Active bool `json:"active"`
Name string `json:"name"`
Limit int `json:"limit"`
ActiveSlots int `json:"active_slots"`
DeniedSlots int `json:"denied_slots"`
SlotDecayPerSecond float64 `json:"slot_decay_per_second"`
AvgSlotOccupancySeconds float64 `json:"avg_slot_occupancy_seconds"`
}
// ConcurrencyLimitV2Create represents the request to create a v2 concurrency limit.
type ConcurrencyLimitV2Create struct {
Active *bool `json:"active,omitempty"`
Name string `json:"name"`
Limit int `json:"limit"`
ActiveSlots *int `json:"active_slots,omitempty"`
DeniedSlots *int `json:"denied_slots,omitempty"`
SlotDecayPerSecond *float64 `json:"slot_decay_per_second,omitempty"`
}
// ConcurrencyLimitV2Update represents the request to update a v2 concurrency limit.
type ConcurrencyLimitV2Update struct {
Active *bool `json:"active,omitempty"`
Name *string `json:"name,omitempty"`
Limit *int `json:"limit,omitempty"`
ActiveSlots *int `json:"active_slots,omitempty"`
DeniedSlots *int `json:"denied_slots,omitempty"`
SlotDecayPerSecond *float64 `json:"slot_decay_per_second,omitempty"`
}
// GlobalConcurrencyLimitResponse represents the response for global concurrency limits.
type GlobalConcurrencyLimitResponse struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Active bool `json:"active"`
Name string `json:"name"`
Limit int `json:"limit"`
ActiveSlots int `json:"active_slots"`
SlotDecayPerSecond float64 `json:"slot_decay_per_second"`
}
// MinimalConcurrencyLimitResponse represents a minimal concurrency limit response.
type MinimalConcurrencyLimitResponse struct {
ID uuid.UUID `json:"id"`
Name string `json:"name"`
Limit int `json:"limit"`
}
// ConcurrencyLimitWithLeaseResponse represents a concurrency limit response with lease info.
type ConcurrencyLimitWithLeaseResponse struct {
LeaseID uuid.UUID `json:"lease_id"`
Limits []MinimalConcurrencyLimitResponse `json:"limits"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (cl *ConcurrencyLimit) UnmarshalJSON(data []byte) error {
type Alias ConcurrencyLimit
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(cl),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
cl.Created = aux.Created.V
cl.Updated = aux.Updated.V
return nil
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (cl *ConcurrencyLimitV2) UnmarshalJSON(data []byte) error {
type Alias ConcurrencyLimitV2
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(cl),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
cl.Created = aux.Created.V
cl.Updated = aux.Updated.V
return nil
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (g *GlobalConcurrencyLimitResponse) UnmarshalJSON(data []byte) error {
type Alias GlobalConcurrencyLimitResponse
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(g),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
g.Created = aux.Created.V
g.Updated = aux.Updated.V
return nil
}

102
pkg/models/deployments.go Normal file
View File

@@ -0,0 +1,102 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// DeploymentStatus represents the status of a deployment.
type DeploymentStatus string
const (
DeploymentStatusReady DeploymentStatus = "READY"
DeploymentStatusNotReady DeploymentStatus = "NOT_READY"
)
// CollisionStrategy represents the strategy for handling concurrent flow runs.
type CollisionStrategy string
const (
CollisionStrategyEnqueue CollisionStrategy = "ENQUEUE"
CollisionStrategyCancelNew CollisionStrategy = "CANCEL_NEW"
)
// ConcurrencyOptions represents concurrency configuration for a deployment.
type ConcurrencyOptions struct {
CollisionStrategy CollisionStrategy `json:"collision_strategy"`
GracePeriodSeconds *int `json:"grace_period_seconds,omitempty"`
}
// Deployment represents a Prefect deployment.
type Deployment struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
FlowID uuid.UUID `json:"flow_id"`
Version *string `json:"version"`
Description *string `json:"description"`
Paused bool `json:"paused"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
WorkQueueName *string `json:"work_queue_name"`
WorkPoolName *string `json:"work_pool_name"`
Path *string `json:"path"`
Entrypoint *string `json:"entrypoint"`
Status *DeploymentStatus `json:"status"`
EnforceParameterSchema bool `json:"enforce_parameter_schema"`
}
// DeploymentCreate represents the request to create a deployment.
type DeploymentCreate struct {
Name string `json:"name"`
FlowID uuid.UUID `json:"flow_id"`
Paused bool `json:"paused,omitempty"`
Description *string `json:"description,omitempty"`
Version *string `json:"version,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
WorkPoolName *string `json:"work_pool_name,omitempty"`
WorkQueueName *string `json:"work_queue_name,omitempty"`
Path *string `json:"path,omitempty"`
Entrypoint *string `json:"entrypoint,omitempty"`
EnforceParameterSchema bool `json:"enforce_parameter_schema,omitempty"`
}
// DeploymentUpdate represents the request to update a deployment.
type DeploymentUpdate struct {
Paused *bool `json:"paused,omitempty"`
Description *string `json:"description,omitempty"`
}
// DeploymentFilter represents filter criteria for querying deployments.
type DeploymentFilter struct {
FlowID *uuid.UUID `json:"flow_id,omitempty"`
Name *string `json:"name,omitempty"`
Tags []string `json:"tags,omitempty"`
Paused *bool `json:"paused,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (d *Deployment) UnmarshalJSON(data []byte) error {
type Alias Deployment
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(d),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
d.Created = aux.Created.V
d.Updated = aux.Updated.V
return nil
}

58
pkg/models/events.go Normal file
View File

@@ -0,0 +1,58 @@
package models
import (
"time"
"github.com/google/uuid"
)
// EventOrder represents sort order for events.
type EventOrder string
const (
EventOrderAsc EventOrder = "ASC"
EventOrderDesc EventOrder = "DESC"
)
// Resource represents an observable business object.
type Resource map[string]string
// RelatedResource represents a resource with a specific role in an event.
type RelatedResource map[string]string
// Event represents a Prefect event.
type Event struct {
Occurred time.Time `json:"occurred"`
Event string `json:"event"`
Resource Resource `json:"resource"`
Related []RelatedResource `json:"related,omitempty"`
ID uuid.UUID `json:"id"`
Payload map[string]interface{} `json:"payload,omitempty"`
Received *time.Time `json:"received,omitempty"`
}
// EventFilter represents filter criteria for querying events.
type EventFilter struct {
Occurred interface{} `json:"occurred,omitempty"`
Event interface{} `json:"event,omitempty"`
Resource interface{} `json:"resource,omitempty"`
Related interface{} `json:"related,omitempty"`
ID interface{} `json:"id,omitempty"`
Order EventOrder `json:"order,omitempty"`
}
// EventPage represents a page of events.
type EventPage struct {
Events []Event `json:"events"`
Total int `json:"total"`
NextPage *string `json:"next_page"`
}
// EventCount represents a count of events for a given filter value.
type EventCount struct {
Value string `json:"value"`
Label string `json:"label"`
Count int `json:"count"`
StartTime time.Time `json:"start_time"`
EndTime time.Time `json:"end_time"`
}

110
pkg/models/flow_runs.go Normal file
View File

@@ -0,0 +1,110 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// FlowRun represents a Prefect flow run.
type FlowRun struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
FlowID uuid.UUID `json:"flow_id"`
Name string `json:"name"`
StateID *uuid.UUID `json:"state_id"`
DeploymentID *uuid.UUID `json:"deployment_id"`
WorkQueueID *uuid.UUID `json:"work_queue_id"`
WorkQueueName *string `json:"work_queue_name"`
FlowVersion *string `json:"flow_version"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
IdempotencyKey *string `json:"idempotency_key"`
Context map[string]interface{} `json:"context,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
ParentTaskRunID *uuid.UUID `json:"parent_task_run_id"`
StateType *StateType `json:"state_type"`
StateName *string `json:"state_name"`
RunCount int `json:"run_count"`
ExpectedStartTime *time.Time `json:"expected_start_time"`
NextScheduledStartTime *time.Time `json:"next_scheduled_start_time"`
StartTime *time.Time `json:"start_time"`
EndTime *time.Time `json:"end_time"`
TotalRunTime float64 `json:"total_run_time"`
State *State `json:"state"`
}
// FlowRunCreate represents the request to create a flow run.
type FlowRunCreate struct {
FlowID uuid.UUID `json:"flow_id"`
Name string `json:"name,omitempty"`
State *StateCreate `json:"state,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Context map[string]interface{} `json:"context,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
IdempotencyKey *string `json:"idempotency_key,omitempty"`
WorkPoolName *string `json:"work_pool_name,omitempty"`
WorkQueueName *string `json:"work_queue_name,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
}
// FlowRunUpdate represents the request to update a flow run.
type FlowRunUpdate struct {
Name *string `json:"name,omitempty"`
}
// FlowRunFilter represents filter criteria for querying flow runs.
type FlowRunFilter struct {
FlowID *uuid.UUID `json:"flow_id,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
StateType *StateType `json:"state_type,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// FlowRunPolicy represents a flow run's retry policy.
type FlowRunPolicy struct {
MaxRetries int `json:"max_retries,omitempty"`
RetryDelaySeconds float64 `json:"retry_delay_seconds,omitempty"`
Retries *int `json:"retries,omitempty"`
RetryDelay *int `json:"retry_delay,omitempty"`
PauseKeys []string `json:"pause_keys,omitempty"`
Resuming *bool `json:"resuming,omitempty"`
}
// CreatedBy represents information about who created an object.
type CreatedBy struct {
ID *uuid.UUID `json:"id,omitempty"`
Type *string `json:"type,omitempty"`
DisplayValue *string `json:"display_value,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (fr *FlowRun) UnmarshalJSON(data []byte) error {
type Alias FlowRun
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
ExpectedStartTime optTime `json:"expected_start_time"`
NextScheduledStartTime optTime `json:"next_scheduled_start_time"`
StartTime optTime `json:"start_time"`
EndTime optTime `json:"end_time"`
*Alias
}{
Alias: (*Alias)(fr),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
fr.Created = aux.Created.V
fr.Updated = aux.Updated.V
fr.ExpectedStartTime = aux.ExpectedStartTime.V
fr.NextScheduledStartTime = aux.NextScheduledStartTime.V
fr.StartTime = aux.StartTime.V
fr.EndTime = aux.EndTime.V
return nil
}

57
pkg/models/flows.go Normal file
View File

@@ -0,0 +1,57 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// Flow represents a Prefect flow.
type Flow struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
}
// FlowCreate represents the request to create a flow.
type FlowCreate struct {
Name string `json:"name"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
}
// FlowUpdate represents the request to update a flow.
type FlowUpdate struct {
Tags *[]string `json:"tags,omitempty"`
Labels *map[string]interface{} `json:"labels,omitempty"`
}
// FlowFilter represents filter criteria for querying flows.
type FlowFilter struct {
Name *string `json:"name,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (f *Flow) UnmarshalJSON(data []byte) error {
type Alias Flow
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(f),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
f.Created = aux.Created.V
f.Updated = aux.Updated.V
return nil
}

49
pkg/models/logs.go Normal file
View File

@@ -0,0 +1,49 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// Log represents a Prefect log entry.
type Log struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Level int `json:"level"`
Message string `json:"message"`
Timestamp time.Time `json:"timestamp"`
FlowRunID *uuid.UUID `json:"flow_run_id"`
TaskRunID *uuid.UUID `json:"task_run_id"`
}
// LogCreate represents the request to create a log.
type LogCreate struct {
Name string `json:"name"`
Level int `json:"level"`
Message string `json:"message"`
Timestamp time.Time `json:"timestamp"`
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
TaskRunID *uuid.UUID `json:"task_run_id,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (l *Log) UnmarshalJSON(data []byte) error {
type Alias Log
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(l),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
l.Created = aux.Created.V
l.Updated = aux.Updated.V
return nil
}

View File

@@ -1,326 +0,0 @@
// Package models contains the data structures for the Prefect API.
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// StateType represents the type of a flow or task run state.
type StateType string
const (
StateTypePending StateType = "PENDING"
StateTypeRunning StateType = "RUNNING"
StateTypeCompleted StateType = "COMPLETED"
StateTypeFailed StateType = "FAILED"
StateTypeCancelled StateType = "CANCELLED"
StateTypeCrashed StateType = "CRASHED"
StateTypePaused StateType = "PAUSED"
StateTypeScheduled StateType = "SCHEDULED"
StateTypeCancelling StateType = "CANCELLING"
)
// DeploymentStatus represents the status of a deployment.
type DeploymentStatus string
const (
DeploymentStatusReady DeploymentStatus = "READY"
DeploymentStatusNotReady DeploymentStatus = "NOT_READY"
)
// CollisionStrategy represents the strategy for handling concurrent flow runs.
type CollisionStrategy string
const (
CollisionStrategyEnqueue CollisionStrategy = "ENQUEUE"
CollisionStrategyCancelNew CollisionStrategy = "CANCEL_NEW"
)
// Flow represents a Prefect flow.
type Flow struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
}
// FlowCreate represents the request to create a flow.
type FlowCreate struct {
Name string `json:"name"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
}
// FlowUpdate represents the request to update a flow.
type FlowUpdate struct {
Tags *[]string `json:"tags,omitempty"`
Labels *map[string]interface{} `json:"labels,omitempty"`
}
// FlowFilter represents filter criteria for querying flows.
type FlowFilter struct {
Name *string `json:"name,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// FlowRun represents a Prefect flow run.
type FlowRun struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
FlowID uuid.UUID `json:"flow_id"`
Name string `json:"name"`
StateID *uuid.UUID `json:"state_id"`
DeploymentID *uuid.UUID `json:"deployment_id"`
WorkQueueID *uuid.UUID `json:"work_queue_id"`
WorkQueueName *string `json:"work_queue_name"`
FlowVersion *string `json:"flow_version"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
IdempotencyKey *string `json:"idempotency_key"`
Context map[string]interface{} `json:"context,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
ParentTaskRunID *uuid.UUID `json:"parent_task_run_id"`
StateType *StateType `json:"state_type"`
StateName *string `json:"state_name"`
RunCount int `json:"run_count"`
ExpectedStartTime *time.Time `json:"expected_start_time"`
NextScheduledStartTime *time.Time `json:"next_scheduled_start_time"`
StartTime *time.Time `json:"start_time"`
EndTime *time.Time `json:"end_time"`
TotalRunTime float64 `json:"total_run_time"`
State *State `json:"state"`
}
// FlowRunCreate represents the request to create a flow run.
type FlowRunCreate struct {
FlowID uuid.UUID `json:"flow_id"`
Name string `json:"name,omitempty"`
State *StateCreate `json:"state,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Context map[string]interface{} `json:"context,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
IdempotencyKey *string `json:"idempotency_key,omitempty"`
WorkPoolName *string `json:"work_pool_name,omitempty"`
WorkQueueName *string `json:"work_queue_name,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
}
// FlowRunUpdate represents the request to update a flow run.
type FlowRunUpdate struct {
Name *string `json:"name,omitempty"`
}
// FlowRunFilter represents filter criteria for querying flow runs.
type FlowRunFilter struct {
FlowID *uuid.UUID `json:"flow_id,omitempty"`
DeploymentID *uuid.UUID `json:"deployment_id,omitempty"`
StateType *StateType `json:"state_type,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// State represents the state of a flow or task run.
type State struct {
ID uuid.UUID `json:"id"`
Type StateType `json:"type"`
Name *string `json:"name"`
Timestamp time.Time `json:"timestamp"`
Message *string `json:"message"`
Data interface{} `json:"data,omitempty"`
StateDetails map[string]interface{} `json:"state_details,omitempty"`
}
// StateCreate represents the request to create a state.
type StateCreate struct {
Type StateType `json:"type"`
Name *string `json:"name,omitempty"`
Message *string `json:"message,omitempty"`
Data interface{} `json:"data,omitempty"`
StateDetails map[string]interface{} `json:"state_details,omitempty"`
}
// Deployment represents a Prefect deployment.
type Deployment struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
FlowID uuid.UUID `json:"flow_id"`
Version *string `json:"version"`
Description *string `json:"description"`
Paused bool `json:"paused"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
WorkQueueName *string `json:"work_queue_name"`
WorkPoolName *string `json:"work_pool_name"`
Path *string `json:"path"`
Entrypoint *string `json:"entrypoint"`
Status *DeploymentStatus `json:"status"`
EnforceParameterSchema bool `json:"enforce_parameter_schema"`
}
// DeploymentCreate represents the request to create a deployment.
type DeploymentCreate struct {
Name string `json:"name"`
FlowID uuid.UUID `json:"flow_id"`
Paused bool `json:"paused,omitempty"`
Description *string `json:"description,omitempty"`
Version *string `json:"version,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Tags []string `json:"tags,omitempty"`
Labels map[string]interface{} `json:"labels,omitempty"`
WorkPoolName *string `json:"work_pool_name,omitempty"`
WorkQueueName *string `json:"work_queue_name,omitempty"`
Path *string `json:"path,omitempty"`
Entrypoint *string `json:"entrypoint,omitempty"`
EnforceParameterSchema bool `json:"enforce_parameter_schema,omitempty"`
}
// DeploymentUpdate represents the request to update a deployment.
type DeploymentUpdate struct {
Paused *bool `json:"paused,omitempty"`
Description *string `json:"description,omitempty"`
}
// TaskRun represents a Prefect task run.
type TaskRun struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
FlowRunID uuid.UUID `json:"flow_run_id"`
TaskKey string `json:"task_key"`
DynamicKey string `json:"dynamic_key"`
CacheKey *string `json:"cache_key"`
StartTime *time.Time `json:"start_time"`
EndTime *time.Time `json:"end_time"`
TotalRunTime float64 `json:"total_run_time"`
Status *StateType `json:"status"`
StateID *uuid.UUID `json:"state_id"`
Tags []string `json:"tags,omitempty"`
State *State `json:"state"`
}
// TaskRunCreate represents the request to create a task run.
type TaskRunCreate struct {
Name string `json:"name"`
FlowRunID uuid.UUID `json:"flow_run_id"`
TaskKey string `json:"task_key"`
DynamicKey string `json:"dynamic_key"`
CacheKey *string `json:"cache_key,omitempty"`
State *StateCreate `json:"state,omitempty"`
Tags []string `json:"tags,omitempty"`
}
// WorkPool represents a Prefect work pool.
type WorkPool struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Description *string `json:"description"`
Type string `json:"type"`
IsPaused bool `json:"is_paused"`
}
// WorkQueue represents a Prefect work queue.
type WorkQueue struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Description *string `json:"description"`
IsPaused bool `json:"is_paused"`
Priority int `json:"priority"`
}
// Variable represents a Prefect variable.
type Variable struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Value string `json:"value"`
Tags []string `json:"tags,omitempty"`
}
// VariableCreate represents the request to create a variable.
type VariableCreate struct {
Name string `json:"name"`
Value string `json:"value"`
Tags []string `json:"tags,omitempty"`
}
// VariableUpdate represents the request to update a variable.
type VariableUpdate struct {
Value *string `json:"value,omitempty"`
Tags []string `json:"tags,omitempty"`
}
// Log represents a Prefect log entry.
type Log struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Level int `json:"level"`
Message string `json:"message"`
Timestamp time.Time `json:"timestamp"`
FlowRunID *uuid.UUID `json:"flow_run_id"`
TaskRunID *uuid.UUID `json:"task_run_id"`
}
// LogCreate represents the request to create a log.
type LogCreate struct {
Name string `json:"name"`
Level int `json:"level"`
Message string `json:"message"`
Timestamp time.Time `json:"timestamp"`
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
TaskRunID *uuid.UUID `json:"task_run_id,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (f *Flow) UnmarshalJSON(data []byte) error {
type Alias Flow
aux := &struct {
Created string `json:"created"`
Updated string `json:"updated"`
*Alias
}{
Alias: (*Alias)(f),
}
if err := json.Unmarshal(data, &aux); err != nil {
return err
}
if aux.Created != "" {
t, err := time.Parse(time.RFC3339, aux.Created)
if err != nil {
return err
}
f.Created = &t
}
if aux.Updated != "" {
t, err := time.Parse(time.RFC3339, aux.Updated)
if err != nil {
return err
}
f.Updated = &t
}
return nil
}

View File

@@ -0,0 +1,50 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// SavedSearch represents a Prefect saved search.
type SavedSearch struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Filters []SavedSearchFilter `json:"filters"`
}
// SavedSearchCreate represents the request to create a saved search.
type SavedSearchCreate struct {
Name string `json:"name"`
Filters []SavedSearchFilter `json:"filters,omitempty"`
}
// SavedSearchFilter represents a filter definition within a saved search.
type SavedSearchFilter struct {
Object string `json:"object"`
Property string `json:"property"`
Type string `json:"type"`
Operation string `json:"operation"`
Value interface{} `json:"value"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (ss *SavedSearch) UnmarshalJSON(data []byte) error {
type Alias SavedSearch
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(ss),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
ss.Created = aux.Created.V
ss.Updated = aux.Updated.V
return nil
}

77
pkg/models/schedules.go Normal file
View File

@@ -0,0 +1,77 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// DeploymentSchedule represents a schedule for a deployment.
type DeploymentSchedule struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
DeploymentID *uuid.UUID `json:"deployment_id"`
Schedule json.RawMessage `json:"schedule"`
Active bool `json:"active"`
MaxScheduledRuns *int `json:"max_scheduled_runs,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Slug *string `json:"slug,omitempty"`
}
// DeploymentScheduleCreate represents the request to create a deployment schedule.
type DeploymentScheduleCreate struct {
Active *bool `json:"active,omitempty"`
Schedule json.RawMessage `json:"schedule"`
MaxScheduledRuns *int `json:"max_scheduled_runs,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Slug *string `json:"slug,omitempty"`
}
// DeploymentScheduleUpdate represents the request to update a deployment schedule.
type DeploymentScheduleUpdate struct {
Active *bool `json:"active,omitempty"`
Schedule json.RawMessage `json:"schedule,omitempty"`
MaxScheduledRuns *int `json:"max_scheduled_runs,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
Slug *string `json:"slug,omitempty"`
}
// IntervalSchedule represents an interval-based schedule.
type IntervalSchedule struct {
Interval float64 `json:"interval"`
AnchorDate *string `json:"anchor_date,omitempty"`
Timezone *string `json:"timezone,omitempty"`
}
// CronSchedule represents a cron-based schedule.
type CronSchedule struct {
Cron string `json:"cron"`
Timezone *string `json:"timezone,omitempty"`
DayOr *bool `json:"day_or,omitempty"`
}
// RRuleSchedule represents an RRule-based schedule.
type RRuleSchedule struct {
RRule string `json:"rrule"`
Timezone *string `json:"timezone,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (ds *DeploymentSchedule) UnmarshalJSON(data []byte) error {
type Alias DeploymentSchedule
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(ds),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
ds.Created = aux.Created.V
ds.Updated = aux.Updated.V
return nil
}

55
pkg/models/states.go Normal file
View File

@@ -0,0 +1,55 @@
package models
import (
"time"
"github.com/google/uuid"
)
// StateType represents the type of a flow or task run state.
type StateType string
const (
StateTypePending StateType = "PENDING"
StateTypeRunning StateType = "RUNNING"
StateTypeCompleted StateType = "COMPLETED"
StateTypeFailed StateType = "FAILED"
StateTypeCancelled StateType = "CANCELLED"
StateTypeCrashed StateType = "CRASHED"
StateTypePaused StateType = "PAUSED"
StateTypeScheduled StateType = "SCHEDULED"
StateTypeCancelling StateType = "CANCELLING"
)
// State represents the state of a flow or task run.
type State struct {
ID uuid.UUID `json:"id"`
Type StateType `json:"type"`
Name *string `json:"name"`
Timestamp time.Time `json:"timestamp"`
Message *string `json:"message"`
Data interface{} `json:"data,omitempty"`
StateDetails map[string]interface{} `json:"state_details,omitempty"`
}
// StateCreate represents the request to create a state.
type StateCreate struct {
Type StateType `json:"type"`
Name *string `json:"name,omitempty"`
Message *string `json:"message,omitempty"`
Data interface{} `json:"data,omitempty"`
StateDetails map[string]interface{} `json:"state_details,omitempty"`
}
// StateDetails represents detailed information about a state.
type StateDetails struct {
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
TaskRunID *uuid.UUID `json:"task_run_id,omitempty"`
ChildFlowRunID *uuid.UUID `json:"child_flow_run_id,omitempty"`
ScheduledTime *time.Time `json:"scheduled_time,omitempty"`
CacheKey *string `json:"cache_key,omitempty"`
CacheExpiration *time.Time `json:"cache_expiration,omitempty"`
Deferred *bool `json:"deferred,omitempty"`
UntrackableResult bool `json:"untrackable_result,omitempty"`
PauseTimeout *time.Time `json:"pause_timeout,omitempty"`
}

83
pkg/models/task_runs.go Normal file
View File

@@ -0,0 +1,83 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// TaskRun represents a Prefect task run.
type TaskRun struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
FlowRunID uuid.UUID `json:"flow_run_id"`
TaskKey string `json:"task_key"`
DynamicKey string `json:"dynamic_key"`
CacheKey *string `json:"cache_key"`
StartTime *time.Time `json:"start_time"`
EndTime *time.Time `json:"end_time"`
TotalRunTime float64 `json:"total_run_time"`
Status *StateType `json:"status"`
StateID *uuid.UUID `json:"state_id"`
Tags []string `json:"tags,omitempty"`
State *State `json:"state"`
}
// TaskRunCreate represents the request to create a task run.
type TaskRunCreate struct {
Name string `json:"name"`
FlowRunID uuid.UUID `json:"flow_run_id"`
TaskKey string `json:"task_key"`
DynamicKey string `json:"dynamic_key"`
CacheKey *string `json:"cache_key,omitempty"`
State *StateCreate `json:"state,omitempty"`
Tags []string `json:"tags,omitempty"`
}
// TaskRunUpdate represents the request to update a task run.
type TaskRunUpdate struct {
Name *string `json:"name,omitempty"`
}
// TaskRunFilter represents filter criteria for querying task runs.
type TaskRunFilter struct {
FlowRunID *uuid.UUID `json:"flow_run_id,omitempty"`
StateType *StateType `json:"state_type,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// TaskRunPolicy represents a task run's retry policy.
type TaskRunPolicy struct {
MaxRetries int `json:"max_retries,omitempty"`
RetryDelaySeconds float64 `json:"retry_delay_seconds,omitempty"`
Retries *int `json:"retries,omitempty"`
RetryDelay interface{} `json:"retry_delay,omitempty"`
RetryJitterFactor *float64 `json:"retry_jitter_factor,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (tr *TaskRun) UnmarshalJSON(data []byte) error {
type Alias TaskRun
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
StartTime optTime `json:"start_time"`
EndTime optTime `json:"end_time"`
*Alias
}{
Alias: (*Alias)(tr),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
tr.Created = aux.Created.V
tr.Updated = aux.Updated.V
tr.StartTime = aux.StartTime.V
tr.EndTime = aux.EndTime.V
return nil
}

View File

@@ -0,0 +1,15 @@
package models
import "time"
// TaskWorkerFilter represents filter criteria for querying task workers.
type TaskWorkerFilter struct {
TaskKeys []string `json:"task_keys"`
}
// TaskWorkerResponse represents a task worker response.
type TaskWorkerResponse struct {
Identifier string `json:"identifier"`
TaskKeys []string `json:"task_keys"`
Timestamp time.Time `json:"timestamp"`
}

45
pkg/models/time.go Normal file
View File

@@ -0,0 +1,45 @@
package models
import (
"encoding/json"
"time"
)
// parseOptionalTime parses a time string, returning nil for empty strings.
// Supports RFC3339Nano and RFC3339 formats used by the Prefect API.
func parseOptionalTime(s string) (*time.Time, error) {
if s == "" {
return nil, nil
}
t, err := time.Parse(time.RFC3339Nano, s)
if err != nil {
t, err = time.Parse(time.RFC3339, s)
if err != nil {
return nil, err
}
}
return &t, nil
}
// optTime handles JSON unmarshaling of nullable time fields that may be
// represented as empty strings or JSON null by the Prefect API.
type optTime struct {
V *time.Time
}
func (o *optTime) UnmarshalJSON(data []byte) error {
if string(data) == "null" {
o.V = nil
return nil
}
var s string
if err := json.Unmarshal(data, &s); err != nil {
return err
}
v, err := parseOptionalTime(s)
if err != nil {
return err
}
o.V = v
return nil
}

57
pkg/models/variables.go Normal file
View File

@@ -0,0 +1,57 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// Variable represents a Prefect variable.
type Variable struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Value string `json:"value"`
Tags []string `json:"tags,omitempty"`
}
// VariableCreate represents the request to create a variable.
type VariableCreate struct {
Name string `json:"name"`
Value string `json:"value"`
Tags []string `json:"tags,omitempty"`
}
// VariableUpdate represents the request to update a variable.
type VariableUpdate struct {
Value *string `json:"value,omitempty"`
Tags []string `json:"tags,omitempty"`
}
// VariableFilter represents filter criteria for querying variables.
type VariableFilter struct {
Name *string `json:"name,omitempty"`
Tags []string `json:"tags,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (v *Variable) UnmarshalJSON(data []byte) error {
type Alias Variable
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(v),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
v.Created = aux.Created.V
v.Updated = aux.Updated.V
return nil
}

72
pkg/models/work_pools.go Normal file
View File

@@ -0,0 +1,72 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// WorkPoolStatus represents the status of a work pool.
type WorkPoolStatus string
const (
WorkPoolStatusReady WorkPoolStatus = "READY"
WorkPoolStatusNotReady WorkPoolStatus = "NOT_READY"
WorkPoolStatusPaused WorkPoolStatus = "PAUSED"
)
// WorkPool represents a Prefect work pool.
type WorkPool struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Description *string `json:"description"`
Type string `json:"type"`
IsPaused bool `json:"is_paused"`
}
// WorkPoolCreate represents the request to create a work pool.
type WorkPoolCreate struct {
Name string `json:"name"`
Description *string `json:"description,omitempty"`
Type string `json:"type"`
BaseJobTemplate map[string]interface{} `json:"base_job_template,omitempty"`
IsPaused bool `json:"is_paused,omitempty"`
ConcurrencyLimit *int `json:"concurrency_limit,omitempty"`
}
// WorkPoolUpdate represents the request to update a work pool.
type WorkPoolUpdate struct {
Description *string `json:"description,omitempty"`
IsPaused *bool `json:"is_paused,omitempty"`
BaseJobTemplate map[string]interface{} `json:"base_job_template,omitempty"`
ConcurrencyLimit *int `json:"concurrency_limit,omitempty"`
}
// WorkPoolFilter represents filter criteria for querying work pools.
type WorkPoolFilter struct {
Name *string `json:"name,omitempty"`
Type *string `json:"type,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (wp *WorkPool) UnmarshalJSON(data []byte) error {
type Alias WorkPool
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(wp),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
wp.Created = aux.Created.V
wp.Updated = aux.Updated.V
return nil
}

62
pkg/models/work_queues.go Normal file
View File

@@ -0,0 +1,62 @@
package models
import (
"encoding/json"
"time"
"github.com/google/uuid"
)
// WorkQueue represents a Prefect work queue.
type WorkQueue struct {
ID uuid.UUID `json:"id"`
Created *time.Time `json:"created"`
Updated *time.Time `json:"updated"`
Name string `json:"name"`
Description *string `json:"description"`
IsPaused bool `json:"is_paused"`
Priority int `json:"priority"`
}
// WorkQueueCreate represents the request to create a work queue.
type WorkQueueCreate struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
IsPaused bool `json:"is_paused,omitempty"`
ConcurrencyLimit *int `json:"concurrency_limit,omitempty"`
Priority *int `json:"priority,omitempty"`
}
// WorkQueueUpdate represents the request to update a work queue.
type WorkQueueUpdate struct {
Name *string `json:"name,omitempty"`
Description *string `json:"description,omitempty"`
IsPaused *bool `json:"is_paused,omitempty"`
ConcurrencyLimit *int `json:"concurrency_limit,omitempty"`
Priority *int `json:"priority,omitempty"`
}
// WorkQueueFilter represents filter criteria for querying work queues.
type WorkQueueFilter struct {
Name *string `json:"name,omitempty"`
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
}
// UnmarshalJSON implements custom JSON unmarshaling for time fields.
func (wq *WorkQueue) UnmarshalJSON(data []byte) error {
type Alias WorkQueue
aux := &struct {
Created optTime `json:"created"`
Updated optTime `json:"updated"`
*Alias
}{
Alias: (*Alias)(wq),
}
if err := json.Unmarshal(data, aux); err != nil {
return err
}
wq.Created = aux.Created.V
wq.Updated = aux.Updated.V
return nil
}

View File

@@ -1,339 +0,0 @@
package pagination
import (
"context"
"errors"
"testing"
)
func TestNewIterator(t *testing.T) {
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
return nil, nil
}
tests := []struct {
name string
limit int
wantLimit int
}{
{"default limit", 0, 100},
{"custom limit", 50, 50},
{"negative limit", -10, 100},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
iter := NewIterator(fetchFunc, tt.limit)
if iter.limit != tt.wantLimit {
t.Errorf("limit = %d, want %d", iter.limit, tt.wantLimit)
}
})
}
}
func TestIterator_Next(t *testing.T) {
items := []string{"item1", "item2", "item3", "item4", "item5"}
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
if offset >= len(items) {
return &PaginatedResponse[string]{
Results: []string{},
Count: len(items),
Limit: limit,
Offset: offset,
HasMore: false,
}, nil
}
end := offset + limit
if end > len(items) {
end = len(items)
}
return &PaginatedResponse[string]{
Results: items[offset:end],
Count: len(items),
Limit: limit,
Offset: offset,
HasMore: end < len(items),
}, nil
}
ctx := context.Background()
iter := NewIterator(fetchFunc, 2)
var results []string
for iter.Next(ctx) {
if item := iter.Value(); item != nil {
results = append(results, *item)
}
}
if err := iter.Err(); err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(results) != len(items) {
t.Errorf("got %d items, want %d", len(results), len(items))
}
for i, want := range items {
if i >= len(results) || results[i] != want {
t.Errorf("item[%d] = %v, want %v", i, results[i], want)
}
}
}
func TestIterator_Error(t *testing.T) {
expectedErr := errors.New("fetch error")
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
return nil, expectedErr
}
ctx := context.Background()
iter := NewIterator(fetchFunc, 10)
if iter.Next(ctx) {
t.Error("Next should return false on error")
}
if err := iter.Err(); err != expectedErr {
t.Errorf("Err() = %v, want %v", err, expectedErr)
}
}
func TestIterator_EmptyResults(t *testing.T) {
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
return &PaginatedResponse[string]{
Results: []string{},
Count: 0,
Limit: limit,
Offset: offset,
HasMore: false,
}, nil
}
ctx := context.Background()
iter := NewIterator(fetchFunc, 10)
if iter.Next(ctx) {
t.Error("Next should return false for empty results")
}
}
func TestIterator_Reset(t *testing.T) {
items := []string{"item1", "item2", "item3"}
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
if offset >= len(items) {
return &PaginatedResponse[string]{
Results: []string{},
HasMore: false,
}, nil
}
return &PaginatedResponse[string]{
Results: items[offset:],
HasMore: false,
}, nil
}
ctx := context.Background()
iter := NewIterator(fetchFunc, 10)
// First iteration
count1 := 0
for iter.Next(ctx) {
count1++
}
// Reset and iterate again
iter.Reset()
count2 := 0
for iter.Next(ctx) {
count2++
}
if count1 != count2 {
t.Errorf("counts don't match after reset: %d vs %d", count1, count2)
}
if count1 != len(items) {
t.Errorf("got %d items, want %d", count1, len(items))
}
}
func TestIterator_Collect(t *testing.T) {
items := []string{"item1", "item2", "item3", "item4", "item5"}
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
if offset >= len(items) {
return &PaginatedResponse[string]{
Results: []string{},
HasMore: false,
}, nil
}
end := offset + limit
if end > len(items) {
end = len(items)
}
return &PaginatedResponse[string]{
Results: items[offset:end],
HasMore: end < len(items),
}, nil
}
ctx := context.Background()
iter := NewIterator(fetchFunc, 2)
results, err := iter.Collect(ctx)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(results) != len(items) {
t.Errorf("got %d items, want %d", len(results), len(items))
}
}
func TestIterator_CollectWithLimit(t *testing.T) {
items := []string{"item1", "item2", "item3", "item4", "item5"}
fetchFunc := func(ctx context.Context, offset, limit int) (*PaginatedResponse[string], error) {
if offset >= len(items) {
return &PaginatedResponse[string]{
Results: []string{},
HasMore: false,
}, nil
}
end := offset + limit
if end > len(items) {
end = len(items)
}
return &PaginatedResponse[string]{
Results: items[offset:end],
HasMore: end < len(items),
}, nil
}
ctx := context.Background()
iter := NewIterator(fetchFunc, 2)
maxItems := 3
results, err := iter.CollectWithLimit(ctx, maxItems)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(results) != maxItems {
t.Errorf("got %d items, want %d", len(results), maxItems)
}
}
func TestPage_HasNextPage(t *testing.T) {
tests := []struct {
name string
page Page[string]
want bool
}{
{
name: "has next page",
page: Page[string]{
Items: []string{"a", "b"},
Offset: 0,
Limit: 2,
Total: 5,
},
want: true,
},
{
name: "no next page",
page: Page[string]{
Items: []string{"d", "e"},
Offset: 3,
Limit: 2,
Total: 5,
},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.page.HasNextPage(); got != tt.want {
t.Errorf("HasNextPage() = %v, want %v", got, tt.want)
}
})
}
}
func TestPage_HasPrevPage(t *testing.T) {
tests := []struct {
name string
page Page[string]
want bool
}{
{
name: "has previous page",
page: Page[string]{
Offset: 2,
},
want: true,
},
{
name: "no previous page",
page: Page[string]{
Offset: 0,
},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.page.HasPrevPage(); got != tt.want {
t.Errorf("HasPrevPage() = %v, want %v", got, tt.want)
}
})
}
}
func TestPage_NextOffset(t *testing.T) {
page := Page[string]{
Offset: 10,
Limit: 5,
}
if got := page.NextOffset(); got != 15 {
t.Errorf("NextOffset() = %v, want 15", got)
}
}
func TestPage_PrevOffset(t *testing.T) {
tests := []struct {
name string
offset int
limit int
want int
}{
{"normal case", 10, 5, 5},
{"at beginning", 3, 5, 0},
{"exact boundary", 5, 5, 0},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
page := Page[string]{
Offset: tt.offset,
Limit: tt.limit,
}
if got := page.PrevOffset(); got != tt.want {
t.Errorf("PrevOffset() = %v, want %v", got, tt.want)
}
})
}
}

View File

@@ -1,219 +0,0 @@
package retry
import (
"context"
"net/http"
"testing"
"time"
)
func TestNew(t *testing.T) {
tests := []struct {
name string
config Config
want Config
}{
{
name: "default config",
config: Config{},
want: DefaultConfig(),
},
{
name: "custom config",
config: Config{
MaxAttempts: 5,
InitialDelay: 2 * time.Second,
MaxDelay: 60 * time.Second,
BackoffFactor: 3.0,
RetryableErrors: []int{500},
},
want: Config{
MaxAttempts: 5,
InitialDelay: 2 * time.Second,
MaxDelay: 60 * time.Second,
BackoffFactor: 3.0,
RetryableErrors: []int{500},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
r := New(tt.config)
if r.config.MaxAttempts != tt.want.MaxAttempts {
t.Errorf("MaxAttempts = %v, want %v", r.config.MaxAttempts, tt.want.MaxAttempts)
}
if r.config.InitialDelay != tt.want.InitialDelay {
t.Errorf("InitialDelay = %v, want %v", r.config.InitialDelay, tt.want.InitialDelay)
}
})
}
}
func TestRetrier_isRetryable(t *testing.T) {
r := New(DefaultConfig())
tests := []struct {
name string
statusCode int
want bool
}{
{"429 rate limit", 429, true},
{"500 server error", 500, true},
{"502 bad gateway", 502, true},
{"503 service unavailable", 503, true},
{"504 gateway timeout", 504, true},
{"200 ok", 200, false},
{"400 bad request", 400, false},
{"404 not found", 404, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resp := &http.Response{StatusCode: tt.statusCode}
if got := r.isRetryable(resp); got != tt.want {
t.Errorf("isRetryable() = %v, want %v", got, tt.want)
}
})
}
}
func TestRetrier_Do_Success(t *testing.T) {
r := New(DefaultConfig())
ctx := context.Background()
attempts := 0
fn := func() (*http.Response, error) {
attempts++
return &http.Response{StatusCode: 200}, nil
}
resp, err := r.Do(ctx, fn)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if resp.StatusCode != 200 {
t.Errorf("status = %d, want 200", resp.StatusCode)
}
if attempts != 1 {
t.Errorf("attempts = %d, want 1", attempts)
}
}
func TestRetrier_Do_Retry(t *testing.T) {
config := Config{
MaxAttempts: 3,
InitialDelay: 10 * time.Millisecond,
MaxDelay: 100 * time.Millisecond,
BackoffFactor: 2.0,
RetryableErrors: []int{500},
}
r := New(config)
ctx := context.Background()
attempts := 0
fn := func() (*http.Response, error) {
attempts++
if attempts < 3 {
return &http.Response{StatusCode: 500}, nil
}
return &http.Response{StatusCode: 200}, nil
}
start := time.Now()
resp, err := r.Do(ctx, fn)
elapsed := time.Since(start)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if resp.StatusCode != 200 {
t.Errorf("status = %d, want 200", resp.StatusCode)
}
if attempts != 3 {
t.Errorf("attempts = %d, want 3", attempts)
}
// Should have waited at least once
if elapsed < 10*time.Millisecond {
t.Errorf("elapsed time %v too short, expected at least 10ms", elapsed)
}
}
func TestRetrier_Do_ContextCancellation(t *testing.T) {
config := Config{
MaxAttempts: 5,
InitialDelay: 100 * time.Millisecond,
MaxDelay: 1 * time.Second,
BackoffFactor: 2.0,
RetryableErrors: []int{500},
}
r := New(config)
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
attempts := 0
fn := func() (*http.Response, error) {
attempts++
return &http.Response{StatusCode: 500}, nil
}
_, err := r.Do(ctx, fn)
if err != context.DeadlineExceeded {
t.Errorf("error = %v, want context.DeadlineExceeded", err)
}
if attempts > 2 {
t.Errorf("attempts = %d, should be cancelled early", attempts)
}
}
func TestRetrier_calculateDelay(t *testing.T) {
config := Config{
InitialDelay: time.Second,
MaxDelay: 30 * time.Second,
BackoffFactor: 2.0,
}
r := New(config)
tests := []struct {
name string
attempt int
resp *http.Response
minWant time.Duration
maxWant time.Duration
}{
{
name: "first retry",
attempt: 1,
resp: &http.Response{StatusCode: 500},
minWant: 500 * time.Millisecond, // with jitter min 0.5
maxWant: 1 * time.Second, // with jitter max 1.0
},
{
name: "second retry",
attempt: 2,
resp: &http.Response{StatusCode: 500},
minWant: 1 * time.Second, // 2^1 * 1s * 0.5
maxWant: 2 * time.Second, // 2^1 * 1s * 1.0
},
{
name: "with Retry-After header",
attempt: 1,
resp: &http.Response{
StatusCode: 429,
Header: http.Header{"Retry-After": []string{"5"}},
},
minWant: 5 * time.Second,
maxWant: 5 * time.Second,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
delay := r.calculateDelay(tt.attempt, tt.resp)
if delay < tt.minWant || delay > tt.maxWant {
t.Errorf("delay = %v, want between %v and %v", delay, tt.minWant, tt.maxWant)
}
})
}
}