tools

type AccessFilterMiddleware struct {
	// contains filtered or unexported fields
}

type AddGlossaryTermInput struct {
	URN        string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	TermURN    string `json:"term_urn" jsonschema_description:"The URN of the glossary term to add (e.g., urn:li:glossaryTerm:Classification)"`
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type AddGlossaryTermOutput struct {
	URN    string `json:"urn"`
	Term   string `json:"term"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}

type AddLinkInput struct {
	URN         string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	URL         string `json:"url" jsonschema_description:"The URL of the link to add"`
	Description string `json:"description" jsonschema_description:"A description of the link"`
	Connection  string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type AddLinkOutput struct {
	URN    string `json:"urn"`
	URL    string `json:"url"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}

type AddTagInput struct {
	URN        string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	TagURN     string `json:"tag_urn" jsonschema_description:"The URN of the tag to add (e.g., urn:li:tag:PII)"`
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type AddTagOutput struct {
	URN    string `json:"urn"`
	Tag    string `json:"tag"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}

type AfterFunc func(ctx context.Context, tc *ToolContext, result *mcp.CallToolResult, err error) (*mcp.CallToolResult, error)

type AuditLoggerMiddleware struct {
	// contains filtered or unexported fields
}

type BeforeFunc func(ctx context.Context, tc *ToolContext) (context.Context, error)

type Config struct {
	// DefaultLimit is the default search result limit. Default: 10.
	DefaultLimit int

	// MaxLimit is the maximum allowed limit. Default: 100.
	MaxLimit int

	// MaxLineageDepth is the maximum lineage traversal depth. Default: 5.
	MaxLineageDepth int

	// WriteEnabled enables write operations (add/remove tags, glossary terms, etc.).
	// Default: false (read-only mode).
	WriteEnabled bool

	// Debug enables debug logging for toolkit operations.
	Debug bool

	// Logger is the logger for debug output. If nil and Debug is true, uses client.StdLogger.
	Logger client.Logger
}

type ConnectionInfoOutput struct {
	Name      string `json:"name"`
	URL       string `json:"url"`
	IsDefault bool   `json:"is_default"`
}

type DataHubClient interface {
	// Search searches for entities.
	Search(ctx context.Context, query string, opts ...client.SearchOption) (*types.SearchResult, error)

	// GetEntity retrieves a single entity by URN.
	GetEntity(ctx context.Context, urn string) (*types.Entity, error)

	// GetSchema retrieves schema for a dataset.
	GetSchema(ctx context.Context, urn string) (*types.SchemaMetadata, error)

	// GetSchemas retrieves schemas for multiple datasets by URN.
	GetSchemas(ctx context.Context, urns []string) (map[string]*types.SchemaMetadata, error)

	// GetLineage retrieves lineage for an entity.
	GetLineage(ctx context.Context, urn string, opts ...client.LineageOption) (*types.LineageResult, error)

	// GetColumnLineage retrieves fine-grained column-level lineage for a dataset.
	GetColumnLineage(ctx context.Context, urn string) (*types.ColumnLineage, error)

	// GetQueries retrieves queries for a dataset.
	GetQueries(ctx context.Context, urn string) (*types.QueryList, error)

	// GetGlossaryTerm retrieves a glossary term.
	GetGlossaryTerm(ctx context.Context, urn string) (*types.GlossaryTerm, error)

	// ListTags lists all tags.
	ListTags(ctx context.Context, filter string) ([]types.Tag, error)

	// ListDomains lists all domains.
	ListDomains(ctx context.Context) ([]types.Domain, error)

	// ListDataProducts lists all data products.
	ListDataProducts(ctx context.Context) ([]types.DataProduct, error)

	// GetDataProduct retrieves a data product by URN.
	GetDataProduct(ctx context.Context, urn string) (*types.DataProduct, error)

	// Ping tests the connection.
	Ping(ctx context.Context) error

	// Close closes the client.
	Close() error

	// UpdateDescription sets the editable description for an entity.
	UpdateDescription(ctx context.Context, urn, description string) error

	// AddTag adds a tag to an entity.
	AddTag(ctx context.Context, urn, tagURN string) error

	// RemoveTag removes a tag from an entity.
	RemoveTag(ctx context.Context, urn, tagURN string) error

	// AddGlossaryTerm adds a glossary term to an entity.
	AddGlossaryTerm(ctx context.Context, urn, termURN string) error

	// RemoveGlossaryTerm removes a glossary term from an entity.
	RemoveGlossaryTerm(ctx context.Context, urn, termURN string) error

	// AddLink adds a link to an entity.
	AddLink(ctx context.Context, urn, linkURL, description string) error

	// RemoveLink removes a link from an entity by URL.
	RemoveLink(ctx context.Context, urn, linkURL string) error
}

type GetColumnLineageInput struct {
	URN string `json:"urn" jsonschema_description:"The DataHub URN of the dataset"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type GetDataProductInput struct {
	URN string `json:"urn" jsonschema_description:"The DataHub URN of the data product"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type GetEntityInput struct {
	URN string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type GetGlossaryTermInput struct {
	URN string `json:"urn" jsonschema_description:"The DataHub URN of the glossary term"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type GetLineageInput struct {
	URN       string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	Direction string `json:"direction,omitempty" jsonschema_description:"Lineage direction: UPSTREAM or DOWNSTREAM (default: DOWNSTREAM)"`
	Depth     int    `json:"depth,omitempty" jsonschema_description:"Maximum depth of lineage traversal (default: 1, max: 5)"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type GetQueriesInput struct {
	URN string `json:"urn" jsonschema_description:"The DataHub URN of the dataset"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type GetSchemaInput struct {
	URN string `json:"urn" jsonschema_description:"The DataHub URN of the dataset"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type ListConnectionsInput struct{}

type ListConnectionsOutput struct {
	Connections []ConnectionInfoOutput `json:"connections"`
	Count       int                    `json:"count"`
}

type ListDataProductsInput struct {
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type ListDomainsInput struct {
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type ListTagsInput struct {
	Filter string `json:"filter,omitempty" jsonschema_description:"Optional filter string to match tag names"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type MetadataEnricherMiddleware struct {
	// contains filtered or unexported fields
}

type RemoveGlossaryTermInput struct {
	URN        string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	TermURN    string `json:"term_urn" jsonschema_description:"The URN of the glossary term to remove"`
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type RemoveGlossaryTermOutput struct {
	URN    string `json:"urn"`
	Term   string `json:"term"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}

type RemoveLinkInput struct {
	URN        string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	URL        string `json:"url" jsonschema_description:"The URL of the link to remove"`
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type RemoveLinkOutput struct {
	URN    string `json:"urn"`
	URL    string `json:"url"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}

type RemoveTagInput struct {
	URN        string `json:"urn" jsonschema_description:"The DataHub URN of the entity"`
	TagURN     string `json:"tag_urn" jsonschema_description:"The URN of the tag to remove (e.g., urn:li:tag:PII)"`
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type RemoveTagOutput struct {
	URN    string `json:"urn"`
	Tag    string `json:"tag"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}

type SearchInput struct {
	Query string `json:"query" jsonschema_description:"Search query string"`
	// EntityType: DATASET, DASHBOARD, DATA_FLOW, DATA_JOB, CONTAINER, TAG, GLOSSARY_TERM, DATA_PRODUCT, etc.
	EntityType string `json:"entity_type,omitempty" jsonschema_description:"Entity type to search. Defaults to DATASET."`
	Limit      int    `json:"limit,omitempty" jsonschema_description:"Maximum number of results (default: 10, max: 100)"`
	Offset     int    `json:"offset,omitempty" jsonschema_description:"Result offset for pagination"`
	// Connection is the named connection to use. Empty uses the default connection.
	Connection string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type ToolContext struct {
	// ToolName is the name of the tool being executed.
	ToolName ToolName

	// Input is the parsed input for the tool.
	Input any

	// StartTime is when the tool execution started.
	StartTime time.Time

	// Extra allows middleware to pass data between Before and After.
	Extra map[string]any
}

type ToolMiddleware interface {
	// Before is called before the tool handler executes.
	// Return a modified context or an error to abort execution.
	Before(ctx context.Context, tc *ToolContext) (context.Context, error)

	// After is called after the tool handler executes.
	// Can modify the result or handle errors.
	After(ctx context.Context, tc *ToolContext, result *mcp.CallToolResult, err error) (*mcp.CallToolResult, error)
}

type ToolName string

type ToolOption func(*toolConfig)

type Toolkit struct {
	// contains filtered or unexported fields
}

type ToolkitOption func(*Toolkit)

type URNResolverMiddleware struct {
	// contains filtered or unexported fields
}

type UpdateDescriptionInput struct {
	URN         string `json:"urn" jsonschema_description:"The DataHub URN of the entity to update"`
	Description string `json:"description" jsonschema_description:"The new description text"`
	Connection  string `json:"connection,omitempty" jsonschema_description:"Named connection to use (see datahub_list_connections)"`
}

type UpdateDescriptionOutput struct {
	URN    string `json:"urn"`
	Aspect string `json:"aspect"`
	Action string `json:"action"`
}