From 81910b6618993a4a76b230c986a8ff598d43924e Mon Sep 17 00:00:00 2001 From: vvaks0 Date: Fri, 26 Jun 2026 23:54:48 -0400 Subject: [PATCH 1/3] feat(dq): add create_dq_job + prepare_create_dq_job MCP tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a guided, multi-turn data-quality job creation flow for DEV-197672. prepare_create_dq_job (read-only) walks the discovery chain — resolve the edge connection, detect the job type (PUSHDOWN/PULLUP) from its capabilities, and enumerate data sources, schemas, tables and columns — returning a ready-to-use plan. It also resolves the location from a catalog Table asset (id / URL / name) for the table-asset-page entry point. create_dq_job builds and submits the job via the PUBLIC DQ API (POST /rest/dq/1.0/jobs) behind a confirm-checkpoint preview. It supports column selection, a single-column row filter, row sampling, time-slice scheduling (${rd}/${rdEnd} composed into a dialect-aware source query), recurring schedules, monitors + adaptive settings, notifications, Pullup sizing / Parallel JDBC, and Pushdown compute. Discovery uses internal BFF endpoints (the public DQ API exposes no connection/edge metadata browse). Validated end-to-end on a live stack: create -> persist -> run FINISHED. Refs: DEV-197672 Co-Authored-By: Claude Opus 4.8 --- pkg/clients/dgc_client.go | 17 + pkg/clients/dgc_dq_client.go | 1081 ++++++++++++++++++ pkg/clients/dgc_dq_client_test.go | 129 +++ pkg/clients/dq_notifications.go | 199 ++++ pkg/clients/dq_notifications_test.go | 67 ++ pkg/clients/dq_source_query.go | 299 +++++ pkg/clients/dq_source_query_test.go | 93 ++ pkg/tools/create_dq_job/tool.go | 915 +++++++++++++++ pkg/tools/create_dq_job/tool_test.go | 1012 ++++++++++++++++ pkg/tools/prepare_create_dq_job/tool.go | 494 ++++++++ pkg/tools/prepare_create_dq_job/tool_test.go | 204 ++++ pkg/tools/register.go | 4 + 12 files changed, 4514 insertions(+) create mode 100644 pkg/clients/dgc_dq_client.go create mode 100644 pkg/clients/dgc_dq_client_test.go create mode 100644 pkg/clients/dq_notifications.go create mode 100644 pkg/clients/dq_notifications_test.go create mode 100644 pkg/clients/dq_source_query.go create mode 100644 pkg/clients/dq_source_query_test.go create mode 100644 pkg/tools/create_dq_job/tool.go create mode 100644 pkg/tools/create_dq_job/tool_test.go create mode 100644 pkg/tools/prepare_create_dq_job/tool.go create mode 100644 pkg/tools/prepare_create_dq_job/tool_test.go diff --git a/pkg/clients/dgc_client.go b/pkg/clients/dgc_client.go index e48156f..2c582e9 100644 --- a/pkg/clients/dgc_client.go +++ b/pkg/clients/dgc_client.go @@ -70,6 +70,23 @@ func GetAssetSummary( return ParseAssetDetailsGraphQLResponse(body) } +// GetAssetWithRelations fetches a single asset (with its first page of incoming/outgoing relations) +// by UUID string. Thin wrapper over GetAssetSummary for callers that have a string id. +func GetAssetWithRelations(ctx context.Context, collibraHttpClient *http.Client, assetID string) (*Asset, error) { + id, err := uuid.Parse(assetID) + if err != nil { + return nil, fmt.Errorf("invalid asset id %q: %w", assetID, err) + } + assets, err := GetAssetSummary(ctx, collibraHttpClient, id, "", "") + if err != nil { + return nil, err + } + if len(assets) == 0 { + return nil, fmt.Errorf("asset %s not found", assetID) + } + return &assets[0], nil +} + func ListAssetTypes(ctx context.Context, collibraHttpClient *http.Client, limit int, offset int) (*AssetTypePagedResponse, error) { slog.InfoContext(ctx, fmt.Sprintf("Listing asset types with limit: %d, offset: %d", limit, offset)) diff --git a/pkg/clients/dgc_dq_client.go b/pkg/clients/dgc_dq_client.go new file mode 100644 index 0000000..aa41b51 --- /dev/null +++ b/pkg/clients/dgc_dq_client.go @@ -0,0 +1,1081 @@ +package clients + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "net/url" + "strconv" + "strings" +) + +// This client wraps the DGC Unified Data Quality (UDQ) APIs used by the +// data-quality job-creation wizard. Discovery (connections, data sources, +// schemas, tables, columns) uses the internal BFF surface +// (/rest/dq/internal/v1/*) — the same calls the UI makes. Job creation uses +// the public DQ API (/rest/dq/1.0/jobs). + +// DqConnection is a data-quality edge connection as returned by +// /rest/dq/internal/v1/connections and /connections/{id}. capabilityTypes +// drives job-type detection: a connection advertises PUSHDOWN, PULLUP, or both. +type DqConnection struct { + ConnectionID string `json:"connectionId"` + ConnectionName string `json:"connectionName"` + CapabilityTypes []string `json:"capabilityTypes"` + DatabaseProductName string `json:"databaseProductName"` + EdgeSiteID string `json:"edgeSiteId"` + EdgeSiteName string `json:"edgeSiteName"` + SourceType *DqSourceType `json:"sourceType,omitempty"` + // SystemAssetID is the DGC System asset this connection ingests from (set via the catalog + // system-asset config). It is the bridge from a catalog asset back to a DQ connection. + SystemAssetID string `json:"systemAssetId,omitempty"` +} + +type DqSourceType struct { + Provider string `json:"provider"` + Type string `json:"type"` +} + +type dqConnectionListResponse struct { + Results []DqConnection `json:"results"` +} + +// DqDataSource is a database/catalog within a connection (e.g. "postgres"). +type DqDataSource struct { + DataSourceName string `json:"dataSourceName"` + SupportsSchemas bool `json:"supportsSchemas"` + TotalJobs int `json:"totalJobs"` +} + +// DqSchema is a schema within a data source. +type DqSchema struct { + Name string `json:"name"` +} + +// DqTable is a table within a schema. +type DqTable struct { + Name string `json:"name"` + Type string `json:"type"` +} + +// DqColumn is a column within a table. +type DqColumn struct { + Name string `json:"name"` + Type string `json:"type"` + Disabled bool `json:"disabled"` +} + +type dqDataSourceListResponse struct { + Results []DqDataSource `json:"results"` + Total int `json:"total"` +} + +type dqSchemaListResponse struct { + Results []DqSchema `json:"results"` + Total int `json:"total"` +} + +type dqTableListResponse struct { + Results []DqTable `json:"results"` + Total int `json:"total"` +} + +type dqColumnListResponse struct { + Results []DqColumn `json:"results"` +} + +// DqDataLocation identifies where a job reads data from. All five fields are +// required by the create API; databaseProductName is optional/read-only. +type DqDataLocation struct { + EdgeSiteName string `json:"edgeSiteName"` + EdgeConnectionName string `json:"edgeConnectionName"` + DataSourceName string `json:"dataSourceName"` + SchemaName string `json:"schemaName"` + TableName string `json:"tableName"` + DatabaseProductName string `json:"databaseProductName,omitempty"` +} + +// CreateDqJobRequest is the body for POST /rest/dq/1.0/jobs — the PUBLIC create +// (JobDefinitionCreateRequest in dq/udq-app-client/oas/dq-v1-public-oas-spec.yaml). It is the full +// job definition: sourceQuery (into which column selection, row filter, sampling and the ${rd}/${rdEnd} +// time-slice predicate are composed — see BuildDqSourceQuery), the runDate window, monitors, schedule, +// back-run, notifications, and pullup/pushdown settings. Unlike the internal BFF endpoint, the public +// server is null-tolerant ("provide only the fields you want to override") and auto-generates + +// auto-increments jobName when it is omitted. queueRun is currently always treated as true server-side +// (create-only is not yet supported), so it is left at its default and not sent. +type CreateDqJobRequest struct { + JobType string `json:"jobType"` + JobName string `json:"jobName,omitempty"` + DataLocation DqDataLocation `json:"dataLocation"` + SourceQuery string `json:"sourceQuery,omitempty"` + RunDate *DqPublicRunDate `json:"runDate,omitempty"` + RunDateEnd *DqPublicRunDate `json:"runDateEnd,omitempty"` + Backrun *DqPublicBackrun `json:"backrun,omitempty"` + JobSettings *DqPublicJobSettings `json:"jobSettings,omitempty"` + MonitoringSettings *DqPublicMonitoringSettings `json:"monitoringSettings,omitempty"` + Notifications *DqJobNotifications `json:"notifications,omitempty"` + SchedulingSettings *DqSchedulingSettings `json:"schedulingSettings,omitempty"` +} + +// DqPublicRunDate is the discriminated runDate/runDateEnd value ({kind, value}) from the public spec. +// kind=DATE => value is yyyy-MM-dd; kind=TIMESTAMP => value is RFC3339 (yyyy-MM-ddTHH:mm:ssZ). The +// engine substitutes ${rd}/${rdEnd} in sourceQuery from these per run (formatted per jobSettings.dateFormat). +type DqPublicRunDate struct { + Kind string `json:"kind"` + Value string `json:"value"` +} + +// DqPublicBackrun is the public Backrun: presence enables it (no `enabled` flag); binValue >= 1. +type DqPublicBackrun struct { + TimeBin string `json:"timeBin"` // DAY | MONTH | YEAR + BinValue int `json:"binValue"` +} + +// DqPublicJobSettings is jobSettings: the run-date dateFormat plus the type-specific tuning. +type DqPublicJobSettings struct { + DateFormat string `json:"dateFormat,omitempty"` // DATE | TIMESTAMP + PushdownSettings *DqPublicPushdownSettings `json:"pushdownSettings,omitempty"` + PullupSettings *DqPublicPullupSettings `json:"pullupSettings,omitempty"` +} + +// DqPublicPushdownSettings is the Pushdown compute (source-system concurrency). +type DqPublicPushdownSettings struct { + Connections int `json:"connections,omitempty"` + Threads int `json:"threads,omitempty"` +} + +// DqPublicPullupSettings is the Pullup tuning. Omitting sparkJobSizing = automatic sizing (the public +// API has no autoSizing flag — absence means auto). +type DqPublicPullupSettings struct { + LoadOptions *DqPublicLoadOptions `json:"loadOptions,omitempty"` + SparkJobSizing *DqPublicSparkJobSizing `json:"sparkJobSizing,omitempty"` + SparkSqlProperties map[string]string `json:"sparkSqlProperties,omitempty"` +} + +// DqPublicLoadOptions mirrors the public LoadOptions (numPartitions 0 = let Spark decide). +type DqPublicLoadOptions struct { + NumPartitions int `json:"numPartitions"` + ParallelJdbcOptions *DqParallelJdbcOptions `json:"parallelJdbcOptions,omitempty"` +} + +// DqPublicSparkJobSizing is manual Spark sizing; memory fields are INTEGER GB (SparkMemoryGB). Send +// this object only for manual sizing — omit it for automatic sizing. +type DqPublicSparkJobSizing struct { + NumExecutors int `json:"numExecutors,omitempty"` + DriverCores int `json:"driverCores,omitempty"` + NumExecutorCores int `json:"numExecutorCores,omitempty"` + ExecutorMemoryGb int `json:"executorMemoryGb,omitempty"` + DriverMemoryGb int `json:"driverMemoryGb,omitempty"` + MemoryOverheadGb int `json:"memoryOverheadGb,omitempty"` +} + +// DqPublicMonitoringSettings is monitoringSettings (currently only adaptive monitors). +type DqPublicMonitoringSettings struct { + AdaptiveMonitors *DqPublicAdaptiveMonitors `json:"adaptiveMonitors,omitempty"` +} + +// DqPublicAdaptiveMonitors mirrors the public AdaptiveMonitors — the same toggle set as the internal +// profileMonitors. settings holds the adaptive lookback/learning tuning. +type DqPublicAdaptiveMonitors struct { + DescriptiveStatistics bool `json:"descriptiveStatistics"` + EmptyFields bool `json:"emptyFields"` + ExecutionTime bool `json:"executionTime"` + Max bool `json:"max"` + Mean bool `json:"mean"` + Min bool `json:"min"` + NullValues bool `json:"nullValues"` + RowCount bool `json:"rowCount"` + Uniqueness bool `json:"uniqueness"` + Settings *DqPublicAdaptiveMonitorSettings `json:"settings,omitempty"` +} + +// DqPublicAdaptiveMonitorSettings is the adaptive tuning. NOTE: the public field is dataLookBack +// (capital B), unlike the internal dataLookback. +type DqPublicAdaptiveMonitorSettings struct { + DataLookBack int `json:"dataLookBack"` + LearningPhase int `json:"learningPhase"` +} + +// PublicAdaptiveMonitorsFromProfile maps the (shared) DqProfileMonitors toggle set onto the public +// AdaptiveMonitors shape, so callers keep using BuildProfileMonitors for the monitor selection. +func PublicAdaptiveMonitorsFromProfile(pm *DqProfileMonitors) *DqPublicAdaptiveMonitors { + if pm == nil { + return nil + } + return &DqPublicAdaptiveMonitors{ + DescriptiveStatistics: pm.DescriptiveStatistics, + EmptyFields: pm.EmptyFields, + ExecutionTime: pm.ExecutionTime, + Max: pm.Max, + Mean: pm.Mean, + Min: pm.Min, + NullValues: pm.NullValues, + RowCount: pm.RowCount, + Uniqueness: pm.Uniqueness, + } +} + +// CreateDqJobResponse is the public create response (JobDefinitionCreateResponse): the created job +// definition plus the queued run id. +type CreateDqJobResponse struct { + JobName string `json:"jobName"` + JobType string `json:"jobType"` + JobRunID string `json:"jobRunId"` + DataLocation DqDataLocation `json:"dataLocation"` + SourceQuery string `json:"sourceQuery,omitempty"` +} + +// ListDqConnections returns all data-quality connections on the instance. +func ListDqConnections(ctx context.Context, collibraHttpClient *http.Client) ([]DqConnection, error) { + req, err := http.NewRequestWithContext(ctx, "GET", "/rest/dq/internal/v1/connections", nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return nil, err + } + var resp dqConnectionListResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse connections response: %w", err) + } + return resp.Results, nil +} + +// GetDqConnection fetches a single connection by id. The response carries the +// job-type (capabilityTypes) and the dataLocation fields edgeSiteName, +// connectionName, and databaseProductName. +func GetDqConnection(ctx context.Context, collibraHttpClient *http.Client, connectionID string) (*DqConnection, error) { + endpoint := "/rest/dq/internal/v1/connections/" + url.PathEscape(connectionID) + req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return nil, err + } + var conn DqConnection + if err := json.Unmarshal(body, &conn); err != nil { + return nil, fmt.Errorf("failed to parse connection response: %w", err) + } + return &conn, nil +} + +// ListDqDataSources lists the databases/catalogs reachable through a connection. +func ListDqDataSources(ctx context.Context, collibraHttpClient *http.Client, connectionID string, limit, offset int) ([]DqDataSource, error) { + endpoint := fmt.Sprintf("/rest/dq/internal/v1/monitoring/edge/connections/%s/dataSources?%s", + url.PathEscape(connectionID), dqPageQuery(nil, limit, offset)) + body, err := dqGet(ctx, collibraHttpClient, endpoint) + if err != nil { + return nil, err + } + var resp dqDataSourceListResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse data sources response: %w", err) + } + return resp.Results, nil +} + +// ListDqSchemas lists schemas in a data source (live edge query). +func ListDqSchemas(ctx context.Context, collibraHttpClient *http.Client, siteID, connectionID, dataSourceName string, limit, offset int) ([]DqSchema, error) { + endpoint := fmt.Sprintf("/rest/dq/internal/v1/monitoring/edge/%s/connections/%s/schemas?%s", + url.PathEscape(siteID), url.PathEscape(connectionID), + dqPageQuery(url.Values{"dataSourceName": {dataSourceName}}, limit, offset)) + body, err := dqGet(ctx, collibraHttpClient, endpoint) + if err != nil { + return nil, err + } + var resp dqSchemaListResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse schemas response: %w", err) + } + return resp.Results, nil +} + +// ListDqTables lists tables in a schema (live edge query). +func ListDqTables(ctx context.Context, collibraHttpClient *http.Client, siteID, connectionID, dataSourceName, schemaName string, limit, offset int) ([]DqTable, error) { + endpoint := fmt.Sprintf("/rest/dq/internal/v1/monitoring/edge/%s/connections/%s/tables?%s", + url.PathEscape(siteID), url.PathEscape(connectionID), + dqPageQuery(url.Values{"dataSourceName": {dataSourceName}, "schemaName": {schemaName}}, limit, offset)) + body, err := dqGet(ctx, collibraHttpClient, endpoint) + if err != nil { + return nil, err + } + var resp dqTableListResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse tables response: %w", err) + } + return resp.Results, nil +} + +// ListDqColumns lists columns in a table (live edge query). +func ListDqColumns(ctx context.Context, collibraHttpClient *http.Client, siteID, connectionID, dataSourceName, schemaName, tableName string, limit, offset int) ([]DqColumn, error) { + endpoint := fmt.Sprintf("/rest/dq/internal/v1/monitoring/edge/%s/connections/%s/columns?%s", + url.PathEscape(siteID), url.PathEscape(connectionID), + dqPageQuery(url.Values{"dataSourceName": {dataSourceName}, "schemaName": {schemaName}, "tableName": {tableName}}, limit, offset)) + body, err := dqGet(ctx, collibraHttpClient, endpoint) + if err != nil { + return nil, err + } + var resp dqColumnListResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse columns response: %w", err) + } + return resp.Results, nil +} + +// CreateDqJob creates a data-quality job and queues an immediate run. +func CreateDqJob(ctx context.Context, collibraHttpClient *http.Client, request CreateDqJobRequest) (*CreateDqJobResponse, error) { + slog.InfoContext(ctx, fmt.Sprintf("Creating DQ job '%s' (%s) for %s.%s", + request.JobName, request.JobType, request.DataLocation.SchemaName, request.DataLocation.TableName)) + + jsonData, err := json.Marshal(request) + if err != nil { + return nil, fmt.Errorf("failed to marshal create job request: %w", err) + } + req, err := http.NewRequestWithContext(ctx, "POST", "/rest/dq/1.0/jobs", bytes.NewBuffer(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return nil, err + } + var resp CreateDqJobResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse create job response: %w", err) + } + return &resp, nil +} + +// DqParallelJdbcOptions is the wizard's "Parallel JDBC" advanced-sizing control (PULLUP). mode is the +// ParallelJdbcMode enum (udq-app): AUTO (column + partition count both auto), AUTO_COLUMN (column auto, +// partitionNumber required), MANUAL (partitionColumn + partitionNumber both required). Validator rules +// (ParallelJdbcOptionsValidator): AUTO needs neither; AUTO_COLUMN needs partitionNumber and no column; +// MANUAL needs both. +type DqParallelJdbcOptions struct { + Mode string `json:"mode"` // AUTO | AUTO_COLUMN | MANUAL + PartitionColumn string `json:"partitionColumn,omitempty"` // only for MANUAL + PartitionNumber int `json:"partitionNumber,omitempty"` // required for AUTO_COLUMN and MANUAL +} + +// Valid ParallelJdbcMode values (udq-app ParallelJdbcMode enum). +const ( + ParallelJdbcAuto = "AUTO" + ParallelJdbcAutoColumn = "AUTO_COLUMN" + ParallelJdbcManual = "MANUAL" +) + +// DqProfileMonitorSettings — the "Advanced monitor settings" (adaptive behavior). dataLookback +// = how many prior runs feed the adaptive baseline (wizard default 10); learningPhase = runs +// before adaptive monitors start alerting (wizard default 4). Null-safe: JobMonitorsMapper reads +// it only when non-null. Contract: ProfileMonitorSettings in ui-v1-private-oas-spec.yaml. +type DqProfileMonitorSettings struct { + DataLookback int `json:"dataLookback"` + LearningPhase int `json:"learningPhase"` +} + +// DqProfileMonitors is the monitor-toggle set (the "Monitors" step). Wizard defaults = +// rowCount/uniqueness/nullValues/emptyFields ON, the rest OFF. BuildProfileMonitors builds it from the +// selected keys; PublicAdaptiveMonitorsFromProfile maps it onto the public adaptiveMonitors shape. +type DqProfileMonitors struct { + DescriptiveStatistics bool `json:"descriptiveStatistics"` + EmptyFields bool `json:"emptyFields"` + ExecutionTime bool `json:"executionTime"` + Max bool `json:"max"` + Mean bool `json:"mean"` + Min bool `json:"min"` + NullValues bool `json:"nullValues"` + RowCount bool `json:"rowCount"` + Uniqueness bool `json:"uniqueness"` + Settings *DqProfileMonitorSettings `json:"settings,omitempty"` +} + +// DqMonitorInfo describes one profile monitor for display and selection. Key matches the +// DqProfileMonitors JSON field and the create_dq_job `monitors` input; DefaultEnabled marks +// the monitors the wizard turns on by default. +type DqMonitorInfo struct { + Key string `json:"key"` + Label string `json:"label"` + Description string `json:"description"` + DefaultEnabled bool `json:"defaultEnabled"` +} + +// DqMonitorCatalog is the ordered list of selectable profile monitors with their defaults. +// Defaults match the DQ wizard: row count, null values, empty values, and uniqueness ON; the +// numeric/timing monitors OFF; descriptiveStatistics OFF because enabling it UNMASKS sensitive +// data (the server sets maskSensitive = !descriptiveStatistics in JobMonitorsMapper). +func DqMonitorCatalog() []DqMonitorInfo { + return []DqMonitorInfo{ + {Key: "rowCount", Label: "Row count", Description: "Track the row count per run.", DefaultEnabled: true}, + {Key: "nullValues", Label: "Null values", Description: "Track the share of NULLs per column.", DefaultEnabled: true}, + {Key: "emptyFields", Label: "Empty values", Description: "Track empty/blank values per column.", DefaultEnabled: true}, + {Key: "uniqueness", Label: "Uniqueness", Description: "Track distinct/duplicate values per column.", DefaultEnabled: true}, + {Key: "min", Label: "Minimum value", Description: "Track the minimum value per numeric column.", DefaultEnabled: false}, + {Key: "mean", Label: "Mean value", Description: "Track the mean value per numeric column.", DefaultEnabled: false}, + {Key: "max", Label: "Maximum value", Description: "Track the maximum value per numeric column.", DefaultEnabled: false}, + {Key: "executionTime", Label: "Execution time", Description: "Track how long each run takes.", DefaultEnabled: false}, + {Key: "descriptiveStatistics", Label: "Descriptive statistics", Description: "Compute descriptive statistics. WARNING: UNMASKS sensitive data (maskSensitive=false) — enable only with explicit user confirmation.", DefaultEnabled: false}, + } +} + +// MonitorKeys returns every valid monitor key, in catalog order. +func MonitorKeys() []string { + cat := DqMonitorCatalog() + keys := make([]string, 0, len(cat)) + for _, m := range cat { + keys = append(keys, m.Key) + } + return keys +} + +// DefaultMonitorKeys returns the keys of the monitors that are enabled by default. +func DefaultMonitorKeys() []string { + var keys []string + for _, m := range DqMonitorCatalog() { + if m.DefaultEnabled { + keys = append(keys, m.Key) + } + } + return keys +} + +// BuildProfileMonitors turns a set of enabled monitor keys (case-insensitive) into a +// DqProfileMonitors. Keys not in the catalog are returned in `unknown` so the caller can +// reject them; the returned monitors reflect only the recognized keys (all others OFF). +func BuildProfileMonitors(enabledKeys []string) (*DqProfileMonitors, []string) { + valid := map[string]bool{} + for _, m := range DqMonitorCatalog() { + valid[strings.ToLower(m.Key)] = true + } + on := map[string]bool{} + var unknown []string + for _, k := range enabledKeys { + lk := strings.ToLower(strings.TrimSpace(k)) + if lk == "" { + continue + } + if !valid[lk] { + unknown = append(unknown, k) + continue + } + on[lk] = true + } + return &DqProfileMonitors{ + DescriptiveStatistics: on["descriptivestatistics"], + EmptyFields: on["emptyfields"], + ExecutionTime: on["executiontime"], + Max: on["max"], + Mean: on["mean"], + Min: on["min"], + NullValues: on["nullvalues"], + RowCount: on["rowcount"], + Uniqueness: on["uniqueness"], + }, unknown +} + +// EnabledMonitorKeys returns the catalog keys enabled in pm, in catalog order (for display). +func EnabledMonitorKeys(pm *DqProfileMonitors) []string { + if pm == nil { + return nil + } + state := map[string]bool{ + "rowcount": pm.RowCount, "nullvalues": pm.NullValues, "emptyfields": pm.EmptyFields, + "uniqueness": pm.Uniqueness, "min": pm.Min, "mean": pm.Mean, "max": pm.Max, + "executiontime": pm.ExecutionTime, "descriptivestatistics": pm.DescriptiveStatistics, + } + var keys []string + for _, m := range DqMonitorCatalog() { + if state[strings.ToLower(m.Key)] { + keys = append(keys, m.Key) + } + } + return keys +} + +// DqAdaptiveMonitorSetting describes one tunable "Advanced monitor setting" (the adaptive +// behavior in the Monitors step). Key matches the create_dq_job input; Default is the wizard +// default applied when the user doesn't override. +type DqAdaptiveMonitorSetting struct { + Key string `json:"key"` + Label string `json:"label"` + Description string `json:"description"` + Default int `json:"default"` +} + +// DqAdaptiveMonitorSettings is the catalog of advanced monitor settings to show the user, +// surfaced alongside the monitor toggles so the agent can offer them explicitly. +func DqAdaptiveMonitorSettings() []DqAdaptiveMonitorSetting { + return []DqAdaptiveMonitorSetting{ + {Key: "dataLookback", Label: "Data lookback", Description: "Number of prior runs used as the adaptive baseline.", Default: 10}, + {Key: "learningPhase", Label: "Learning phase", Description: "Number of runs before adaptive monitors begin alerting.", Default: 4}, + } +} + +// DqSchedulingSettings — schedulerMode HOURLY/DAILY/MONTHLY + scheduledRunTime (HH:mm:ss UTC), +// plus the mode-specific sub-object the server's SchedulingSettingsMapper requires: +// - DAILY (also Weekly/Weekdays): daily.daysOfWeek + daily.dailyOffset +// - HOURLY: hourly.hourlyOffset +// - MONTHLY: monthly.{monthlyRepeat, dayNumber, monthlyOffset} +// +// The mapper dereferences these sub-objects unguarded, so the chosen mode's sub-object MUST be +// populated. The *Offset enums set the run-date offset that drives ${rd}/${rdEnd} per run. +type DqSchedulingSettings struct { + SchedulerMode string `json:"schedulerMode"` + ScheduledRunTime string `json:"scheduledRunTime"` + IsActive bool `json:"isActive"` + Daily *DqDailySchedule `json:"daily,omitempty"` + Hourly *DqHourlySchedule `json:"hourly,omitempty"` + Monthly *DqMonthlySchedule `json:"monthly,omitempty"` +} + +// DqDailySchedule drives DAILY/Weekly/Weekdays. daysOfWeek is required (>=1); dailyOffset is the +// run-date offset (SCHEDULED = the run's own day, ONE_DAY..SEVEN_DAYS = that many days back). +type DqDailySchedule struct { + DailyOffset string `json:"dailyOffset"` + DaysOfWeek []string `json:"daysOfWeek"` +} + +// DqHourlySchedule drives HOURLY. hourlyOffset: SCHEDULED | ONE_HOUR | TWO_HOURS. +type DqHourlySchedule struct { + HourlyOffset string `json:"hourlyOffset"` +} + +// DqMonthlySchedule drives MONTHLY. monthlyRepeat FIRST|LAST|DAY (DAY uses dayNumber 1-28); +// monthlyOffset: SCHEDULED | FIRST_OF_CURRENT_MONTH | FIRST_OF_PRIOR_MONTH | LAST_OF_PRIOR_MONTH. +type DqMonthlySchedule struct { + DayNumber int `json:"dayNumber,omitempty"` + MonthlyOffset string `json:"monthlyOffset"` + MonthlyRepeat string `json:"monthlyRepeat"` +} + +// Valid offset/day enum values (from ui-v1-private-oas-spec.yaml). +var ( + dqHourlyOffsets = []string{"SCHEDULED", "ONE_HOUR", "TWO_HOURS"} + dqDailyOffsets = []string{"SCHEDULED", "ONE_DAY", "TWO_DAYS", "THREE_DAYS", "FOUR_DAYS", "FIVE_DAYS", "SIX_DAYS", "SEVEN_DAYS"} + dqMonthlyOffsets = []string{"SCHEDULED", "FIRST_OF_CURRENT_MONTH", "FIRST_OF_PRIOR_MONTH", "LAST_OF_PRIOR_MONTH"} + dqDaysOfWeek = []string{"MONDAY", "TUESDAY", "WEDNESDAY", "THURSDAY", "FRIDAY", "SATURDAY", "SUNDAY"} + dqWeekdays = []string{"MONDAY", "TUESDAY", "WEDNESDAY", "THURSDAY", "FRIDAY"} +) + +// DqScheduleInput is the friendly schedule request that BuildSchedulingSettings validates and maps +// onto the API's mode-specific sub-objects. +type DqScheduleInput struct { + Repeat string // NEVER (default) | HOURLY | DAILY | WEEKLY | WEEKDAYS | MONTHLY + RunTime string // HH:mm[:ss] UTC; defaults to 00:00:00 + DaysOfWeek []string // for WEEKLY + DayOfMonth int // for MONTHLY DAY mode (1-28) + MonthlyMode string // for MONTHLY: DAY (default) | FIRST | LAST + RunDateOffset string // mode-specific offset; defaults to SCHEDULED +} + +// BuildSchedulingSettings validates the friendly schedule input and maps it onto the API shape. +// Returns (nil, nil) for NEVER/empty (run once now), the populated settings on success, or a +// descriptive error the caller can surface as needs_input. +func BuildSchedulingSettings(in DqScheduleInput) (*DqSchedulingSettings, error) { + repeat := strings.ToUpper(strings.TrimSpace(in.Repeat)) + if repeat == "" || repeat == "NEVER" { + return nil, nil + } + offset := strings.ToUpper(strings.TrimSpace(in.RunDateOffset)) + if offset == "" { + offset = "SCHEDULED" + } + settings := &DqSchedulingSettings{ScheduledRunTime: normalizeRunTime(in.RunTime), IsActive: true} + + switch repeat { + case "HOURLY": + if !dqContains(dqHourlyOffsets, offset) { + return nil, fmt.Errorf("runDateOffset %q is invalid for HOURLY (use one of: %s)", offset, strings.Join(dqHourlyOffsets, ", ")) + } + settings.SchedulerMode = "HOURLY" + settings.Hourly = &DqHourlySchedule{HourlyOffset: offset} + + case "DAILY", "WEEKLY", "WEEKDAYS": + if !dqContains(dqDailyOffsets, offset) { + return nil, fmt.Errorf("runDateOffset %q is invalid for %s (use one of: %s)", offset, repeat, strings.Join(dqDailyOffsets, ", ")) + } + days, err := resolveDaysOfWeek(repeat, in.DaysOfWeek) + if err != nil { + return nil, err + } + settings.SchedulerMode = "DAILY" + settings.Daily = &DqDailySchedule{DailyOffset: offset, DaysOfWeek: days} + + case "MONTHLY": + if !dqContains(dqMonthlyOffsets, offset) { + return nil, fmt.Errorf("runDateOffset %q is invalid for MONTHLY (use one of: %s)", offset, strings.Join(dqMonthlyOffsets, ", ")) + } + mode := strings.ToUpper(strings.TrimSpace(in.MonthlyMode)) + if mode == "" { + mode = "DAY" + } + monthly := &DqMonthlySchedule{MonthlyOffset: offset, MonthlyRepeat: mode} + switch mode { + case "DAY": + if in.DayOfMonth < 1 || in.DayOfMonth > 28 { + return nil, fmt.Errorf("scheduleDayOfMonth must be 1-28 for MONTHLY DAY mode (got %d)", in.DayOfMonth) + } + monthly.DayNumber = in.DayOfMonth + case "FIRST", "LAST": + // run on the first/last day of the month; no day number needed + default: + return nil, fmt.Errorf("scheduleMonthlyMode %q is invalid (use DAY, FIRST, or LAST)", mode) + } + settings.SchedulerMode = "MONTHLY" + settings.Monthly = monthly + + default: + return nil, fmt.Errorf("scheduleRepeat %q is invalid (use NEVER, HOURLY, DAILY, WEEKLY, WEEKDAYS, or MONTHLY)", repeat) + } + return settings, nil +} + +// resolveDaysOfWeek returns the active days for a daily-family mode: all 7 for DAILY, Mon-Fri for +// WEEKDAYS, or the validated user set for WEEKLY (which requires at least one valid day). +func resolveDaysOfWeek(repeat string, userDays []string) ([]string, error) { + switch repeat { + case "DAILY": + return dqDaysOfWeek, nil + case "WEEKDAYS": + return dqWeekdays, nil + default: // WEEKLY + var days []string + for _, d := range userDays { + up := strings.ToUpper(strings.TrimSpace(d)) + if up == "" { + continue + } + if !dqContains(dqDaysOfWeek, up) { + return nil, fmt.Errorf("scheduleDaysOfWeek contains invalid day %q (use: %s)", d, strings.Join(dqDaysOfWeek, ", ")) + } + days = append(days, up) + } + if len(days) == 0 { + return nil, fmt.Errorf("WEEKLY requires at least one day in scheduleDaysOfWeek (e.g. MONDAY)") + } + return days, nil + } +} + +// normalizeRunTime pads HH:mm to HH:mm:ss and defaults empty to 00:00:00. +func normalizeRunTime(t string) string { + t = strings.TrimSpace(t) + if t == "" { + return "00:00:00" + } + if strings.Count(t, ":") == 1 { + return t + ":00" + } + return t +} + +func dqContains(set []string, v string) bool { + for _, s := range set { + if s == v { + return true + } + } + return false +} + +// GetDqDataDistribution wraps the wizard's "Show distribution" / days-with-data helper. +// WHY INTERNAL: there is no public equivalent — this BFF endpoint runs a live edge query to +// return the row distribution over a date column, so the agent can show the user which days +// actually have data before choosing a time-slice / backrun range. +// GET /rest/dq/internal/v1/explorer/{connId}/data/distribution +func GetDqDataDistribution(ctx context.Context, collibraHttpClient *http.Client, connectionID, dataSourceName, schemaName, tableName, columnName, groupBy string, isDate bool) (json.RawMessage, error) { + v := url.Values{ + "databaseName": {dataSourceName}, + "schemaName": {schemaName}, + "tableName": {tableName}, + "columnName": {columnName}, + "groupBy": {groupBy}, // e.g. DAY + "isDate": {strconv.FormatBool(isDate)}, + } + endpoint := fmt.Sprintf("/rest/dq/internal/v1/explorer/%s/data/distribution?%s", url.PathEscape(connectionID), v.Encode()) + body, err := dqGet(ctx, collibraHttpClient, endpoint) + if err != nil { + return nil, err + } + return json.RawMessage(body), nil +} + +// ===================================================================================== +// Job-name helpers (the wizard's Step 1 name handling) — all on the internal jobs surface. +// ===================================================================================== + +type dqJobNameRequest struct { + SchemaName string `json:"schemaName"` + TableName string `json:"tableName"` +} + +type dqJobNameResponse struct { + JobName string `json:"jobName"` +} + +// GenerateUniqueJobName asks the server for a collision-free default job name for schema.table. +// The server auto-increments (e.g. "sales.orders_2") when the base name is taken, so the user is +// never asked to resolve a name conflict manually — POST /rest/dq/internal/v1/jobs/name. +func GenerateUniqueJobName(ctx context.Context, collibraHttpClient *http.Client, schemaName, tableName string) (string, error) { + payload, err := json.Marshal(dqJobNameRequest{SchemaName: schemaName, TableName: tableName}) + if err != nil { + return "", fmt.Errorf("failed to marshal job-name request: %w", err) + } + req, err := http.NewRequestWithContext(ctx, "POST", "/rest/dq/internal/v1/jobs/name", bytes.NewBuffer(payload)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return "", err + } + var resp dqJobNameResponse + if err := json.Unmarshal(body, &resp); err != nil { + return "", fmt.Errorf("failed to parse job-name response: %w", err) + } + return resp.JobName, nil +} + +// IsValidDqJobName asks the server whether a job name is syntactically valid — the wizard rejects +// special characters other than - and _. GET /rest/dq/internal/v1/jobs/{name}/validJobName. +func IsValidDqJobName(ctx context.Context, collibraHttpClient *http.Client, jobName string) (bool, error) { + return dqGetBool(ctx, collibraHttpClient, "/rest/dq/internal/v1/jobs/"+url.PathEscape(jobName)+"/validJobName") +} + +// DqJobExists reports whether a job with the given name already exists — +// GET /rest/dq/internal/v1/jobs/{name}/exists. +func DqJobExists(ctx context.Context, collibraHttpClient *http.Client, jobName string) (bool, error) { + return dqGetBool(ctx, collibraHttpClient, "/rest/dq/internal/v1/jobs/"+url.PathEscape(jobName)+"/exists") +} + +func dqGetBool(ctx context.Context, collibraHttpClient *http.Client, endpoint string) (bool, error) { + body, err := dqGet(ctx, collibraHttpClient, endpoint) + if err != nil { + return false, err + } + var b bool + if err := json.Unmarshal(body, &b); err != nil { + return false, fmt.Errorf("failed to parse boolean response from %s: %w", endpoint, err) + } + return b, nil +} + +// ===================================================================================== +// DQ permission preflight (DGC-core global permissions). +// +// The DQ wizard gates the Create/Schedule/Run actions on the invoking user's GLOBAL permissions +// (the same identifiers the frontend checks via currentUser.global). We read them from the public +// core endpoint GET /rest/2.0/users/current/globalPermissions, which returns the enum names below. +// Source of identifiers: iam-authorization Permission enum (DATA_QUALITY_JOB_*). +// ===================================================================================== + +const ( + PermDqJobCreate = "DATA_QUALITY_JOB_CREATE" + PermDqJobRun = "DATA_QUALITY_JOB_RUN" + PermDqJobSchedule = "DATA_QUALITY_JOB_SCHEDULE" + PermDqJobEdit = "DATA_QUALITY_JOB_EDIT" + PermResourceManageAll = "RESOURCE_MANAGE_ALL" +) + +// dqConnectionResourcePrefix is the resource-id prefix the DQ UI uses for connection-scoped +// permissions (frontend DATA_QUALITY_CONNECTION_RESOURCE_PREFIX). +const dqConnectionResourcePrefix = "DataQualityConnection" + +// HasPermission reports whether perm is present in perms (case-insensitive). +func HasPermission(perms []string, perm string) bool { + for _, p := range perms { + if strings.EqualFold(strings.TrimSpace(p), perm) { + return true + } + } + return false +} + +// GetCurrentUserGlobalPermissions returns the invoking user's GLOBAL permission identifiers +// (e.g. DATA_QUALITY_JOB_CREATE) — GET /rest/2.0/users/current/globalPermissions. NOTE: DQ +// create/run/schedule are usually granted as CONNECTION-resource permissions, not global ones — +// use GetDqConnectionPermissions for the DQ preflight; this is kept for global-only checks. +func GetCurrentUserGlobalPermissions(ctx context.Context, collibraHttpClient *http.Client) ([]string, error) { + req, err := http.NewRequestWithContext(ctx, "GET", "/rest/2.0/users/current/globalPermissions", nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Accept", "application/json") + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return nil, err + } + var resp struct { + GlobalPermissions []string `json:"globalPermissions"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse permissions response: %w", err) + } + return resp.GlobalPermissions, nil +} + +// dqPermissionsQuery mirrors the DQ UI's permission lookup (frontend permissions query): it returns +// the current user's global permissions plus the permissions on a specific resource. The DQ wizard +// gates create/run/schedule on the CONNECTION resource (`global || resource`). +const dqPermissionsQuery = `query permissions($shouldIncludeResource: Boolean!, $resourceId: String) { + api { + currentUser { + global: permissions + resource: permissions(resourceId: $resourceId) @include(if: $shouldIncludeResource) + } + } +}` + +type dqPermissionsResponse struct { + Data *struct { + Api *struct { + CurrentUser *struct { + Global []string `json:"global"` + Resource []string `json:"resource"` + } `json:"currentUser"` + } `json:"api"` + } `json:"data"` + Errors []Error `json:"errors"` +} + +// GetDqConnectionPermissions returns the invoking user's (global, connectionResource) permission +// identifiers for the given DQ connection — POST /graphql, the same query the DQ UI uses. Check a +// permission with `HasPermission(global, p) || HasPermission(resource, p)`. +func GetDqConnectionPermissions(ctx context.Context, collibraHttpClient *http.Client, connectionID string) (global, resource []string, err error) { + reqBody := Request{ + Query: dqPermissionsQuery, + Variables: map[string]interface{}{ + "shouldIncludeResource": true, + "resourceId": dqConnectionResourcePrefix + ":" + connectionID, + }, + } + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal permissions query: %w", err) + } + req, err := http.NewRequestWithContext(ctx, "POST", "/graphql", bytes.NewBuffer(jsonData)) + if err != nil { + return nil, nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return nil, nil, err + } + var resp dqPermissionsResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, nil, fmt.Errorf("failed to parse permissions response: %w", err) + } + if len(resp.Errors) > 0 { + return nil, nil, fmt.Errorf("permissions query errors: %v", resp.Errors) + } + if resp.Data == nil || resp.Data.Api == nil || resp.Data.Api.CurrentUser == nil { + return nil, nil, fmt.Errorf("permissions response missing currentUser") + } + return resp.Data.Api.CurrentUser.Global, resp.Data.Api.CurrentUser.Resource, nil +} + +// DqJobDetailsPath returns the Job Details deep-link path for a created job. The DQ SPA route +// /data-quality/jobs takes the jobName and resolves the latest run itself. It is relative to the +// Collibra instance base URL (the chip client reaches DGC over an internal URL, so the public host +// is prepended by the calling surface). +func DqJobDetailsPath(jobName string) string { + return "/data-quality/jobs?jobName=" + url.QueryEscape(jobName) +} + +// CatalogAssetPath returns the catalog deep-link path for an asset (relative to the instance URL). +func CatalogAssetPath(assetID string) string { + return "/asset/" + url.PathEscape(assetID) +} + +// DqTableAssetLocation is a catalog Table asset resolved to its DQ data location. +type DqTableAssetLocation struct { + TableAssetID string + SystemAssetID string + ConnectionID string + ConnectionName string + EdgeSiteName string + DataSourceName string + SchemaName string + TableName string +} + +// ResolveDqLocationFromTableAsset maps a DGC catalog Table asset to the DQ connection/dataSource/ +// schema/table that backs it. It walks the catalog hierarchy up from the table +// (Table -> Schema -> Database -> System) via incoming relations, then matches the DQ connection +// whose systemAssetId equals the table's System ancestor — the same mapping DQ uses for its own +// catalog asset links (so if asset links work from DQ, this resolves). +func ResolveDqLocationFromTableAsset(ctx context.Context, collibraHttpClient *http.Client, tableAssetID string) (*DqTableAssetLocation, error) { + tableAssetID = strings.TrimSpace(tableAssetID) + table, err := GetAssetWithRelations(ctx, collibraHttpClient, tableAssetID) + if err != nil { + return nil, err + } + if table.Type == nil || !strings.EqualFold(table.Type.Name, "Table") { + return nil, fmt.Errorf("asset %s is not a Table asset (type=%s)", tableAssetID, assetTypeName(table.Type)) + } + schemaRef := parentAssetOfType(table, "Schema") + if schemaRef == nil { + return nil, fmt.Errorf("table asset %s has no parent Schema in the catalog hierarchy", tableAssetID) + } + schema, err := GetAssetWithRelations(ctx, collibraHttpClient, schemaRef.ID) + if err != nil { + return nil, err + } + dbRef := parentAssetOfType(schema, "Database") + if dbRef == nil { + return nil, fmt.Errorf("schema asset %s has no parent Database in the catalog hierarchy", schemaRef.ID) + } + database, err := GetAssetWithRelations(ctx, collibraHttpClient, dbRef.ID) + if err != nil { + return nil, err + } + sysRef := parentAssetOfType(database, "System") + if sysRef == nil { + return nil, fmt.Errorf("database asset %s has no parent System in the catalog hierarchy", dbRef.ID) + } + + conns, err := ListDqConnections(ctx, collibraHttpClient) + if err != nil { + return nil, err + } + for i := range conns { + sysID := conns[i].SystemAssetID + if sysID == "" { + // The list response may omit systemAssetId; fetch the connection detail to fill it. + if full, e := GetDqConnection(ctx, collibraHttpClient, conns[i].ConnectionID); e == nil && full != nil { + sysID = full.SystemAssetID + } + } + if sysID != "" && strings.EqualFold(sysID, sysRef.ID) { + return &DqTableAssetLocation{ + TableAssetID: tableAssetID, + SystemAssetID: sysRef.ID, + ConnectionID: conns[i].ConnectionID, + ConnectionName: conns[i].ConnectionName, + EdgeSiteName: conns[i].EdgeSiteName, + DataSourceName: dbRef.DisplayName, + SchemaName: schemaRef.DisplayName, + TableName: table.DisplayName, + }, nil + } + } + return nil, fmt.Errorf("no DQ connection is mapped to the table's System asset %s — map the catalog system on the connection first", sysRef.ID) +} + +// parentAssetOfType returns the first incoming-relation source whose asset type matches typeName. +func parentAssetOfType(a *Asset, typeName string) *RelatedAsset { + if a == nil { + return nil + } + for i := range a.IncomingRelations { + s := a.IncomingRelations[i].Source + if s != nil && s.Type != nil && strings.EqualFold(s.Type.Name, typeName) { + return s + } + } + return nil +} + +func assetTypeName(t *AssetType) string { + if t == nil { + return "unknown" + } + return t.Name +} + +// TableAssetMatch is a catalog Table asset candidate from a by-name lookup. +type TableAssetMatch struct { + ID string + DisplayName string + FullName string + DomainName string +} + +// FindTableAssetsByName looks up catalog Table assets by exact signifier (displayName) via the +// public assets API (GET /rest/2.0/assets) — no search index required. Returns all matches so the +// caller can disambiguate. +func FindTableAssetsByName(ctx context.Context, collibraHttpClient *http.Client, name string, limit int) ([]TableAssetMatch, error) { + tableType, err := GetAssetTypeByPublicID(ctx, collibraHttpClient, "Table") + if err != nil { + return nil, fmt.Errorf("resolve Table asset type: %w", err) + } + if limit <= 0 { + limit = 50 + } + params := url.Values{} + params.Set("name", name) + params.Set("nameMatchMode", "EXACT") + params.Set("typeId", tableType.ID) + params.Set("limit", strconv.Itoa(limit)) + req, err := http.NewRequestWithContext(ctx, "GET", "/rest/2.0/assets?"+params.Encode(), nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Accept", "application/json") + body, err := executeRequest(collibraHttpClient, req) + if err != nil { + return nil, err + } + var resp struct { + Results []struct { + ID string `json:"id"` + DisplayName string `json:"displayName"` + Name string `json:"name"` + Domain struct { + Name string `json:"name"` + } `json:"domain"` + } `json:"results"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("failed to parse assets response: %w", err) + } + out := make([]TableAssetMatch, 0, len(resp.Results)) + for _, r := range resp.Results { + out = append(out, TableAssetMatch{ID: r.ID, DisplayName: r.DisplayName, FullName: r.Name, DomainName: r.Domain.Name}) + } + return out, nil +} + +func dqGet(ctx context.Context, collibraHttpClient *http.Client, endpoint string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + return executeRequest(collibraHttpClient, req) +} + +func dqPageQuery(values url.Values, limit, offset int) string { + if values == nil { + values = url.Values{} + } + values.Set("limit", strconv.Itoa(limit)) + values.Set("offset", strconv.Itoa(offset)) + return values.Encode() +} + +// UnsupportedWizardOptions lists the data-quality wizard configuration steps the +// create-DQ-job tools do NOT expose yet, each paired with the default the server +// applies. Shared by prepare_create_dq_job (to disclose proactively at the ready +// step) and create_dq_job (to disclose again at preview), so the user is never +// misled about scope. jobType selects the type-specific entry (Sizing for Pullup, +// Compute for Pushdown). +func UnsupportedWizardOptions(jobType string) []string { + // NOTE: schedule, time-slice, back-runs, column selection, monitor selection, adaptive monitor + // settings, row filtering, row sampling, notifications (incl. per-type messages), manual Spark + // sizing + Parallel JDBC (Pullup), and compute settings (Pushdown) ARE all supported now. The + // only no-code wizard option still not wired is the optional date-cast expression for a + // non-standard date column (used by the time-slice and row-filter date inputs). + _ = jobType // retained for API symmetry; the remaining gap is the same for both job types + return []string{ + "Date cast expression for a non-standard date column (time-slice / row-filter) — use a real date/timestamp column instead", + } +} diff --git a/pkg/clients/dgc_dq_client_test.go b/pkg/clients/dgc_dq_client_test.go new file mode 100644 index 0000000..9f242d3 --- /dev/null +++ b/pkg/clients/dgc_dq_client_test.go @@ -0,0 +1,129 @@ +package clients + +import "testing" + +func TestBuildProfileMonitorsMapsKeysToFields(t *testing.T) { + pm, unknown := BuildProfileMonitors([]string{"rowCount", "MIN", " uniqueness ", "descriptiveStatistics"}) + if len(unknown) != 0 { + t.Fatalf("expected no unknown keys, got %v", unknown) + } + // Case-insensitive + trimmed keys map to the right fields. + if !pm.RowCount || !pm.Min || !pm.Uniqueness || !pm.DescriptiveStatistics { + t.Errorf("selected monitors not all ON: %+v", pm) + } + // Everything not selected is OFF. + if pm.NullValues || pm.EmptyFields || pm.Mean || pm.Max || pm.ExecutionTime { + t.Errorf("unselected monitors should be OFF: %+v", pm) + } +} + +func TestBuildProfileMonitorsReportsUnknown(t *testing.T) { + _, unknown := BuildProfileMonitors([]string{"rowCount", "bogus"}) + if len(unknown) != 1 || unknown[0] != "bogus" { + t.Fatalf("expected unknown=[bogus], got %v", unknown) + } +} + +func TestDefaultMonitorKeysMatchCatalog(t *testing.T) { + defaults := map[string]bool{} + for _, k := range DefaultMonitorKeys() { + defaults[k] = true + } + for _, m := range DqMonitorCatalog() { + if m.DefaultEnabled != defaults[m.Key] { + t.Errorf("monitor %q: DefaultEnabled=%v but DefaultMonitorKeys membership=%v", m.Key, m.DefaultEnabled, defaults[m.Key]) + } + } +} + +func TestBuildSchedulingSettingsNeverIsNil(t *testing.T) { + for _, repeat := range []string{"", "NEVER", "never"} { + s, err := BuildSchedulingSettings(DqScheduleInput{Repeat: repeat}) + if err != nil || s != nil { + t.Errorf("repeat=%q: expected (nil,nil), got (%+v,%v)", repeat, s, err) + } + } +} + +func TestBuildSchedulingSettingsDailyDefaults(t *testing.T) { + s, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "DAILY"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s.SchedulerMode != "DAILY" || s.Daily == nil || len(s.Daily.DaysOfWeek) != 7 { + t.Errorf("expected DAILY with 7 days, got %+v", s) + } + if s.ScheduledRunTime != "00:00:00" || s.Daily.DailyOffset != "SCHEDULED" || !s.IsActive { + t.Errorf("unexpected defaults: %+v / daily %+v", s, s.Daily) + } +} + +func TestBuildSchedulingSettingsWeekdays(t *testing.T) { + s, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "WEEKDAYS", RunTime: "08:00:00"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s.SchedulerMode != "DAILY" || len(s.Daily.DaysOfWeek) != 5 { + t.Errorf("expected Mon-Fri (5 days), got %+v", s.Daily) + } +} + +func TestBuildSchedulingSettingsWeeklyInvalidDay(t *testing.T) { + if _, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "WEEKLY", DaysOfWeek: []string{"FUNDAY"}}); err == nil { + t.Errorf("expected error for invalid day FUNDAY") + } + if _, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "WEEKLY"}); err == nil { + t.Errorf("expected error for WEEKLY with no days") + } +} + +func TestBuildSchedulingSettingsHourly(t *testing.T) { + s, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "HOURLY", RunDateOffset: "TWO_HOURS"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s.SchedulerMode != "HOURLY" || s.Hourly == nil || s.Hourly.HourlyOffset != "TWO_HOURS" { + t.Errorf("expected HOURLY/TWO_HOURS, got %+v", s) + } + if _, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "HOURLY", RunDateOffset: "SEVEN_DAYS"}); err == nil { + t.Errorf("expected error: SEVEN_DAYS is not a valid hourly offset") + } +} + +func TestHasPermission(t *testing.T) { + perms := []string{"DATA_QUALITY_JOB_CREATE", "DATA_QUALITY_JOB_RUN"} + if !HasPermission(perms, "DATA_QUALITY_JOB_CREATE") { + t.Error("expected create permission present") + } + if !HasPermission(perms, "data_quality_job_run") { // case-insensitive + t.Error("expected run permission present (case-insensitive)") + } + if HasPermission(perms, "DATA_QUALITY_JOB_SCHEDULE") { + t.Error("did not expect schedule permission") + } +} + +func TestDqJobDetailsPath(t *testing.T) { + if got := DqJobDetailsPath("sales.orders"); got != "/data-quality/jobs?jobName=sales.orders" { + t.Errorf("unexpected path: %q", got) + } + if p := DqJobDetailsPath("my job"); p != "/data-quality/jobs?jobName=my+job" { + t.Errorf("expected query-escaped name, got %q", p) + } +} + +func TestBuildSchedulingSettingsMonthly(t *testing.T) { + s, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "MONTHLY", MonthlyMode: "LAST", RunDateOffset: "FIRST_OF_PRIOR_MONTH"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s.Monthly == nil || s.Monthly.MonthlyRepeat != "LAST" || s.Monthly.MonthlyOffset != "FIRST_OF_PRIOR_MONTH" { + t.Errorf("expected MONTHLY LAST/FIRST_OF_PRIOR_MONTH, got %+v", s.Monthly) + } + if _, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "MONTHLY", MonthlyMode: "DAY", DayOfMonth: 0}); err == nil { + t.Errorf("expected error: DAY mode needs dayOfMonth 1-28") + } + if _, err := BuildSchedulingSettings(DqScheduleInput{Repeat: "MONTHLY", MonthlyMode: "DAY", DayOfMonth: 31}); err == nil { + t.Errorf("expected error: dayOfMonth 31 is out of range") + } +} diff --git a/pkg/clients/dq_notifications.go b/pkg/clients/dq_notifications.go new file mode 100644 index 0000000..ecbce7b --- /dev/null +++ b/pkg/clients/dq_notifications.go @@ -0,0 +1,199 @@ +package clients + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +// DQ job notifications. Unlike rowFilter/sampleSetting (which are persisted-but-not-applied on the +// scan), the structured `notifications` field is APPLIED: the server's JobNotificationsMapper turns +// each NotificationOption into a real AlertCond (verified). So this is a clean structured feature — +// no SQL/dialect involvement. Contract: JobNotifications/NotificationOption in ui-v1-private-oas-spec.yaml. + +// DqNotificationOption is one notification rule (→ one AlertCond server-side). +type DqNotificationOption struct { + NotificationType string `json:"notificationType"` + Enabled bool `json:"enabled"` + Message string `json:"message,omitempty"` + Quantity int `json:"quantity,omitempty"` +} + +// DqJobNotifications is the public `notifications` object (dq-v1-public-oas-spec.yaml JobNotifications). +// Recipients are delivered via channels (currently EMAIL) carrying platform USERNAMES — not UUIDs. The +// public schema requires both notificationOptions and channels (>=1) when notifications are configured. +type DqJobNotifications struct { + NotificationOptions []DqNotificationOption `json:"notificationOptions"` + GlobalMessage string `json:"globalMessage,omitempty"` + UseIndividualMessages bool `json:"useIndividualMessages"` + Channels []DqNotificationChannel `json:"channels"` +} + +// DqNotificationChannel is one delivery channel and its recipients. Channel is EMAIL today; recipients +// are platform usernames. +type DqNotificationChannel struct { + Channel string `json:"channel"` + Recipients []string `json:"recipients"` +} + +// DqNotificationInfo describes one selectable notification for display/selection. +type DqNotificationInfo struct { + Key string `json:"key"` + Label string `json:"label"` + NotificationType string `json:"notificationType"` + DefaultEnabled bool `json:"defaultEnabled"` + TakesQuantity bool `json:"takesQuantity"` + DefaultQuantity int `json:"defaultQuantity,omitempty"` +} + +// DqNotificationCatalog is the selectable notifications with wizard defaults (Step Notifications): +// Job failed, Rows<=, Score<=, Run time> ON; Job completed, Runs/Days without data OFF. +func DqNotificationCatalog() []DqNotificationInfo { + return []DqNotificationInfo{ + {Key: "jobFailed", Label: "Job failed", NotificationType: "JOB_FAILED", DefaultEnabled: true}, + {Key: "rowsBelow", Label: "Rows below limit", NotificationType: "ROWS_LESS_THAN_LIMIT", DefaultEnabled: true, TakesQuantity: true, DefaultQuantity: 1}, + {Key: "scoreBelow", Label: "Score below limit", NotificationType: "SCORE_LESS_THAN_LIMIT", DefaultEnabled: true, TakesQuantity: true, DefaultQuantity: 75}, + {Key: "runTimeAbove", Label: "Run time above (minutes)", NotificationType: "RUN_TIME_MORE_THEN_LIMIT", DefaultEnabled: true, TakesQuantity: true, DefaultQuantity: 60}, + {Key: "jobCompleted", Label: "Job completed", NotificationType: "JOB_COMPLETED", DefaultEnabled: false}, + {Key: "runsWithoutData", Label: "Runs without data", NotificationType: "RUNS_WITHOUT_DATA", DefaultEnabled: false, TakesQuantity: true, DefaultQuantity: 1}, + {Key: "daysWithoutData", Label: "Days without data", NotificationType: "DAYS_WITHOUT_DATA", DefaultEnabled: false, TakesQuantity: true, DefaultQuantity: 1}, + } +} + +// NotificationKeys returns every valid notification key, in catalog order. +func NotificationKeys() []string { + cat := DqNotificationCatalog() + keys := make([]string, 0, len(cat)) + for _, n := range cat { + keys = append(keys, n.Key) + } + return keys +} + +// DefaultNotificationKeys returns the keys enabled by default. +func DefaultNotificationKeys() []string { + var keys []string + for _, n := range DqNotificationCatalog() { + if n.DefaultEnabled { + keys = append(keys, n.Key) + } + } + return keys +} + +// BuildNotificationOptions turns enabled keys (case-insensitive) into NotificationOptions. quantities +// overrides a key's threshold (keyed by lower-case key; <=0 uses the catalog default). messages sets a +// per-notification message (keyed by lower-case key) that overrides the global message for that one +// notification — pass nil for none. Unknown keys are returned so the caller can reject them. +func BuildNotificationOptions(enabledKeys []string, quantities map[string]int, messages map[string]string) ([]DqNotificationOption, []string) { + valid := map[string]DqNotificationInfo{} + for _, n := range DqNotificationCatalog() { + valid[strings.ToLower(n.Key)] = n + } + var opts []DqNotificationOption + var unknown []string + for _, k := range enabledKeys { + lk := strings.ToLower(strings.TrimSpace(k)) + if lk == "" { + continue + } + info, ok := valid[lk] + if !ok { + unknown = append(unknown, k) + continue + } + opt := DqNotificationOption{NotificationType: info.NotificationType, Enabled: true} + if info.TakesQuantity { + q := quantities[lk] + if q <= 0 { + q = info.DefaultQuantity + } + opt.Quantity = q + } + if m := strings.TrimSpace(messages[lk]); m != "" { + opt.Message = m + } + opts = append(opts, opt) + } + return opts, unknown +} + +// RecipientResolution is the outcome of resolving notification recipients to active users. UserIDs and +// Usernames are positionally aligned (same resolved user); Usernames feed the public notification +// channels (which take usernames), UserIDs are kept for callers that need the UUID. +type RecipientResolution struct { + UserIDs []string + Usernames []string + Unresolved []string +} + +// ResolveNotificationRecipients resolves each username/email to an active user's UUID. An entry +// containing '@' is looked up by email, otherwise by username. Not-found or disabled accounts land +// in Unresolved (the list endpoint excludes disabled users and the email lookup 404s), so the +// caller can warn and decide whether to proceed without them. Duplicates are de-duped. +func ResolveNotificationRecipients(ctx context.Context, client *http.Client, recipients []string) (RecipientResolution, error) { + var res RecipientResolution + seen := map[string]bool{} + for _, r := range recipients { + r = strings.TrimSpace(r) + if r == "" { + continue + } + var ( + u *EditAssetUser + err error + ) + if strings.Contains(r, "@") { + u, err = FindUserByEmail(ctx, client, r) + } else { + u, err = FindUserByUsername(ctx, client, r) + } + if err != nil { + return res, err + } + if u == nil || u.ID == "" { + res.Unresolved = append(res.Unresolved, r) + continue + } + if !seen[u.ID] { + seen[u.ID] = true + res.UserIDs = append(res.UserIDs, u.ID) + username := strings.TrimSpace(u.UserName) + if username == "" { + username = r // fall back to what the caller provided if the account has no username + } + res.Usernames = append(res.Usernames, username) + } + } + return res, nil +} + +// GetCurrentUser returns the invoking user (GET /rest/2.0/users/current) — the default notification +// recipient. +func GetCurrentUser(ctx context.Context, client *http.Client) (*EditAssetUser, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "/rest/2.0/users/current", nil) + if err != nil { + return nil, fmt.Errorf("get current user: building request: %w", err) + } + req.Header.Set("Accept", "application/json") + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("get current user: sending request: %w", err) + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("get current user: reading response: %w", err) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("get current user: status %d: %s", resp.StatusCode, string(body)) + } + var user EditAssetUser + if err := json.Unmarshal(body, &user); err != nil { + return nil, fmt.Errorf("get current user: decoding response: %w", err) + } + return &user, nil +} diff --git a/pkg/clients/dq_notifications_test.go b/pkg/clients/dq_notifications_test.go new file mode 100644 index 0000000..bc1815d --- /dev/null +++ b/pkg/clients/dq_notifications_test.go @@ -0,0 +1,67 @@ +package clients + +import "testing" + +func TestBuildNotificationOptionsDefaultsAndQuantities(t *testing.T) { + opts, unknown := BuildNotificationOptions( + []string{"jobFailed", "scoreBelow", "rowsBelow"}, + map[string]int{"rowsbelow": 500}, // override rowsBelow; scoreBelow uses its default + nil, + ) + if len(unknown) != 0 { + t.Fatalf("unexpected unknown keys: %v", unknown) + } + byType := map[string]DqNotificationOption{} + for _, o := range opts { + byType[o.NotificationType] = o + } + if o, ok := byType["JOB_FAILED"]; !ok || !o.Enabled || o.Quantity != 0 { + t.Errorf("jobFailed wrong: %+v", o) + } + if o := byType["SCORE_LESS_THAN_LIMIT"]; o.Quantity != 75 { + t.Errorf("expected scoreBelow default quantity 75, got %d", o.Quantity) + } + if o := byType["ROWS_LESS_THAN_LIMIT"]; o.Quantity != 500 { + t.Errorf("expected rowsBelow override quantity 500, got %d", o.Quantity) + } +} + +func TestBuildNotificationOptionsReportsUnknown(t *testing.T) { + _, unknown := BuildNotificationOptions([]string{"jobFailed", "bogus"}, nil, nil) + if len(unknown) != 1 || unknown[0] != "bogus" { + t.Fatalf("expected unknown=[bogus], got %v", unknown) + } +} + +func TestBuildNotificationOptionsPerTypeMessages(t *testing.T) { + opts, unknown := BuildNotificationOptions( + []string{"jobFailed", "scoreBelow"}, + nil, + map[string]string{"jobfailed": "ping me"}, // per-type message overrides global for jobFailed only + ) + if len(unknown) != 0 { + t.Fatalf("unexpected unknown keys: %v", unknown) + } + byType := map[string]DqNotificationOption{} + for _, o := range opts { + byType[o.NotificationType] = o + } + if o := byType["JOB_FAILED"]; o.Message != "ping me" { + t.Errorf("expected jobFailed message 'ping me', got %q", o.Message) + } + if o := byType["SCORE_LESS_THAN_LIMIT"]; o.Message != "" { + t.Errorf("expected scoreBelow to have no per-type message, got %q", o.Message) + } +} + +func TestDefaultNotificationKeysMatchCatalog(t *testing.T) { + defaults := map[string]bool{} + for _, k := range DefaultNotificationKeys() { + defaults[k] = true + } + for _, n := range DqNotificationCatalog() { + if n.DefaultEnabled != defaults[n.Key] { + t.Errorf("%s: DefaultEnabled=%v but membership=%v", n.Key, n.DefaultEnabled, defaults[n.Key]) + } + } +} diff --git a/pkg/clients/dq_source_query.go b/pkg/clients/dq_source_query.go new file mode 100644 index 0000000..7b6b14e --- /dev/null +++ b/pkg/clients/dq_source_query.go @@ -0,0 +1,299 @@ +package clients + +import ( + "fmt" + "strconv" + "strings" +) + +// ───────────────────────────────────────────────────────────────────────────── +// Dialect-aware DQ source-query builder — PORTED from the wizard. +// +// WHY THIS EXISTS / WHY IT'S A DUPLICATE: +// The DQ engine scans from the job's `sourceQuery` (opt_load.query). The structured +// rowFilter / sampleSetting fields on the create request are persisted to job_settings +// for UI re-display but are NOT compiled into the scan (verified live: a structured +// rowFilter never reaches opt_load.filter; sampleSetting never reaches opt_load.sample). +// So column-selection, row-filter, time-slice and sampling only take effect if they are +// written into `sourceQuery`. The ONLY complete implementation that composes all four, +// per database dialect, is the frontend wizard's `applyTimeSliceAndSampling`. +// +// This is a faithful Go port of that function (+ its helpers). It is a THIRD copy of the +// same logic and WILL drift unless maintained. Keep it in sync with — or, better, delete it +// once the server composes the query — these canonical sources: +// - FE: frontend/libs/bound-shared/data-quality/src/job-explorer/utils/utils.ts +// (applyTimeSliceAndSampling, createColumnsPart, createTimeSlicePart, createRowFilterPart) +// + .../utils/common.tsx (escapeColumnName) + constants/shared-constants.ts (SQL_ESCAPE_CHARACTERS) +// - BE (partial, sample-only): dq common-domain CommonJavaSqlUtils (getQueryWithSample/buildSampledQuery) +// +// The proper fix is server-side: have JobDefinitionMapperPublicV1 / getQueryWithSample compose +// filter+slice+sample from the structured fields, then this file goes away. +// ───────────────────────────────────────────────────────────────────────────── + +// rd/rdEnd substitution tokens. Kept single-quoted in the SQL (per the public OAS guidance and +// verified persisted form); the scheduler/engine replaces them per run. +const ( + dqRunIDToken = "'${rd}'" + dqRunIDEndToken = "'${rdEnd}'" +) + +// DqSourceQueryInput is everything needed to build the scan query. Filter is enabled when +// FilterColumn+FilterOperator are set; time-slice when TimeSliceColumn is set; sampling when +// SampleSize > 0. SelectedColumns empty => SELECT *. +type DqSourceQueryInput struct { + DatabaseProduct string // e.g. POSTGRES (connection.databaseProductName) + SchemaName string + TableName string + SelectedColumns []string + FilterColumn string + FilterOperator string + FilterValue string + TimeSliceColumn string + TimeSliceColumnCast string // optional explicit cast expression + TimeSliceColumnType string // optional source type (drives ATHENA/TRINO/ORACLE casts) + RunDateFormat string // "DATE" (default) | "TIMESTAMP" + SampleSize int +} + +// BuildDqSourceQuery returns the dialect-correct scan SQL, composing column selection, row +// filter, time-slice (${rd}/${rdEnd}) and sampling — mirroring the wizard exactly. +func BuildDqSourceQuery(in DqSourceQueryInput) string { + dialect := strings.ToUpper(strings.TrimSpace(in.DatabaseProduct)) + columns := dqColumnsPart(dialect, in.SelectedColumns) + from := dqFromPart(dialect, in.SchemaName, in.TableName) + where := dqTimeSliceAndRowFilterClause(in, dialect) + + whereClause := "" + if where != "" { + whereClause = "WHERE " + where + } + + if in.SampleSize <= 0 { + return dqJoin(fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), whereClause) + } + + size := strconv.Itoa(in.SampleSize) + // For the RANDOM()/RAND() forms the predicate is folded into the sampled WHERE and the + // total_rows count, so the sample is taken from the filtered+sliced subset (not the table). + whereWithPrependedSpace := "" + if where != "" { + whereWithPrependedSpace = " WHERE " + where + } + predicateAnd := "" + if where != "" { + predicateAnd = where + " AND" + } + + switch dialect { + case "DENODO", "SYBASE": + return dqJoin( + "SELECT * FROM (", + fmt.Sprintf(" SELECT %s, ROW_NUMBER() OVER (ORDER BY 1) AS rn_limit_alias", columns), + fmt.Sprintf(" FROM %s", from), + " "+whereClause, + ") AS temp_table", + fmt.Sprintf("WHERE rn_limit_alias <= %s", size), + ) + case "SNOWFLAKE": + return dqJoin(fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), + fmt.Sprintf("TABLESAMPLE (%s ROWS)", size), whereClause) + case "SAP", "SQLSERVER", "TERADATA", "AZURESYNAPSE": + return dqJoin(fmt.Sprintf("SELECT TOP %s %s", size, columns), fmt.Sprintf("FROM %s", from), whereClause) + case "ATHENA", "DATABRICKS": + return dqJoin( + fmt.Sprintf("WITH total_rows AS (SELECT COUNT(*) AS total FROM %s%s)", from, whereWithPrependedSpace), + fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), + fmt.Sprintf("WHERE %s RAND() < (CAST(%s AS DOUBLE) / CASE WHEN (SELECT total FROM total_rows) = 0 THEN 1 ELSE (SELECT total FROM total_rows) END) LIMIT %s", predicateAnd, size, size), + ) + case "BIGQUERY": + return dqJoin( + fmt.Sprintf("WITH total_rows AS (SELECT COUNT(*) AS total FROM %s%s)", from, whereWithPrependedSpace), + fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), + fmt.Sprintf("WHERE %s RAND() < (CAST(%s AS FLOAT64) / CASE WHEN (SELECT total FROM total_rows) = 0 THEN 1 ELSE (SELECT total FROM total_rows) END) LIMIT %s", predicateAnd, size, size), + ) + case "REDSHIFT": + return dqJoin( + fmt.Sprintf("WITH total_rows AS (SELECT COUNT(*) AS total FROM %s%s)", from, whereWithPrependedSpace), + fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), + fmt.Sprintf("WHERE %s RANDOM() < (CAST(%s AS FLOAT) / CASE WHEN (SELECT total FROM total_rows) = 0 THEN 1 ELSE (SELECT total FROM total_rows) END) LIMIT %s", predicateAnd, size, size), + ) + case "TRINO": + return dqJoin( + fmt.Sprintf("WITH total_rows AS (SELECT COUNT(*) AS total FROM %s%s)", from, whereWithPrependedSpace), + fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), + fmt.Sprintf("WHERE %s RANDOM() < (CAST(%s AS DOUBLE) / CASE WHEN (SELECT total FROM total_rows) = 0 THEN 1 ELSE (SELECT total FROM total_rows) END) limit %s", predicateAnd, size, size), + ) + case "POSTGRES": + return dqJoin( + fmt.Sprintf("WITH total_rows AS (SELECT COUNT(*) AS total FROM %s%s)", from, whereWithPrependedSpace), + fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), + fmt.Sprintf("WHERE %s RANDOM() < (%s::float / CASE WHEN (SELECT total FROM total_rows) = 0 THEN 1 ELSE (SELECT total FROM total_rows) END) LIMIT %s", predicateAnd, size, size), + ) + case "DB2": + return dqJoin(fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), whereClause, + fmt.Sprintf("FETCH FIRST %s ROWS ONLY", size)) + case "ORACLE": + return dqJoin(fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), whereClause, + "ORDER BY DBMS_RANDOM.VALUE", fmt.Sprintf("FETCH FIRST %s ROWS ONLY", size)) + case "MYSQL": + return dqJoin("SELECT *", fmt.Sprintf("FROM %s", from), whereClause, fmt.Sprintf("LIMIT %s", size)) + case "HIVE": + return dqJoin(fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), whereClause, fmt.Sprintf("LIMIT %s", size)) + default: + return dqJoin(fmt.Sprintf("SELECT %s", columns), fmt.Sprintf("FROM %s", from), whereClause, fmt.Sprintf("LIMIT %s", size)) + } +} + +// dqTimeSliceAndRowFilterClause builds the combined WHERE body (slice AND filter), each part dialect-aware. +func dqTimeSliceAndRowFilterClause(in DqSourceQueryInput, dialect string) string { + var slice, filter string + if strings.TrimSpace(in.TimeSliceColumn) != "" { + slice = dqTimeSlicePart(in, dialect) + } + if strings.TrimSpace(in.FilterColumn) != "" && strings.TrimSpace(in.FilterOperator) != "" { + filter = dqRowFilterPart(in, dialect) + } + switch { + case slice != "" && filter != "": + return slice + " AND " + filter + case slice != "": + return slice + default: + return filter + } +} + +// dqTimeSlicePart mirrors the wizard's createTimeSlicePart (per-dialect casts). +func dqTimeSlicePart(in DqSourceQueryInput, dialect string) string { + col := strings.TrimSpace(in.TimeSliceColumnCast) + if col == "" { + col = dqEscapeIdentifier(dialect, in.TimeSliceColumn) + } + dflt := fmt.Sprintf("%s >= %s AND %s < %s", col, dqRunIDToken, col, dqRunIDEndToken) + ts := strings.ToLower(strings.TrimSpace(in.TimeSliceColumnType)) + isTimestamp := strings.EqualFold(strings.TrimSpace(in.RunDateFormat), "TIMESTAMP") + + switch dialect { + case "ATHENA": + if !dqIsDateTimeType(ts) { + return dflt + } + castType := "DATE" + if isTimestamp { + castType = "TIMESTAMP" + } + return fmt.Sprintf("%s >= cast(%s as %s) AND %s < cast(%s as %s)", col, dqRunIDToken, castType, col, dqRunIDEndToken, castType) + case "TRINO": + if ts == "date" { + return fmt.Sprintf("%s >= CAST(%s AS DATE) AND %s < CAST(%s AS DATE)", col, dqRunIDToken, col, dqRunIDEndToken) + } + if ts == "timestamp" { + return fmt.Sprintf("%s >= CAST(%s AS TIMESTAMP) AND %s < CAST(%s AS TIMESTAMP)", col, dqRunIDToken, col, dqRunIDEndToken) + } + return fmt.Sprintf("%s >= TIMESTAMP %s AND %s < TIMESTAMP %s", col, dqRunIDToken, col, dqRunIDEndToken) + case "TERADATA": + if isTimestamp { + return fmt.Sprintf("%s >= TIMESTAMP %s AND %s < TIMESTAMP %s", col, dqRunIDToken, col, dqRunIDEndToken) + } + return dflt + case "ORACLE": + if !dqIsDateTimeType(ts) { + return dflt + } + if isTimestamp { + return fmt.Sprintf("%s >= TO_TIMESTAMP(%s, 'YYYY-MM-DD HH24:MI:SS') AND %s < TO_TIMESTAMP(%s, 'YYYY-MM-DD HH24:MI:SS')", col, dqRunIDToken, col, dqRunIDEndToken) + } + return fmt.Sprintf("%s >= TO_DATE(%s, 'YYYY-MM-DD HH24:MI:SS') AND %s < TO_DATE(%s, 'YYYY-MM-DD HH24:MI:SS')", col, dqRunIDToken, col, dqRunIDEndToken) + default: + return dflt + } +} + +// dqRowFilterPart mirrors the wizard's createRowFilterPart: `col op 'value'` (value single-quoted; +// BIGQUERY bare for numeric). For valueless operators (IS NULL / IS NOT NULL) the value is omitted. +func dqRowFilterPart(in DqSourceQueryInput, dialect string) string { + col := dqEscapeIdentifier(dialect, in.FilterColumn) + op := strings.TrimSpace(in.FilterOperator) + val := strings.TrimSpace(in.FilterValue) + // The value is meant to be BARE — the builder adds the surrounding single quotes (matching the + // wizard). Be forgiving if a caller already wrapped it, so 'US' doesn't become ''US'' (a syntax + // error). Only strips a matched outer pair; embedded quotes are left as-is. + if len(val) >= 2 && strings.HasPrefix(val, "'") && strings.HasSuffix(val, "'") { + val = val[1 : len(val)-1] + } + if val == "" { + // e.g. IS NULL / IS NOT NULL — no right-hand value. + return strings.TrimSpace(fmt.Sprintf("%s %s", col, op)) + } + if dialect == "BIGQUERY" { + if _, err := strconv.ParseFloat(val, 64); err == nil { + return fmt.Sprintf("%s %s %s", col, op, val) + } + } + return fmt.Sprintf("%s %s '%s'", col, op, val) +} + +// dqColumnsPart mirrors createColumnsPart: '*' when no subset, else the escaped selected columns. +func dqColumnsPart(dialect string, selected []string) string { + cols := make([]string, 0, len(selected)) + for _, c := range selected { + if c = strings.TrimSpace(c); c != "" { + cols = append(cols, dqEscapeIdentifier(dialect, c)) + } + } + if len(cols) == 0 { + return "*" + } + return strings.Join(cols, ", ") +} + +// dqFromPart mirrors createFromPart: escaped "schema"."table"; BIGQUERY wraps the whole ref in backticks. +func dqFromPart(dialect, schema, table string) string { + parts := make([]string, 0, 2) + if s := dqEscapeIdentifier(dialect, schema); strings.TrimSpace(schema) != "" { + parts = append(parts, s) + } + parts = append(parts, dqEscapeIdentifier(dialect, table)) + from := strings.Join(parts, ".") + if dialect == "BIGQUERY" { + return "`" + from + "`" + } + return from +} + +// dqEscapeIdentifier mirrors escapeColumnName / escapeTableOrSchemaName (enforced) per dialect. +func dqEscapeIdentifier(dialect, name string) string { + if name == "" { + return name + } + switch dialect { + case "SNOWFLAKE", "ORACLE", "ATHENA": // '"' quoting; SNOWFLAKE/ORACLE double embedded quotes + return `"` + strings.ReplaceAll(name, `"`, `""`) + `"` + case "POSTGRES", "REDSHIFT", "DENODO", "SAP": + return `"` + name + `"` + case "HIVE": + return "`" + strings.ReplaceAll(name, "`", "``") + "`" + case "DATABRICKS": + return "`" + name + "`" + case "SQLSERVER", "AZURESYNAPSE", "SYBASE": // bracket quoting; escape only the closing bracket + return "[" + strings.ReplaceAll(name, "]", "]]") + "]" + default: // MYSQL, BIGQUERY, DB2, TRINO, TERADATA, … — unescaped (matches the wizard's default) + return name + } +} + +func dqIsDateTimeType(t string) bool { + t = strings.ToLower(t) + return strings.Contains(t, "date") || strings.Contains(t, "time") +} + +// dqJoin joins non-empty SQL fragments with a single space and trims (whitespace is SQL-insignificant). +func dqJoin(parts ...string) string { + out := make([]string, 0, len(parts)) + for _, p := range parts { + if s := strings.TrimSpace(p); s != "" { + out = append(out, s) + } + } + return strings.Join(out, " ") +} diff --git a/pkg/clients/dq_source_query_test.go b/pkg/clients/dq_source_query_test.go new file mode 100644 index 0000000..75f077a --- /dev/null +++ b/pkg/clients/dq_source_query_test.go @@ -0,0 +1,93 @@ +package clients + +import ( + "strings" + "testing" +) + +func TestBuildDqSourceQueryPostgresComposesAll(t *testing.T) { + q := BuildDqSourceQuery(DqSourceQueryInput{ + DatabaseProduct: "POSTGRES", + SchemaName: "sales", + TableName: "customers", + SelectedColumns: []string{"id", "balance"}, + FilterColumn: "is_active", + FilterOperator: "=", + FilterValue: "true", + TimeSliceColumn: "created_at", + SampleSize: 137, + }) + for _, want := range []string{ + `"id", "balance"`, // column subset, escaped + `FROM "sales"."customers"`, // escaped relation + `"is_active" = 'true'`, // filter, value single-quoted + `"created_at" >= '${rd}' AND "created_at" < '${rdEnd}'`, // time slice + `RANDOM() < (137::float`, // postgres sampling + `LIMIT 137`, + } { + if !strings.Contains(q, want) { + t.Errorf("postgres query missing %q\n got: %s", want, q) + } + } + // The sample's total_rows count must include the filter+slice (sample of the filtered subset). + if !strings.Contains(q, `total_rows AS (SELECT COUNT(*) AS total FROM "sales"."customers" WHERE "created_at" >= '${rd}'`) { + t.Errorf("expected total_rows count to include the WHERE predicate, got: %s", q) + } +} + +func TestBuildDqSourceQueryDialectSampling(t *testing.T) { + cases := map[string]string{ + "SNOWFLAKE": "TABLESAMPLE (50 ROWS)", + "SQLSERVER": "SELECT TOP 50", + "DB2": "FETCH FIRST 50 ROWS ONLY", + "MYSQL": "LIMIT 50", + "ORACLE": "ORDER BY DBMS_RANDOM.VALUE", + } + for dialect, want := range cases { + q := BuildDqSourceQuery(DqSourceQueryInput{ + DatabaseProduct: dialect, SchemaName: "s", TableName: "t", SampleSize: 50, + }) + if !strings.Contains(q, want) { + t.Errorf("%s: expected %q in query, got: %s", dialect, want, q) + } + } +} + +func TestBuildDqSourceQueryEscapingAndBigQuery(t *testing.T) { + // SQL Server bracket-escapes identifiers. + q := BuildDqSourceQuery(DqSourceQueryInput{DatabaseProduct: "SQLSERVER", SchemaName: "dbo", TableName: "orders", SelectedColumns: []string{"id"}}) + if !strings.Contains(q, "[dbo].[orders]") || !strings.Contains(q, "[id]") { + t.Errorf("sqlserver escaping wrong: %s", q) + } + // BigQuery wraps the relation in backticks and leaves numeric filter values bare. + q = BuildDqSourceQuery(DqSourceQueryInput{ + DatabaseProduct: "BIGQUERY", SchemaName: "ds", TableName: "evt", + FilterColumn: "amount", FilterOperator: ">", FilterValue: "100", + }) + if !strings.Contains(q, "`ds.evt`") { + t.Errorf("bigquery relation should be backtick-wrapped: %s", q) + } + if !strings.Contains(q, "amount > 100") || strings.Contains(q, "amount > '100'") { + t.Errorf("bigquery numeric filter should be bare: %s", q) + } +} + +func TestBuildDqSourceQueryDoesNotDoubleQuote(t *testing.T) { + // A caller that pre-quotes the value must not produce ''US'' (the live failure). + for _, v := range []string{"US", "'US'"} { + q := BuildDqSourceQuery(DqSourceQueryInput{ + DatabaseProduct: "POSTGRES", SchemaName: "s", TableName: "t", + FilterColumn: "country", FilterOperator: "=", FilterValue: v, + }) + if !strings.Contains(q, `"country" = 'US'`) || strings.Contains(q, `''US''`) { + t.Errorf("value %q: expected single-quoted 'US', got: %s", v, q) + } + } +} + +func TestBuildDqSourceQueryPlainWhenNoOptions(t *testing.T) { + q := BuildDqSourceQuery(DqSourceQueryInput{DatabaseProduct: "POSTGRES", SchemaName: "sales", TableName: "customers"}) + if q != `SELECT * FROM "sales"."customers"` { + t.Errorf("unexpected plain query: %q", q) + } +} diff --git a/pkg/tools/create_dq_job/tool.go b/pkg/tools/create_dq_job/tool.go new file mode 100644 index 0000000..aead255 --- /dev/null +++ b/pkg/tools/create_dq_job/tool.go @@ -0,0 +1,915 @@ +// Package create_dq_job implements the create_dq_job MCP tool — it creates a +// Collibra data-quality job for a table and (per the wizard) queues a run. +// +// IT POSTS TO THE PUBLIC DQ API: POST /rest/dq/1.0/jobs (clients.CreateDqJob / +// clients.CreateDqJobRequest, contract dq/udq-app-client/oas/dq-v1-public-oas-spec.yaml). +// The public create accepts the full job definition — sourceQuery, runDate window, monitors, +// schedule, back-runs, notifications, and pullup/pushdown settings — so the internal BFF +// endpoint is no longer needed. (Discovery in prepare_create_dq_job still uses internal +// endpoints; the public DQ API exposes no connection/edge metadata browse.) +// +// HOW SCHEDULING + TIME-SLICE + ${rd} FIT TOGETHER: +// A scheduled job re-runs on a cadence; each run scans one date-column slice. The slice lives +// in the job SQL's WHERE as ` >= '${rd}' AND < '${rdEnd}'`. The public API has NO +// structured time-slice object — instead the engine substitutes ${rd}/${rdEnd} from runDate/ +// runDateEnd (formatted per jobSettings.dateFormat), and the scheduler adjusts the run date per +// run from the schedule's *Offset enums. So this tool writes the ${rd} predicate into sourceQuery +// whenever a timeSliceColumn is set and seeds runDate/runDateEnd to the most recent slice; without +// a timeSliceColumn every scheduled run rescans the whole table. +// +// It is built around a confirm checkpoint: confirm=false (default) returns a PREVIEW of +// the exact request without creating; confirm=true submits. +package create_dq_job + +import ( + "context" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "github.com/collibra/chip/pkg/chip" + "github.com/collibra/chip/pkg/clients" +) + +type Status string + +const ( + StatusPreview Status = "preview" + StatusCreated Status = "created" + StatusNeedsInput Status = "needs_input" + StatusError Status = "error" +) + +// Input — the five dataLocation fields are required (easiest via prepare_create_dq_job's +// `resolved` block). Time-slice / schedule / back-run fields are optional and structured. +type Input struct { + EdgeSiteName string `json:"edgeSiteName" jsonschema:"Edge site name. From prepare_create_dq_job resolved.edgeSiteName."` + EdgeConnectionName string `json:"edgeConnectionName" jsonschema:"Edge connection name, e.g. 'POSTGRES-SOURCE'. The tool resolves its databaseProductName + job type from the connection for the create request."` + DataSourceName string `json:"dataSourceName" jsonschema:"Data source / database, e.g. 'postgres'."` + SchemaName string `json:"schemaName" jsonschema:"Schema, e.g. 'sales'."` + TableName string `json:"tableName" jsonschema:"Table, e.g. 'transactions'."` + + JobType string `json:"jobType,omitempty" jsonschema:"PUSHDOWN or PULLUP. If omitted, resolved from the connection's capabilities."` + JobName string `json:"jobName,omitempty" jsonschema:"Job name to assign — offer the user the chance to set this. Defaults to '.'."` + + // --- Column selection. Omit to monitor ALL columns (the default). Provide a subset to + // scope profiling/monitoring to just those columns (each sent with selected=true). --- + SelectedColumns []string `json:"selectedColumns,omitempty" jsonschema:"Columns to monitor. Omit for ALL columns (default). Provide a subset — exact names from prepare_create_dq_job's columns — to profile only those. Does not change the source query (still SELECT *); only scopes what is profiled."` + + // --- Monitors. Omit to use the default set; provide an authoritative set of monitor keys + // to enable (anything not listed is turned off). descriptiveStatistics unmasks sensitive + // data — only include it after explicit user confirmation. --- + Monitors []string `json:"monitors,omitempty" jsonschema:"Monitor keys to enable — authoritative (anything omitted is turned OFF). Omit to use the defaults (rowCount, nullValues, emptyFields, uniqueness). Valid keys (see prepare_create_dq_job's monitors): rowCount, nullValues, emptyFields, uniqueness, min, mean, max, executionTime, descriptiveStatistics. descriptiveStatistics UNMASKS sensitive data — only include it after explicit user confirmation."` + + // --- Advanced monitor settings (adaptive behavior). Omit both to use the wizard defaults + // (lookback 10, learning phase 4). Setting either sends a structured `settings` object. --- + DataLookback int `json:"dataLookback,omitempty" jsonschema:"Adaptive monitors: number of prior runs used as the baseline. Default 10 when omitted. Set with learningPhase to override the adaptive 'Advanced monitor settings'."` + LearningPhase int `json:"learningPhase,omitempty" jsonschema:"Adaptive monitors: number of runs before adaptive monitors begin alerting. Default 4 when omitted."` + + // --- Row filter (the wizard's single-column predicate). Scopes the job to rows matching + // "filterColumn filterOperator filterValue" (e.g. amount > 100). Omit to scan all rows. + // This is ONE predicate — not compound AND/OR or free-form SQL (mirrors the DQ wizard). --- + FilterColumn string `json:"filterColumn,omitempty" jsonschema:"Column for a single-predicate row filter (the wizard's row filter). Use a name from prepare_create_dq_job's columns. Set filterColumn + filterOperator (+ filterValue) to scope the job to matching rows; omit to scan all rows."` + FilterOperator string `json:"filterOperator,omitempty" jsonschema:"Comparison operator for the row filter: = != <> > >= < <= (the wizard's set) or LIKE; the value is treated as a single literal. For IS NULL / IS NOT NULL leave filterValue empty. Required when filterColumn is set."` + FilterValue string `json:"filterValue,omitempty" jsonschema:"Right-hand value for the row filter, passed BARE — do NOT add quotes. The tool wraps it in single quotes for you (matching the DQ wizard), e.g. pass US (not 'US'), 100, 2024-01-01. Leave empty for valueless operators (IS NULL / IS NOT NULL). Applied by writing it into the job's source-query WHERE."` + + // --- Row sampling. Omit to scan all matching rows; a positive size enables sampling. The + // engine turns this into a dialect-correct RANDOM()/LIMIT query at dispatch. --- + SampleSize int `json:"sampleSize,omitempty" jsonschema:"Row sampling: number of rows to sample. Omit (or 0) to scan ALL matching rows (default). A positive value enables sampling at that size (the engine generates a RANDOM()/LIMIT query for it)."` + + // --- Time slice (incremental scanning). Pick a DATE COLUMN; when set, the tool writes a + // ">= '${rd}' AND < '${rdEnd}'" predicate into the job SQL's WHERE. The scheduler + // substitutes ${rd}/${rdEnd} per run, so each run scans only that slice, not the whole table. --- + TimeSliceColumn string `json:"timeSliceColumn,omitempty" jsonschema:"Date/timestamp column to slice on (e.g. 'txn_ts'). When set, the tool adds a WHERE \"\" >= '${rd}' AND \"\" < '${rdEnd}' predicate to the job SQL so each scheduled/backrun run scans only that slice. REQUIRED for incremental scheduling — without it, every scheduled run rescans the whole table. Pick a real date/timestamp column from prepare_create_dq_job's columns."` + TimeSliceSize int `json:"timeSliceSize,omitempty" jsonschema:"Time-slice window width (default 1 when timeSliceColumn set). Combined with timeSliceUnit, e.g. size=1 unit=DAYS = one day per run."` + TimeSliceUnit string `json:"timeSliceUnit,omitempty" jsonschema:"HOURS | DAYS | WEEKS | MONTHS (default DAYS)."` + + // --- Schedule (recurrence). Omit/NEVER = run once now. To slice incrementally per run, set a + // timeSliceColumn too (the schedule provides the cadence; the ${rd} window provides the slice). --- + ScheduleRepeat string `json:"scheduleRepeat,omitempty" jsonschema:"NEVER (default, run once now) | HOURLY | DAILY | WEEKLY | WEEKDAYS | MONTHLY. WEEKLY needs scheduleDaysOfWeek; MONTHLY uses scheduleDayOfMonth or scheduleMonthlyMode."` + ScheduleRunTime string `json:"scheduleRunTime,omitempty" jsonschema:"Scheduled run time, UTC HH:mm[:ss] (e.g. '00:01:00'). Defaults to 00:00:00 when a schedule is set."` + ScheduleDaysOfWeek []string `json:"scheduleDaysOfWeek,omitempty" jsonschema:"For WEEKLY: the days to run, e.g. ['MONDAY','THURSDAY'] (MONDAY..SUNDAY). Ignored for other modes (DAILY runs every day; WEEKDAYS runs Mon-Fri)."` + ScheduleDayOfMonth int `json:"scheduleDayOfMonth,omitempty" jsonschema:"For MONTHLY with scheduleMonthlyMode=DAY (the default): day of month 1-28."` + ScheduleMonthlyMode string `json:"scheduleMonthlyMode,omitempty" jsonschema:"For MONTHLY: DAY (default — run on scheduleDayOfMonth) | FIRST (first day of month) | LAST (last day of month)."` + RunDateOffset string `json:"runDateOffset,omitempty" jsonschema:"How far back the slice's run date (${rd}) is from the execution time. Default SCHEDULED (the run's own period). DAILY/WEEKLY/WEEKDAYS: SCHEDULED | ONE_DAY..SEVEN_DAYS. HOURLY: SCHEDULED | ONE_HOUR | TWO_HOURS. MONTHLY: SCHEDULED | FIRST_OF_CURRENT_MONTH | FIRST_OF_PRIOR_MONTH | LAST_OF_PRIOR_MONTH. Use a non-SCHEDULED offset so each run scans the prior, already-complete period."` + + // --- Back runs (historical backfill, walks forward one slice/run from runDate - binValue*timeBin). --- + BackrunEnabled bool `json:"backrunEnabled,omitempty" jsonschema:"Enable historical backfill runs."` + BackrunBinValue int `json:"backrunBinValue,omitempty" jsonschema:"Number of prior bins to backfill."` + BackrunTimeBin string `json:"backrunTimeBin,omitempty" jsonschema:"DAY | MONTH | YEAR."` + + // --- Notifications (optional). Configured when `notify` or `notifyRecipients` is set; otherwise + // no notifications. The invoking user is always the default recipient. --- + Notify []string `json:"notify,omitempty" jsonschema:"Notification keys to enable (authoritative; omit but set notifyRecipients to use the defaults jobFailed/rowsBelow/scoreBelow/runTimeAbove). Keys (see prepare_create_dq_job notifications): jobFailed, rowsBelow, scoreBelow, runTimeAbove, jobCompleted, runsWithoutData, daysWithoutData."` + NotifyRowsBelow int `json:"notifyRowsBelow,omitempty" jsonschema:"Threshold for rowsBelow — alert when row count <= this. Default 1."` + NotifyScoreBelow int `json:"notifyScoreBelow,omitempty" jsonschema:"Threshold for scoreBelow — alert when score (0-100) <= this. Default 75."` + NotifyRunTimeAboveMinutes int `json:"notifyRunTimeAboveMinutes,omitempty" jsonschema:"Threshold for runTimeAbove — alert when run time minutes > this. Default 60."` + NotifyRunsWithoutData int `json:"notifyRunsWithoutData,omitempty" jsonschema:"Threshold for runsWithoutData — alert when runs without data >= this. Default 1."` + NotifyDaysWithoutData int `json:"notifyDaysWithoutData,omitempty" jsonschema:"Threshold for daysWithoutData — alert when days without data >= this. Default 1."` + NotifyMessage string `json:"notifyMessage,omitempty" jsonschema:"Optional global message applied to the enabled notifications."` + NotifyRecipients []string `json:"notifyRecipients,omitempty" jsonschema:"Additional recipients by username or email (the invoking user is always included). Each is validated against active Collibra accounts; unresolved ones are reported."` + NotifyProceedWithoutUnresolved bool `json:"notifyProceedWithoutUnresolved,omitempty" jsonschema:"If some notifyRecipients can't be resolved to an active account, set true to create anyway with the resolvable recipients. Default false: the tool returns needs_input listing the unresolved ones so you can fix or confirm."` + + // --- Per-notification message overrides. Keyed by notification key (see prepare's notifications); + // a per-key message overrides notifyMessage for just that notification. --- + NotifyMessages map[string]string `json:"notifyMessages,omitempty" jsonschema:"Per-notification message overrides, keyed by notification key (jobFailed, rowsBelow, scoreBelow, runTimeAbove, jobCompleted, runsWithoutData, daysWithoutData). A per-key message overrides notifyMessage for just that notification (the wizard's individual-message mode)."` + + // --- Sizing (PULLUP only — the wizard's "Size Job Resources" step). Default = automatic + // ("Automate job resources"). Setting ANY manual field below switches to manual sizing (advanced; + // for large datasets). Memory fields are GB (e.g. "1", "2"). --- + SizingMaxExecutors int `json:"sizingMaxExecutors,omitempty" jsonschema:"PULLUP manual sizing: number of executors. Setting any sizing* field switches off automatic sizing. Default 1 when manual sizing is engaged."` + SizingExecutorCores int `json:"sizingExecutorCores,omitempty" jsonschema:"PULLUP manual sizing: cores per executor (maxExecutorCores). Default 1 when manual sizing is engaged."` + SizingExecutorMemoryGb string `json:"sizingExecutorMemoryGb,omitempty" jsonschema:"PULLUP manual sizing: memory per executor in GB, as a string (e.g. '1'). Default '1' when manual sizing is engaged."` + SizingDriverCores int `json:"sizingDriverCores,omitempty" jsonschema:"PULLUP manual sizing: driver cores. Default 1 when manual sizing is engaged."` + SizingDriverMemoryGb string `json:"sizingDriverMemoryGb,omitempty" jsonschema:"PULLUP manual sizing: driver memory in GB, as a string (e.g. '1'). Default '1' when manual sizing is engaged."` + SizingMemoryOverheadGb string `json:"sizingMemoryOverheadGb,omitempty" jsonschema:"PULLUP manual sizing: memory overhead in GB, as a string (e.g. '1'). Default '1' when manual sizing is engaged."` + SizingNumPartitions int `json:"sizingNumPartitions,omitempty" jsonschema:"PULLUP: number of partitions (load options). Omit/0 lets Spark decide."` + + // --- Parallel JDBC (PULLUP advanced sizing). mode AUTO (column+count auto) | AUTO_COLUMN (column + // auto, count required) | MANUAL (column+count required). Mode is inferred when omitted: a + // partition column implies MANUAL, a partition count alone implies AUTO_COLUMN. --- + ParallelJdbcMode string `json:"parallelJdbcMode,omitempty" jsonschema:"PULLUP Parallel JDBC mode: AUTO | AUTO_COLUMN | MANUAL. AUTO needs nothing else; AUTO_COLUMN requires parallelJdbcPartitionNumber; MANUAL requires both parallelJdbcPartitionColumn and parallelJdbcPartitionNumber. Omit to infer from the other fields."` + ParallelJdbcPartitionColumn string `json:"parallelJdbcPartitionColumn,omitempty" jsonschema:"PULLUP Parallel JDBC: the column to partition on (MANUAL mode only). Setting a specific column requires a manual parallelJdbcPartitionNumber (auto-calculate is not allowed)."` + ParallelJdbcPartitionNumber int `json:"parallelJdbcPartitionNumber,omitempty" jsonschema:"PULLUP Parallel JDBC: number of partitions. Required for AUTO_COLUMN and MANUAL modes."` + + SparkSqlProperties map[string]string `json:"sparkSqlProperties,omitempty" jsonschema:"PULLUP: additional Spark SQL key/value properties to set on the job."` + + // --- Compute settings (PUSHDOWN only — the wizard's Review step). Defaults connections 10 + // (range 1-50), threads 2 (range 1-10). --- + PushdownConnections int `json:"pushdownConnections,omitempty" jsonschema:"PUSHDOWN compute: number of connections (1-50). Default 10 when omitted."` + PushdownThreads int `json:"pushdownThreads,omitempty" jsonschema:"PUSHDOWN compute: number of threads (1-10). Default 2 when omitted."` + + // --- Catalog linkage (optional). When set (e.g. from prepare_create_dq_job's resolved table + // asset), the success result includes a deep link to the catalog Table asset. --- + TableAssetID string `json:"tableAssetId,omitempty" jsonschema:"Catalog Table asset UUID this job monitors (from prepare_create_dq_job). When set, the success result includes a catalog deep link to the asset."` + + // AcknowledgeDescriptiveStatistics must be true to enable the descriptiveStatistics monitor — it + // UNMASKS sensitive values. Without it the tool refuses to proceed (the wizard's explicit-confirm). + AcknowledgeDescriptiveStatistics bool `json:"acknowledgeDescriptiveStatistics,omitempty" jsonschema:"Required true when monitors includes descriptiveStatistics — confirms you accept that it UNMASKS sensitive values. Without it the tool refuses to proceed."` + + Confirm bool `json:"confirm,omitempty" jsonschema:"Safety checkpoint. false (default) returns a PREVIEW of the exact request WITHOUT creating it — review with the user. true submits."` +} + +type Output struct { + Status Status `json:"status" jsonschema:"preview | created | needs_input | error."` + Message string `json:"message" jsonschema:"Human-readable outcome."` + Request *clients.CreateDqJobRequest `json:"request,omitempty" jsonschema:"The exact public-API payload (POST /rest/dq/1.0/jobs) that will be / was submitted. Review on preview."` + JobName string `json:"jobName,omitempty" jsonschema:"Final job name."` + JobID string `json:"jobId,omitempty" jsonschema:"Created job id (the internal create returns jobId, not a run id)."` + JobType string `json:"jobType,omitempty" jsonschema:"Resolved job type."` + JobDetailsLink string `json:"jobDetailsLink,omitempty" jsonschema:"On success: the Job Details deep-link path, relative to the Collibra instance URL (e.g. /data-quality/jobs?jobName=...)."` + TableAssetLink string `json:"tableAssetLink,omitempty" jsonschema:"On success: catalog deep-link path to the related Table asset (when tableAssetId was provided), relative to the instance URL (/asset/)."` + Warnings []string `json:"warnings,omitempty" jsonschema:"Non-fatal warnings to surface to the user (e.g. missing Schedule/Run permission, dropped schedule)."` + AffectedStep string `json:"affectedStep,omitempty" jsonschema:"On a submission error: the configuration step the error maps to, so you can return the user there and re-call with the other inputs preserved."` + UnsupportedOptions []string `json:"unsupportedOptions,omitempty" jsonschema:"Wizard options this tool still does NOT set (server defaults apply). Surface to the user."` + Guidance string `json:"guidance,omitempty" jsonschema:"On needs_input/error, what to do next."` +} + +func NewTool(collibraClient *http.Client) *chip.Tool[Input, Output] { + return &chip.Tool[Input, Output]{ + Name: "create_dq_job", + Title: "Create Data Quality Job", + Description: "Creates a Collibra data-quality job for a table (and queues a run). Provide the table location " + + "(edgeSiteName, edgeConnectionName, dataSourceName, schemaName, tableName — easiest from prepare_create_dq_job's " + + "`resolved`). Optional structured config: column selection (selectedColumns — subset to monitor; omit for all), monitors (monitors — " + + "set of monitor keys to enable; omit for defaults) with adaptive settings (dataLookback/learningPhase), row sampling (sampleSize — sample N " + + "rows; omit for all), row filter (filterColumn/filterOperator/filterValue — single-column predicate like amount > 100, inlined into the " + + "source-query WHERE), time slice (timeSliceColumn/Size/Unit — writes a ${rd}/${rdEnd} WHERE so each run scans one slice), schedule " + + "(scheduleRepeat/RunTime/DaysOfWeek/DayOfMonth/MonthlyMode + runDateOffset — pair with a timeSliceColumn for incremental runs), " + + "back runs (backrun*), notifications (notify keys + thresholds + notifyRecipients — the invoking user is always a recipient). jobType resolves from the connection if omitted. " + + "Posts to the PUBLIC DQ API (POST /rest/dq/1.0/jobs). Defaults to a PREVIEW: " + + "confirm=false returns the exact payload without creating; call again with confirm=true after the user approves.", + Handler: handler(collibraClient), + Permissions: []string{}, + } +} + +func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { + return func(ctx context.Context, input Input) (Output, error) { + if missing := missingLocationFields(input); len(missing) > 0 { + return Output{ + Status: StatusNeedsInput, + Message: fmt.Sprintf("Missing required field(s): %s.", strings.Join(missing, ", ")), + Guidance: "Call prepare_create_dq_job to resolve connection/data source/schema/table, then pass its `resolved` block here.", + }, nil + } + + // A row filter needs at least a column and an operator (value may be empty for + // valueless operators like IS NULL). Catch a half-specified filter before any API call. + if hasFilterFields(input) { + if strings.TrimSpace(input.FilterColumn) == "" { + return Output{Status: StatusNeedsInput, Message: "filterOperator/filterValue given without filterColumn.", Guidance: "Set filterColumn (a column from prepare_create_dq_job) to apply a row filter, or omit all filter fields."}, nil + } + if strings.TrimSpace(input.FilterOperator) == "" { + return Output{Status: StatusNeedsInput, Message: "filterColumn set but filterOperator is missing.", Guidance: "Provide filterOperator (e.g. =, >, <, LIKE, IN, IS NULL), plus filterValue unless the operator takes none."}, nil + } + } + + // Reject unknown monitor keys before any API call. + if len(input.Monitors) > 0 { + if _, unknown := clients.BuildProfileMonitors(input.Monitors); len(unknown) > 0 { + return Output{ + Status: StatusNeedsInput, + Message: fmt.Sprintf("Unknown monitor(s): %s.", strings.Join(unknown, ", ")), + Guidance: "Use monitor keys from prepare_create_dq_job's `monitors`: " + strings.Join(clients.MonitorKeys(), ", ") + ".", + }, nil + } + } + + // Descriptive statistics UNMASKS sensitive data — require explicit acknowledgement (the + // wizard's "explicit confirmation"). Hard-gate before preview or submit. + if monitorsIncludeDescriptiveStats(input.Monitors) && !input.AcknowledgeDescriptiveStatistics { + return Output{ + Status: StatusNeedsInput, + AffectedStep: stepMonitors, + Message: "Enabling descriptive statistics may expose sensitive values if they are present in the columns included in the scan.", + Guidance: "To proceed with descriptiveStatistics, set acknowledgeDescriptiveStatistics=true; otherwise remove it from monitors.", + }, nil + } + + // Validate + build the schedule (nil = run once now) before any API call. + schedule, schedErr := clients.BuildSchedulingSettings(clients.DqScheduleInput{ + Repeat: input.ScheduleRepeat, + RunTime: input.ScheduleRunTime, + DaysOfWeek: input.ScheduleDaysOfWeek, + DayOfMonth: input.ScheduleDayOfMonth, + MonthlyMode: input.ScheduleMonthlyMode, + RunDateOffset: input.RunDateOffset, + }) + if schedErr != nil { + return Output{ + Status: StatusNeedsInput, + Message: schedErr.Error(), + Guidance: "Fix the schedule input: scheduleRepeat (NEVER/HOURLY/DAILY/WEEKLY/WEEKDAYS/MONTHLY), scheduleRunTime, scheduleDaysOfWeek (WEEKLY), scheduleDayOfMonth/scheduleMonthlyMode (MONTHLY), runDateOffset.", + }, nil + } + + // Notifications: configured only when the caller engages them (notify keys, recipients, or a + // global message). The invoking user is always the default recipient. Unlike rowFilter/sample, + // this structured field IS applied by the server (→ alert conditions). + var notifications *clients.DqJobNotifications + if len(input.Notify) > 0 || len(input.NotifyRecipients) > 0 || strings.TrimSpace(input.NotifyMessage) != "" { + keys := input.Notify + if len(keys) == 0 { + keys = clients.DefaultNotificationKeys() + } + notifyMessages := map[string]string{} + for k, v := range input.NotifyMessages { + notifyMessages[strings.ToLower(strings.TrimSpace(k))] = v + } + opts, unknown := clients.BuildNotificationOptions(keys, map[string]int{ + "rowsbelow": input.NotifyRowsBelow, + "scorebelow": input.NotifyScoreBelow, + "runtimeabove": input.NotifyRunTimeAboveMinutes, + "runswithoutdata": input.NotifyRunsWithoutData, + "dayswithoutdata": input.NotifyDaysWithoutData, + }, notifyMessages) + if len(unknown) > 0 { + return Output{ + Status: StatusNeedsInput, + Message: fmt.Sprintf("Unknown notification(s): %s.", strings.Join(unknown, ", ")), + Guidance: "Use notification keys from prepare_create_dq_job's `notifications`: " + strings.Join(clients.NotificationKeys(), ", ") + ".", + }, nil + } + // Recipients: invoking user (always) + any additional usernames/emails, de-duped. The public + // notification channel takes platform USERNAMES, not UUIDs (no UUID resolution needed). + var recipients []string + seenUser := map[string]bool{} + addUser := func(username string) { + username = strings.TrimSpace(username) + if username != "" && !seenUser[username] { + seenUser[username] = true + recipients = append(recipients, username) + } + } + if cu, err := clients.GetCurrentUser(ctx, collibraClient); err == nil && cu != nil { + addUser(cu.UserName) + } + res, err := clients.ResolveNotificationRecipients(ctx, collibraClient, input.NotifyRecipients) + if err != nil { + return Output{Status: StatusError, Message: fmt.Sprintf("Failed to resolve notification recipients: %v", err), Guidance: "Check the recipient usernames/emails and retry."}, nil + } + if len(res.Unresolved) > 0 && !input.NotifyProceedWithoutUnresolved { + return Output{ + Status: StatusNeedsInput, + Message: fmt.Sprintf("These notification recipients have no active Collibra account: %s.", strings.Join(res.Unresolved, ", ")), + Guidance: "Fix the username/email, or set notifyProceedWithoutUnresolved=true to create anyway with the valid recipients (the unresolved ones are dropped).", + }, nil + } + for _, username := range res.Usernames { + addUser(username) + } + useIndividual := false + for _, o := range opts { + if o.Message != "" { + useIndividual = true + break + } + } + notifications = &clients.DqJobNotifications{ + NotificationOptions: opts, + GlobalMessage: strings.TrimSpace(input.NotifyMessage), + UseIndividualMessages: useIndividual, + Channels: []clients.DqNotificationChannel{{Channel: "EMAIL", Recipients: recipients}}, + } + } + + // Resolve connectionId + edgeSiteId (+ jobType fallback) from the connection — the + // internal request needs the IDs, not just names. findConnectionByName hits the + // INTERNAL connections list (clients.ListDqConnections); the public API has no + // equivalent that returns capabilityTypes/edgeSiteId. + conn, connErr := findConnectionByName(ctx, collibraClient, input.EdgeConnectionName) + jobType := strings.ToUpper(strings.TrimSpace(input.JobType)) + if jobType == "" && connErr == nil && conn != nil { + switch len(conn.CapabilityTypes) { + case 1: + jobType = conn.CapabilityTypes[0] + case 0: + return Output{Status: StatusNeedsInput, Message: fmt.Sprintf("Connection %q advertises no DQ capability.", input.EdgeConnectionName), Guidance: "Confirm a Data Quality Pushdown/Pullup capability is enabled, or set jobType."}, nil + default: + return Output{Status: StatusNeedsInput, Message: fmt.Sprintf("Connection %q supports multiple job types (%s).", input.EdgeConnectionName, strings.Join(conn.CapabilityTypes, ", ")), Guidance: "Set jobType explicitly."}, nil + } + } + if jobType == "" { + return Output{Status: StatusNeedsInput, Message: "Could not determine job type.", Guidance: "Set jobType to PUSHDOWN or PULLUP, or call prepare_create_dq_job."}, nil + } + if connErr != nil || conn == nil { + return Output{Status: StatusNeedsInput, Message: fmt.Sprintf("Could not resolve connection %q (needed for connectionId/edgeSiteId).", input.EdgeConnectionName), Guidance: "Verify the connection name via prepare_create_dq_job."}, nil + } + + // Job name: default from the server's collision-free generator (auto-increments, e.g. "..._2"); + // validate a user-provided name against the server's rules (rejects special chars). + jobName := strings.TrimSpace(input.JobName) + if jobName == "" { + if generated, err := clients.GenerateUniqueJobName(ctx, collibraClient, input.SchemaName, input.TableName); err == nil && strings.TrimSpace(generated) != "" { + jobName = generated + } else { + jobName = input.SchemaName + "." + input.TableName + } + } else if ok, err := clients.IsValidDqJobName(ctx, collibraClient, jobName); err == nil && !ok { + guidance := "Use only letters, numbers, '-' and '_'." + if suggestion, sErr := clients.GenerateUniqueJobName(ctx, collibraClient, input.SchemaName, input.TableName); sErr == nil && strings.TrimSpace(suggestion) != "" { + guidance += fmt.Sprintf(" A valid default is %q.", suggestion) + } + return Output{Status: StatusNeedsInput, AffectedStep: stepSelectData, Message: fmt.Sprintf("Job name %q is invalid (special characters other than - and _ are not allowed).", jobName), Guidance: guidance}, nil + } + + // Type-specific config must match the job type (sizing/Parallel JDBC are PULLUP-only; compute is PUSHDOWN-only). + if err := validateTypeSpecificConfig(input, jobType); err != nil { + return Output{Status: StatusNeedsInput, AffectedStep: stepForTypeConfigError(jobType), Message: err.Error(), Guidance: "Provide settings that match the resolved job type, or omit them."}, nil + } + + // Build the type-specific settings (manual sizing / Parallel JDBC for PULLUP, compute for PUSHDOWN). + jobSettings, settingsErr := buildJobSettings(input, jobType) + if settingsErr != nil { + return Output{Status: StatusNeedsInput, AffectedStep: stepForTypeConfigError(jobType), Message: settingsErr.Error(), Guidance: "Fix the sizing / Parallel JDBC inputs (see parallelJdbcMode rules), then retry."}, nil + } + + // Permission preflight: DQ create/run/schedule are CONNECTION-resource permissions (granted by + // the Data Quality Editor/Manager roles), checked as `global || resource` like the DQ UI. Hard-gate + // on Create; degrade gracefully on Schedule (drop it) and warn on Run. Best-effort: if the + // permission lookup itself fails, proceed with a warning and let the server enforce — never block + // on a failed lookup. + var warnings []string + if global, resource, permErr := clients.GetDqConnectionPermissions(ctx, collibraClient, conn.ConnectionID); permErr == nil { + has := func(p string) bool { return clients.HasPermission(global, p) || clients.HasPermission(resource, p) } + manageAll := has(clients.PermResourceManageAll) + if !manageAll && !has(clients.PermDqJobCreate) { + return Output{ + Status: StatusError, + JobType: jobType, + Message: fmt.Sprintf("You do not have the Data Quality Job > Create permission (DATA_QUALITY_JOB_CREATE) on connection %q required to create a DQ job.", conn.ConnectionName), + Guidance: "Ask an administrator to grant you the Data Quality Editor or Data Quality Manager role on this connection, then retry.", + }, nil + } + if schedule != nil && !manageAll && !has(clients.PermDqJobSchedule) { + warnings = append(warnings, "You lack the Data Quality Job > Schedule permission on this connection — the schedule was dropped; the job will run once on demand.") + schedule = nil + } + if !manageAll && !has(clients.PermDqJobRun) { + warnings = append(warnings, "You lack the Data Quality Job > Run permission on this connection. This tool's create endpoint always queues a run (no save-without-run via the internal API), so the run may be rejected server-side.") + } + } else { + warnings = append(warnings, fmt.Sprintf("Could not verify your DQ permissions (%v) — proceeding; the server will enforce them.", permErr)) + } + + req := buildPublicRequest(input, conn, jobType, jobName, jobSettings, schedule, notifications) + + if !input.Confirm { + columnsDesc := "all" + if len(input.SelectedColumns) > 0 { + columnsDesc = fmt.Sprintf("%d selected", len(input.SelectedColumns)) + } + filterDesc := "none" + if strings.TrimSpace(input.FilterColumn) != "" { + filterDesc = strings.TrimSpace(fmt.Sprintf("%s %s %s", input.FilterColumn, input.FilterOperator, input.FilterValue)) + } + monitorKeys := input.Monitors + if len(monitorKeys) == 0 { + monitorKeys = clients.DefaultMonitorKeys() + } + pm, _ := clients.BuildProfileMonitors(monitorKeys) + enabledMonitors := clients.EnabledMonitorKeys(pm) + am := req.MonitoringSettings.AdaptiveMonitors + sensitiveWarning := "" + if am != nil && am.DescriptiveStatistics { + sensitiveWarning = " WARNING: descriptive statistics is ON — this UNMASKS sensitive data in profiling output; confirm the user explicitly accepts this." + } + hasSlice := strings.TrimSpace(input.TimeSliceColumn) != "" + rdNote := "" + if hasSlice { + rdNote = " The source query uses ${rd}/${rdEnd} on the time-slice column (see request.sourceQuery) so each run scans one slice; review it." + } else if req.SchedulingSettings != nil { + rdNote = " NOTE: no timeSliceColumn set, so every scheduled run rescans the WHOLE table (no ${rd} slice). Set timeSliceColumn for incremental runs." + } + sampleDesc := "all rows" + if input.SampleSize > 0 { + sampleDesc = fmt.Sprintf("sample %d", input.SampleSize) + } + adaptiveDesc := "defaults (lookback 10, learning 4)" + if am != nil && am.Settings != nil { + adaptiveDesc = fmt.Sprintf("lookback %d, learning %d", am.Settings.DataLookBack, am.Settings.LearningPhase) + } + notifyDesc := "none" + if req.Notifications != nil { + notifyDesc = fmt.Sprintf("%d alert(s) -> %d recipient(s)", len(req.Notifications.NotificationOptions), countRecipients(req.Notifications)) + } + warnNote := "" + if len(warnings) > 0 { + warnNote = " WARNINGS: " + strings.Join(warnings, " ") + } + return Output{ + Status: StatusPreview, + JobType: jobType, + JobName: jobName, + Request: req, + Warnings: warnings, + UnsupportedOptions: clients.UnsupportedWizardOptions(jobType), + Message: fmt.Sprintf("Preview only — nothing created. Will create a %s job %q for %s.%s (columns=%s, rows=%s, filter=%q, monitors=[%s], adaptive=%s, timeSlice=%v, schedule=%s, backrun=%v, notifications=%s, %s). "+ + "Review with the user (job name is overridable via jobName), then call again with confirm=true.%s%s%s", + jobType, jobName, input.SchemaName, input.TableName, columnsDesc, sampleDesc, filterDesc, strings.Join(enabledMonitors, ", "), adaptiveDesc, hasSlice, describeSchedule(req.SchedulingSettings), req.Backrun != nil, notifyDesc, describeCompute(req.JobSettings, jobType), sensitiveWarning, rdNote, warnNote), + }, nil + } + + // confirm=true -> submit to the PUBLIC endpoint POST /rest/dq/1.0/jobs. + resp, err := clients.CreateDqJob(ctx, collibraClient, *req) + if err != nil { + return Output{ + Status: StatusError, + JobType: jobType, + Request: req, + Warnings: warnings, + AffectedStep: affectedStepForError(err), + Message: fmt.Sprintf("Create failed: %v", err), + Guidance: "A validation error (400) from the public API names the offending field; affectedStep points to the most likely step — return the user there, fix it, and re-call preserving the other inputs.", + }, nil + } + finalName := jobName + if strings.TrimSpace(resp.JobName) != "" { + finalName = resp.JobName + } + tableAssetLink := "" + if id := strings.TrimSpace(input.TableAssetID); id != "" { + tableAssetLink = clients.CatalogAssetPath(id) + } + msg := fmt.Sprintf("Created %s job %q (jobRunId %s). Job Details: %s", jobType, finalName, resp.JobRunID, clients.DqJobDetailsPath(finalName)) + if tableAssetLink != "" { + msg += " | Catalog asset: " + tableAssetLink + } + msg += " (links are relative to your Collibra instance URL)." + return Output{ + Status: StatusCreated, + JobName: finalName, + JobID: resp.JobRunID, + JobType: jobType, + JobDetailsLink: clients.DqJobDetailsPath(finalName), + TableAssetLink: tableAssetLink, + Warnings: warnings, + UnsupportedOptions: clients.UnsupportedWizardOptions(jobType), + Message: msg, + }, nil + } +} + +// buildPublicRequest assembles the public JobDefinitionCreateRequest (POST /rest/dq/1.0/jobs). Column +// selection, row filter, sampling and the time-slice ${rd}/${rdEnd} predicate are all composed INTO +// sourceQuery (the public API has no structured fields for them; sourceQuery drives the scan). When a +// time slice is set, the initial runDate/runDateEnd window is seeded to the most recent slice so the +// queued run has concrete ${rd}/${rdEnd} values. Monitors map onto monitoringSettings.adaptiveMonitors. +func buildPublicRequest(input Input, conn *clients.DqConnection, jobType, jobName string, jobSettings *clients.DqPublicJobSettings, schedule *clients.DqSchedulingSettings, notifications *clients.DqJobNotifications) *clients.CreateDqJobRequest { + // Monitors: caller's authoritative set, or the wizard defaults when omitted. Unknown keys are + // already rejected in the handler. Mapped onto the public adaptiveMonitors shape. + monitorKeys := input.Monitors + if len(monitorKeys) == 0 { + monitorKeys = clients.DefaultMonitorKeys() + } + profileMonitors, _ := clients.BuildProfileMonitors(monitorKeys) + adaptive := clients.PublicAdaptiveMonitorsFromProfile(profileMonitors) + + // Adaptive monitor settings: only send `settings` if the caller customized one of them; + // default the unset one to the wizard default (lookback 10, learning phase 4). + if input.DataLookback > 0 || input.LearningPhase > 0 { + lookback := input.DataLookback + if lookback <= 0 { + lookback = 10 + } + learning := input.LearningPhase + if learning <= 0 { + learning = 4 + } + adaptive.Settings = &clients.DqPublicAdaptiveMonitorSettings{DataLookBack: lookback, LearningPhase: learning} + } + + // Build the scan query with the dialect-aware builder (clients.BuildDqSourceQuery). The engine + // scans from sourceQuery, so column selection, row filter, time-slice (${rd}/${rdEnd}) and sampling + // are all composed INTO it, per the source dialect. + tsCol := strings.TrimSpace(input.TimeSliceColumn) + sourceQuery := clients.BuildDqSourceQuery(clients.DqSourceQueryInput{ + DatabaseProduct: conn.DatabaseProductName, + SchemaName: input.SchemaName, + TableName: input.TableName, + SelectedColumns: input.SelectedColumns, + FilterColumn: strings.TrimSpace(input.FilterColumn), + FilterOperator: strings.TrimSpace(input.FilterOperator), + FilterValue: input.FilterValue, + TimeSliceColumn: tsCol, + RunDateFormat: "DATE", + SampleSize: input.SampleSize, + }) + + // Run-date window: when slicing, seed runDate/runDateEnd to the most recent slice and set the + // matching dateFormat so the immediate queued run has concrete ${rd}/${rdEnd}. Otherwise stamp + // runDate = today (the server defaults it when omitted, but seeding it matches prior behavior). + var runDate, runDateEnd *clients.DqPublicRunDate + if tsCol != "" { + tsUnit := strings.ToUpper(strings.TrimSpace(input.TimeSliceUnit)) + if tsUnit == "" { + tsUnit = "DAYS" + } + tsSize := input.TimeSliceSize + if tsSize <= 0 { + tsSize = 1 + } + start, end, kind, dateFormat := sliceWindowPublic(tsSize, tsUnit) + runDate = &clients.DqPublicRunDate{Kind: kind, Value: start} + runDateEnd = &clients.DqPublicRunDate{Kind: kind, Value: end} + if jobSettings != nil { + jobSettings.DateFormat = dateFormat + } + } else { + runDate = &clients.DqPublicRunDate{Kind: "DATE", Value: time.Now().UTC().Format("2006-01-02")} + } + + req := &clients.CreateDqJobRequest{ + JobType: jobType, + JobName: jobName, + DataLocation: clients.DqDataLocation{ + EdgeSiteName: input.EdgeSiteName, + EdgeConnectionName: conn.ConnectionName, + DataSourceName: input.DataSourceName, + SchemaName: input.SchemaName, + TableName: input.TableName, + DatabaseProductName: conn.DatabaseProductName, + }, + SourceQuery: sourceQuery, + RunDate: runDate, + RunDateEnd: runDateEnd, + JobSettings: jobSettings, + MonitoringSettings: &clients.DqPublicMonitoringSettings{AdaptiveMonitors: adaptive}, + Notifications: notifications, + SchedulingSettings: schedule, + } + + if input.BackrunEnabled { + bin := strings.ToUpper(strings.TrimSpace(input.BackrunTimeBin)) + if bin == "" { + bin = "DAY" + } + binValue := input.BackrunBinValue + if binValue < 1 { + binValue = 1 // the public Backrun requires binValue >= 1 + } + req.Backrun = &clients.DqPublicBackrun{TimeBin: bin, BinValue: binValue} + } + + return req +} + +// sliceWindowPublic returns the initial [start, end) run window for the immediate queued run plus the +// matching runDate kind and jobSettings dateFormat: the most recent slice of width size*unit ending +// now. HOURS uses an RFC3339 TIMESTAMP; the rest use a DATE (yyyy-MM-dd). +func sliceWindowPublic(size int, unit string) (start, end, kind, dateFormat string) { + now := time.Now().UTC() + switch strings.ToUpper(unit) { + case "HOURS": + const f = "2006-01-02T15:04:05Z" + return now.Add(time.Duration(-size) * time.Hour).Format(f), now.Format(f), "TIMESTAMP", "TIMESTAMP" + case "WEEKS": + return now.AddDate(0, 0, -7*size).Format("2006-01-02"), now.Format("2006-01-02"), "DATE", "DATE" + case "MONTHS": + return now.AddDate(0, -size, 0).Format("2006-01-02"), now.Format("2006-01-02"), "DATE", "DATE" + default: // DAYS + return now.AddDate(0, 0, -size).Format("2006-01-02"), now.Format("2006-01-02"), "DATE", "DATE" + } +} + +// countRecipients totals the recipients across a notifications object's channels. +func countRecipients(n *clients.DqJobNotifications) int { + total := 0 + for _, c := range n.Channels { + total += len(c.Recipients) + } + return total +} + +// describeSchedule renders a one-line summary of the resolved schedule for the preview. +func describeSchedule(s *clients.DqSchedulingSettings) string { + if s == nil { + return "run once now" + } + switch s.SchedulerMode { + case "HOURLY": + return fmt.Sprintf("HOURLY (offset %s)", s.Hourly.HourlyOffset) + case "DAILY": + days := "every day" + if s.Daily != nil && len(s.Daily.DaysOfWeek) > 0 && len(s.Daily.DaysOfWeek) < 7 { + days = strings.Join(s.Daily.DaysOfWeek, ",") + } + return fmt.Sprintf("DAILY [%s] @ %s UTC (offset %s)", days, s.ScheduledRunTime, s.Daily.DailyOffset) + case "MONTHLY": + when := s.Monthly.MonthlyRepeat + if s.Monthly.MonthlyRepeat == "DAY" { + when = fmt.Sprintf("day %d", s.Monthly.DayNumber) + } + return fmt.Sprintf("MONTHLY (%s) @ %s UTC (offset %s)", when, s.ScheduledRunTime, s.Monthly.MonthlyOffset) + } + return s.SchedulerMode +} + +// buildJobSettings populates the type-specific public jobSettings. PUSHDOWN carries compute settings +// (connections/threads). PULLUP carries loadOptions (+ optional Parallel JDBC) and, for manual sizing, +// sparkJobSizing — omitting sparkJobSizing entirely tells the public API to size automatically (it has +// no autoSizing flag). Returns an error on out-of-range compute, invalid Parallel JDBC, or non-integer +// memory GB. +func buildJobSettings(input Input, jobType string) (*clients.DqPublicJobSettings, error) { + if strings.EqualFold(jobType, "PUSHDOWN") { + connections := orDefaultInt(input.PushdownConnections, 10) + threads := orDefaultInt(input.PushdownThreads, 2) + if connections < 1 || connections > 50 { + return nil, fmt.Errorf("pushdownConnections must be between 1 and 50 (got %d)", input.PushdownConnections) + } + if threads < 1 || threads > 10 { + return nil, fmt.Errorf("pushdownThreads must be between 1 and 10 (got %d)", input.PushdownThreads) + } + return &clients.DqPublicJobSettings{PushdownSettings: &clients.DqPublicPushdownSettings{Connections: connections, Threads: threads}}, nil + } + // PULLUP. + loadOptions := &clients.DqPublicLoadOptions{NumPartitions: input.SizingNumPartitions} // 0 = let Spark decide + pj, err := buildParallelJdbc(input) + if err != nil { + return nil, err + } + loadOptions.ParallelJdbcOptions = pj + sizing, err := buildSparkJobSizing(input) + if err != nil { + return nil, err + } + return &clients.DqPublicJobSettings{PullupSettings: &clients.DqPublicPullupSettings{ + LoadOptions: loadOptions, + SparkJobSizing: sizing, // nil => automatic sizing + SparkSqlProperties: input.SparkSqlProperties, + }}, nil +} + +// buildSparkJobSizing returns nil for automatic sizing (the public API auto-sizes when sparkJobSizing +// is omitted) unless the caller set any manual sizing field, in which case manual sizing is used with +// wizard defaults (1) for the unset fields. Memory GB strings are parsed to integers (SparkMemoryGB). +func buildSparkJobSizing(input Input) (*clients.DqPublicSparkJobSizing, error) { + if !hasManualSizing(input) { + return nil, nil + } + execMem, err := gbToInt(input.SizingExecutorMemoryGb, 1) + if err != nil { + return nil, fmt.Errorf("sizingExecutorMemoryGb %w", err) + } + driverMem, err := gbToInt(input.SizingDriverMemoryGb, 1) + if err != nil { + return nil, fmt.Errorf("sizingDriverMemoryGb %w", err) + } + overhead, err := gbToInt(input.SizingMemoryOverheadGb, 1) + if err != nil { + return nil, fmt.Errorf("sizingMemoryOverheadGb %w", err) + } + return &clients.DqPublicSparkJobSizing{ + NumExecutors: orDefaultInt(input.SizingMaxExecutors, 1), + NumExecutorCores: orDefaultInt(input.SizingExecutorCores, 1), + ExecutorMemoryGb: execMem, + DriverCores: orDefaultInt(input.SizingDriverCores, 1), + DriverMemoryGb: driverMem, + MemoryOverheadGb: overhead, + }, nil +} + +// gbToInt parses a GB string to a positive integer (the public SparkMemoryGB is an integer; fractional +// values are not supported). An empty string uses def. +func gbToInt(s string, def int) (int, error) { + s = strings.TrimSpace(s) + if s == "" { + return def, nil + } + v, err := strconv.Atoi(s) + if err != nil { + return 0, fmt.Errorf("must be an integer number of GB, got %q", s) + } + if v < 1 { + return 0, fmt.Errorf("must be >= 1 GB, got %q", s) + } + return v, nil +} + +func hasManualSizing(input Input) bool { + return input.SizingMaxExecutors > 0 || input.SizingExecutorCores > 0 || strings.TrimSpace(input.SizingExecutorMemoryGb) != "" || + input.SizingDriverCores > 0 || strings.TrimSpace(input.SizingDriverMemoryGb) != "" || strings.TrimSpace(input.SizingMemoryOverheadGb) != "" +} + +// buildParallelJdbc maps the caller's Parallel JDBC inputs onto the ParallelJdbcOptions contract, +// inferring the mode when omitted and enforcing the wizard's rules: AUTO needs nothing; AUTO_COLUMN +// requires a partition count (and no column); MANUAL requires both a column and a count (a specific +// partition column disables auto-calculate). Returns (nil, nil) when Parallel JDBC isn't configured. +func buildParallelJdbc(input Input) (*clients.DqParallelJdbcOptions, error) { + mode := strings.ToUpper(strings.TrimSpace(input.ParallelJdbcMode)) + col := strings.TrimSpace(input.ParallelJdbcPartitionColumn) + num := input.ParallelJdbcPartitionNumber + if mode == "" { + switch { + case col == "" && num == 0: + return nil, nil // not configured + case col != "": + mode = clients.ParallelJdbcManual + default: + mode = clients.ParallelJdbcAutoColumn + } + } + switch mode { + case clients.ParallelJdbcAuto: + return &clients.DqParallelJdbcOptions{Mode: clients.ParallelJdbcAuto}, nil + case clients.ParallelJdbcAutoColumn: + if col != "" { + return nil, fmt.Errorf("parallelJdbcMode AUTO_COLUMN auto-selects the partition column — omit parallelJdbcPartitionColumn, or use MANUAL to choose one") + } + if num <= 0 { + return nil, fmt.Errorf("parallelJdbcMode AUTO_COLUMN requires parallelJdbcPartitionNumber (auto-calculate is off once a partition count is set)") + } + return &clients.DqParallelJdbcOptions{Mode: clients.ParallelJdbcAutoColumn, PartitionNumber: num}, nil + case clients.ParallelJdbcManual: + if col == "" { + return nil, fmt.Errorf("parallelJdbcMode MANUAL requires parallelJdbcPartitionColumn") + } + if num <= 0 { + return nil, fmt.Errorf("parallelJdbcMode MANUAL requires a manual parallelJdbcPartitionNumber (auto-calculate is not allowed with a specific partition column)") + } + return &clients.DqParallelJdbcOptions{Mode: clients.ParallelJdbcManual, PartitionColumn: col, PartitionNumber: num}, nil + default: + return nil, fmt.Errorf("parallelJdbcMode %q is invalid (use AUTO, AUTO_COLUMN, or MANUAL)", mode) + } +} + +func orDefaultInt(v, def int) int { + if v > 0 { + return v + } + return def +} + +// validateTypeSpecificConfig rejects sizing/Parallel JDBC inputs on a PUSHDOWN job and compute inputs +// on a PULLUP job, so the caller can't silently set options that don't apply to the resolved type. +func validateTypeSpecificConfig(input Input, jobType string) error { + if strings.EqualFold(jobType, "PUSHDOWN") { + if hasManualSizing(input) || strings.TrimSpace(input.ParallelJdbcMode) != "" || + strings.TrimSpace(input.ParallelJdbcPartitionColumn) != "" || input.ParallelJdbcPartitionNumber > 0 || + input.SizingNumPartitions > 0 || len(input.SparkSqlProperties) > 0 { + return fmt.Errorf("sizing / Parallel JDBC settings apply to PULLUP jobs only, but this is a PUSHDOWN job") + } + return nil + } + if input.PushdownConnections > 0 || input.PushdownThreads > 0 { + return fmt.Errorf("compute settings (connections/threads) apply to PUSHDOWN jobs only, but this is a PULLUP job") + } + return nil +} + +// Conversational step labels — used for affectedStep so the caller can route the user back to the +// failing step and re-call with the other inputs preserved. +const ( + stepSelectData = "Select the data (Step 1)" + stepMonitors = "Add monitors (Step 2)" + stepSizing = "Size job resources (Pullup)" + stepSchedule = "Set a run schedule" + stepNotifications = "Add notifications" + stepReview = "Review and run" +) + +// stepForTypeConfigError points type-specific config errors at the right step (compute lives in the +// Pushdown Review step; sizing is the Pullup Sizing step). +func stepForTypeConfigError(jobType string) string { + if strings.EqualFold(jobType, "PUSHDOWN") { + return stepReview + } + return stepSizing +} + +// affectedStepForError maps a submission error message to the most likely configuration step. +func affectedStepForError(err error) string { + if err == nil { + return "" + } + msg := strings.ToLower(err.Error()) + switch { + case strings.Contains(msg, "sizing") || strings.Contains(msg, "executor") || strings.Contains(msg, "partition") || strings.Contains(msg, "spark"): + return stepSizing + case strings.Contains(msg, "schedul"): + return stepSchedule + case strings.Contains(msg, "notif") || strings.Contains(msg, "recipient") || strings.Contains(msg, "alert"): + return stepNotifications + case strings.Contains(msg, "monitor") || strings.Contains(msg, "profile"): + return stepMonitors + case strings.Contains(msg, "column") || strings.Contains(msg, "filter") || strings.Contains(msg, "sample") || strings.Contains(msg, "query") || strings.Contains(msg, "name"): + return stepSelectData + default: + return stepReview + } +} + +// describeCompute renders a one-line summary of the type-specific settings for the preview. For PULLUP, +// a nil sparkJobSizing means automatic sizing (the public API auto-sizes when it is omitted). +func describeCompute(js *clients.DqPublicJobSettings, jobType string) string { + if js == nil { + return "compute=defaults" + } + if strings.EqualFold(jobType, "PUSHDOWN") { + if js.PushdownSettings != nil { + return fmt.Sprintf("compute=connections %d/threads %d", js.PushdownSettings.Connections, js.PushdownSettings.Threads) + } + return "compute=defaults" + } + if js.PullupSettings == nil { + return "sizing=auto" + } + sizing := "auto" + if s := js.PullupSettings.SparkJobSizing; s != nil { + sizing = fmt.Sprintf("manual(executors %d, cores %d, execMem %dGB, driverMem %dGB)", s.NumExecutors, s.NumExecutorCores, s.ExecutorMemoryGb, s.DriverMemoryGb) + } + pj := "off" + if lo := js.PullupSettings.LoadOptions; lo != nil && lo.ParallelJdbcOptions != nil { + pj = lo.ParallelJdbcOptions.Mode + } + return fmt.Sprintf("sizing=%s, parallelJdbc=%s", sizing, pj) +} + +// monitorsIncludeDescriptiveStats reports whether the caller asked to enable descriptiveStatistics. +func monitorsIncludeDescriptiveStats(keys []string) bool { + for _, k := range keys { + if strings.EqualFold(strings.TrimSpace(k), "descriptiveStatistics") { + return true + } + } + return false +} + +// hasFilterFields reports whether the caller supplied any row-filter input at all. +func hasFilterFields(in Input) bool { + return strings.TrimSpace(in.FilterColumn) != "" || + strings.TrimSpace(in.FilterOperator) != "" || + strings.TrimSpace(in.FilterValue) != "" +} + +func missingLocationFields(in Input) []string { + var missing []string + for name, val := range map[string]string{ + "edgeSiteName": in.EdgeSiteName, + "edgeConnectionName": in.EdgeConnectionName, + "dataSourceName": in.DataSourceName, + "schemaName": in.SchemaName, + "tableName": in.TableName, + } { + if strings.TrimSpace(val) == "" { + missing = append(missing, name) + } + } + return missing +} + +func findConnectionByName(ctx context.Context, client *http.Client, name string) (*clients.DqConnection, error) { + conns, err := clients.ListDqConnections(ctx, client) + if err != nil { + return nil, err + } + for i := range conns { + if strings.EqualFold(conns[i].ConnectionName, name) { + return &conns[i], nil + } + } + return nil, fmt.Errorf("connection %q not found", name) +} diff --git a/pkg/tools/create_dq_job/tool_test.go b/pkg/tools/create_dq_job/tool_test.go new file mode 100644 index 0000000..8ba74e7 --- /dev/null +++ b/pkg/tools/create_dq_job/tool_test.go @@ -0,0 +1,1012 @@ +package create_dq_job_test + +import ( + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + + tools "github.com/collibra/chip/pkg/tools/create_dq_job" + "github.com/collibra/chip/pkg/tools/testutil" +) + +// muxWithPerms returns the base connections mux plus the GraphQL permissions endpoint, reporting the +// given permission identifiers (e.g. DATA_QUALITY_JOB_CREATE) on the CONNECTION resource — the scope +// the DQ UI checks. Returned under `resource` to exercise the resource-aware preflight. +func muxWithPerms(perms ...string) *http.ServeMux { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + mux.HandleFunc("/graphql", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + quoted := make([]string, 0, len(perms)) + for _, p := range perms { + quoted = append(quoted, strconv.Quote(p)) + } + _, _ = w.Write([]byte(`{"data":{"api":{"currentUser":{"global":[],"resource":[` + strings.Join(quoted, ",") + `]}}}}`)) + }) + return mux +} + +// connectionsHandler mocks the INTERNAL connections list create_dq_job uses to resolve the connection +// name + databaseProductName (and jobType from capabilities) for the public create request. +func connectionsHandler() http.Handler { + return testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + return http.StatusOK, map[string]any{ + "results": []map[string]any{{ + "connectionId": "conn-1", + "connectionName": "POSTGRES-SOURCE", + "capabilityTypes": []string{"PULLUP"}, + "databaseProductName": "POSTGRES", + "edgeSiteId": "site-1", + "edgeSiteName": "EDGE-1", + }}, + } + }) +} + +// createdJobHandler mocks the PUBLIC create POST /rest/dq/1.0/jobs, echoing a job name + queued run id +// (JobDefinitionCreateResponse). +func createdJobHandler() http.Handler { + return testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + if r.Method != http.MethodPost { + return http.StatusMethodNotAllowed, map[string]any{} + } + return http.StatusCreated, map[string]any{"jobName": "sales.transactions", "jobType": "PULLUP", "jobRunId": "run-1"} + }) +} + +func baseInput() tools.Input { + return tools.Input{ + EdgeSiteName: "EDGE-1", + EdgeConnectionName: "POSTGRES-SOURCE", + DataSourceName: "postgres", + SchemaName: "sales", + TableName: "transactions", + JobType: "PULLUP", + } +} + +func TestPreviewDoesNotCreate(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + mux.HandleFunc("/rest/dq/1.0/jobs", func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("create endpoint must NOT be called during preview") + }) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusPreview { + t.Fatalf("expected preview, got %q (%s)", out.Status, out.Message) + } + if out.Request == nil { + t.Fatalf("expected a preview request") + } + // The public body uses dataLocation with NAMES (no connectionId/edgeSiteId). + dl := out.Request.DataLocation + if dl.EdgeConnectionName != "POSTGRES-SOURCE" || dl.EdgeSiteName != "EDGE-1" || dl.DatabaseProductName != "POSTGRES" { + t.Errorf("dataLocation not resolved into request: %+v", dl) + } + if out.Request.JobType != "PULLUP" || out.Request.JobName != "sales.transactions" { + t.Errorf("unexpected request: jobType=%s jobName=%s", out.Request.JobType, out.Request.JobName) + } + // PULLUP defaults to automatic sizing: pullupSettings present, sparkJobSizing omitted (nil). + if out.Request.JobSettings == nil || out.Request.JobSettings.PullupSettings == nil { + t.Fatalf("PULLUP request missing jobSettings.pullupSettings: %+v", out.Request.JobSettings) + } + if out.Request.JobSettings.PullupSettings.SparkJobSizing != nil { + t.Errorf("expected automatic sizing (nil sparkJobSizing) by default, got %+v", out.Request.JobSettings.PullupSettings.SparkJobSizing) + } + if out.JobID != "" { + t.Errorf("expected no jobId on preview, got %q", out.JobID) + } + // No selected columns / filter / sample -> a plain whole-table scan. + q := out.Request.SourceQuery + if !strings.HasPrefix(q, "SELECT * FROM") || strings.Contains(q, "RANDOM()") || strings.Contains(q, "${rd}") { + t.Errorf("expected a plain SELECT * scan, got %q", q) + } + if len(out.UnsupportedOptions) == 0 { + t.Errorf("preview should disclose still-unsupported wizard options") + } +} + +func TestRowFilterInlinedInQuery(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.FilterColumn = "amount" + in.FilterOperator = ">" + in.FilterValue = "100" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // The filter must be INLINED into the source query WHERE (value single-quoted, per the wizard). + if !strings.Contains(out.Request.SourceQuery, `"amount" > '100'`) { + t.Errorf("expected filter inlined into sourceQuery WHERE, got %q", out.Request.SourceQuery) + } +} + +func TestRowFilterAndSliceCombineInQuery(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.FilterColumn = "is_active" + in.FilterOperator = "=" + in.FilterValue = "true" + in.TimeSliceColumn = "txn_ts" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + q := out.Request.SourceQuery + // Filter AND slice predicates must both be present, ANDed together (filter value single-quoted). + if !strings.Contains(q, `"is_active" = 'true'`) { + t.Errorf("filter predicate missing from query: %q", q) + } + if !strings.Contains(q, `"txn_ts" >= '${rd}'`) || !strings.Contains(q, `"txn_ts" < '${rdEnd}'`) { + t.Errorf("slice predicate missing from query: %q", q) + } + if !strings.Contains(q, " AND ") { + t.Errorf("expected predicates ANDed together: %q", q) + } +} + +func TestRowFilterRequiresOperator(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.FilterColumn = "amount" // operator omitted -> incomplete filter + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input when filterColumn set without operator, got %q (%s)", out.Status, out.Message) + } +} + +func TestMonitorsDefaultWhenOmitted(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + am := out.Request.MonitoringSettings.AdaptiveMonitors + // Defaults: rowCount/nullValues/emptyFields/uniqueness ON; everything else OFF. + if !am.RowCount || !am.NullValues || !am.EmptyFields || !am.Uniqueness { + t.Errorf("expected the four default monitors ON, got %+v", am) + } + if am.Min || am.Mean || am.Max || am.ExecutionTime || am.DescriptiveStatistics { + t.Errorf("expected non-default monitors OFF, got %+v", am) + } +} + +func TestMonitorsCustomSelection(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Monitors = []string{"rowCount", "min"} // authoritative -> only these ON + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + am := out.Request.MonitoringSettings.AdaptiveMonitors + if !am.RowCount || !am.Min { + t.Errorf("expected rowCount+min ON, got %+v", am) + } + if am.NullValues || am.EmptyFields || am.Uniqueness || am.Mean || am.Max { + t.Errorf("expected unselected default monitors turned OFF, got %+v", am) + } +} + +func TestUnknownMonitorNeedsInput(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Monitors = []string{"rowCount", "bogusMonitor"} + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for unknown monitor, got %q (%s)", out.Status, out.Message) + } +} + +func TestDescriptiveStatisticsWarnsInPreview(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Monitors = []string{"rowCount", "descriptiveStatistics"} + in.AcknowledgeDescriptiveStatistics = true // required by the hard gate + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusPreview { + t.Fatalf("expected preview, got %q", out.Status) + } + if !out.Request.MonitoringSettings.AdaptiveMonitors.DescriptiveStatistics { + t.Errorf("expected descriptiveStatistics ON in request") + } + if !strings.Contains(out.Message, "UNMASKS") { + t.Errorf("expected a sensitive-data warning in the preview message, got %q", out.Message) + } +} + +func TestDescriptiveStatisticsRequiresAcknowledgement(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Monitors = []string{"rowCount", "descriptiveStatistics"} // no acknowledgement + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input (hard gate) for descriptiveStatistics without acknowledgement, got %q (%s)", out.Status, out.Message) + } + if !strings.Contains(out.Message, "sensitive") { + t.Errorf("expected the sensitive-data warning, got %q", out.Message) + } +} + +func TestSuccessReturnsTableAssetLink(t *testing.T) { + mux := muxWithPerms("DATA_QUALITY_JOB_CREATE", "DATA_QUALITY_JOB_SCHEDULE", "DATA_QUALITY_JOB_RUN") + mux.Handle("/rest/dq/1.0/jobs", createdJobHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Confirm = true + in.TableAssetID = "019f062a-18dc-7489-a373-c928e59f1fc4" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusCreated { + t.Fatalf("expected created, got %q (%s)", out.Status, out.Message) + } + if out.TableAssetLink != "/asset/019f062a-18dc-7489-a373-c928e59f1fc4" { + t.Errorf("expected catalog asset link, got %q", out.TableAssetLink) + } +} + +func TestTimeSliceGeneratesRdQuery(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.TimeSliceColumn = "txn_ts" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + q := out.Request.SourceQuery + // The slice predicate must be in the WHERE for the scheduler to substitute it per run. + if !strings.Contains(q, `"txn_ts" >= '${rd}'`) || !strings.Contains(q, `"txn_ts" < '${rdEnd}'`) { + t.Errorf("expected ${rd}/${rdEnd} predicate on the slice column, got query %q", q) + } + // An initial run-date window must be set so the immediate queued run has concrete ${rd}/${rdEnd}. + if out.Request.RunDate == nil || out.Request.RunDateEnd == nil { + t.Fatalf("expected runDate+runDateEnd window for the initial run, got runDate=%+v runDateEnd=%+v", out.Request.RunDate, out.Request.RunDateEnd) + } + if out.Request.RunDate.Kind != "DATE" || out.Request.RunDateEnd.Kind != "DATE" { + t.Errorf("expected DATE-kind run dates for a DAYS slice, got %+v / %+v", out.Request.RunDate, out.Request.RunDateEnd) + } + // dateFormat must match the run-date kind. + if out.Request.JobSettings.DateFormat != "DATE" { + t.Errorf("expected jobSettings.dateFormat=DATE, got %q", out.Request.JobSettings.DateFormat) + } +} + +func TestScheduleDailyPopulatesSubObject(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.ScheduleRepeat = "DAILY" + in.ScheduleRunTime = "01:30" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + s := out.Request.SchedulingSettings + if s == nil || s.SchedulerMode != "DAILY" || s.Daily == nil { + t.Fatalf("expected a DAILY schedule with a daily sub-object, got %+v", s) + } + if len(s.Daily.DaysOfWeek) != 7 { + t.Errorf("expected all 7 days for DAILY, got %v", s.Daily.DaysOfWeek) + } + if s.ScheduledRunTime != "01:30:00" { + t.Errorf("expected HH:mm normalized to HH:mm:ss, got %q", s.ScheduledRunTime) + } + if s.Daily.DailyOffset != "SCHEDULED" { + t.Errorf("expected default offset SCHEDULED, got %q", s.Daily.DailyOffset) + } +} + +func TestScheduleWeeklyWithDays(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.ScheduleRepeat = "WEEKLY" + in.ScheduleDaysOfWeek = []string{"monday", "THURSDAY"} + in.RunDateOffset = "ONE_DAY" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + s := out.Request.SchedulingSettings + if s == nil || s.SchedulerMode != "DAILY" || s.Daily == nil { + t.Fatalf("expected WEEKLY mapped to DAILY+daily, got %+v", s) + } + if len(s.Daily.DaysOfWeek) != 2 || s.Daily.DaysOfWeek[0] != "MONDAY" || s.Daily.DaysOfWeek[1] != "THURSDAY" { + t.Errorf("expected [MONDAY THURSDAY], got %v", s.Daily.DaysOfWeek) + } + if s.Daily.DailyOffset != "ONE_DAY" { + t.Errorf("expected offset ONE_DAY, got %q", s.Daily.DailyOffset) + } +} + +func TestScheduleWeeklyRequiresDays(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.ScheduleRepeat = "WEEKLY" // no days + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for WEEKLY without days, got %q (%s)", out.Status, out.Message) + } +} + +func TestScheduleMonthlyByDay(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.ScheduleRepeat = "MONTHLY" + in.ScheduleDayOfMonth = 15 + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + s := out.Request.SchedulingSettings + if s == nil || s.SchedulerMode != "MONTHLY" || s.Monthly == nil { + t.Fatalf("expected a MONTHLY schedule with a monthly sub-object, got %+v", s) + } + if s.Monthly.MonthlyRepeat != "DAY" || s.Monthly.DayNumber != 15 { + t.Errorf("expected DAY/15, got %+v", s.Monthly) + } +} + +func TestScheduleInvalidOffsetForMode(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.ScheduleRepeat = "DAILY" + in.RunDateOffset = "ONE_HOUR" // hourly offset on a daily schedule + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for invalid offset, got %q (%s)", out.Status, out.Message) + } +} + +func TestRowSampling(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.SampleSize = 5000 + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Sampling must be APPLIED in the source query (Postgres RANDOM()/LIMIT form). + q := out.Request.SourceQuery + if !strings.Contains(q, "RANDOM() <") || !strings.Contains(q, "LIMIT 5000") { + t.Errorf("expected dialect sampling in sourceQuery, got %q", q) + } +} + +func TestNoSamplingByDefault(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if strings.Contains(out.Request.SourceQuery, "RANDOM()") { + t.Errorf("expected no sampling by default (scan all rows), got %q", out.Request.SourceQuery) + } +} + +func TestNoAdaptiveSettingsByDefault(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // No adaptive settings unless asked (server applies its defaults). + if out.Request.MonitoringSettings.AdaptiveMonitors.Settings != nil { + t.Errorf("expected no adaptive monitor settings by default, got %+v", out.Request.MonitoringSettings.AdaptiveMonitors.Settings) + } +} + +func TestAdaptiveMonitorSettings(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.DataLookback = 30 // only lookback set -> learningPhase defaults to 4 + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + s := out.Request.MonitoringSettings.AdaptiveMonitors.Settings + if s == nil || s.DataLookBack != 30 || s.LearningPhase != 4 { + t.Fatalf("expected settings lookback=30 learning=4 (defaulted), got %+v", s) + } +} + +// usersMux returns a mux with the DQ connections handler plus mocked core-api user endpoints: +// users/current resolves to admin, and "jane" resolves by name; everything else is unknown. +func usersMux() *http.ServeMux { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + mux.HandleFunc("/rest/2.0/users/current", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"u-current","userName":"admin"}`)) + }) + mux.HandleFunc("/rest/2.0/users", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Query().Get("name") == "jane" { + _, _ = w.Write([]byte(`{"results":[{"id":"u-jane","userName":"jane"}]}`)) + return + } + _, _ = w.Write([]byte(`{"results":[]}`)) + }) + return mux +} + +func TestNotificationsConfigured(t *testing.T) { + server := httptest.NewServer(usersMux()) + defer server.Close() + + in := baseInput() + in.Notify = []string{"jobFailed", "scoreBelow"} + in.NotifyRecipients = []string{"jane"} + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + n := out.Request.Notifications + if n == nil || len(n.NotificationOptions) != 2 { + t.Fatalf("expected 2 notification options, got %+v", n) + } + var foundScore bool + for _, o := range n.NotificationOptions { + if o.NotificationType == "SCORE_LESS_THAN_LIMIT" { + foundScore = true + if o.Quantity != 75 { + t.Errorf("expected scoreBelow default quantity 75, got %d", o.Quantity) + } + } + } + if !foundScore { + t.Errorf("scoreBelow alert missing: %+v", n.NotificationOptions) + } + // Recipients are usernames in a single EMAIL channel: invoking user (admin) + jane. + if len(n.Channels) != 1 || n.Channels[0].Channel != "EMAIL" { + t.Fatalf("expected one EMAIL channel, got %+v", n.Channels) + } + got := strings.Join(n.Channels[0].Recipients, ",") + if got != "admin,jane" { + t.Errorf("expected recipients [admin jane], got %v", n.Channels[0].Recipients) + } +} + +func TestNotificationUnresolvedRecipient(t *testing.T) { + server := httptest.NewServer(usersMux()) + defer server.Close() + + in := baseInput() + in.Notify = []string{"jobFailed"} + in.NotifyRecipients = []string{"ghost"} // not found + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for unresolved recipient, got %q (%s)", out.Status, out.Message) + } + + // With the proceed flag, it creates anyway with only the resolvable recipient (invoking user). + in.NotifyProceedWithoutUnresolved = true + out, err = tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusPreview { + t.Fatalf("expected preview when proceeding without unresolved, got %q (%s)", out.Status, out.Message) + } + n := out.Request.Notifications + if n == nil || len(n.Channels) != 1 || len(n.Channels[0].Recipients) != 1 || n.Channels[0].Recipients[0] != "admin" { + t.Errorf("expected only the invoking user (admin) as recipient, got %+v", n) + } +} + +func TestUnknownNotificationNeedsInput(t *testing.T) { + server := httptest.NewServer(usersMux()) + defer server.Close() + + in := baseInput() + in.Notify = []string{"bogusAlert"} + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for unknown notification, got %q (%s)", out.Status, out.Message) + } +} + +func TestNoNotificationsByDefault(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Request.Notifications != nil { + t.Errorf("expected no notifications by default, got %+v", out.Request.Notifications) + } +} + +func TestNoScheduleByDefault(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Request.SchedulingSettings != nil { + t.Errorf("expected no schedule by default, got %+v", out.Request.SchedulingSettings) + } + // No time slice -> plain whole-table query, no ${rd}. + if strings.Contains(out.Request.SourceQuery, "${rd}") { + t.Errorf("did not expect ${rd} without a timeSliceColumn, got %q", out.Request.SourceQuery) + } +} + +func TestSelectedColumnsInQuery(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.SelectedColumns = []string{"txn_ts", " ", "amount"} // blank entry must be dropped + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Selected columns scope the SELECT list of the source query (blank dropped, Postgres-escaped). + if !strings.Contains(out.Request.SourceQuery, `SELECT "txn_ts", "amount" FROM`) { + t.Errorf("expected selected columns in the source query SELECT list, got %q", out.Request.SourceQuery) + } +} + +func TestConfirmCreatesAndReturnsRunID(t *testing.T) { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", connectionsHandler()) + mux.Handle("/rest/dq/1.0/jobs", createdJobHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Confirm = true + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusCreated { + t.Fatalf("expected created, got %q (%s)", out.Status, out.Message) + } + if out.JobID != "run-1" { + t.Errorf("expected jobRunId run-1, got %q", out.JobID) + } + if out.JobName != "sales.transactions" { + t.Errorf("expected server-echoed job name, got %q", out.JobName) + } +} + +func TestMissingFieldsNeedsInput(t *testing.T) { + server := httptest.NewServer(http.NewServeMux()) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), tools.Input{EdgeSiteName: "EDGE-1", JobType: "PULLUP"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input, got %q (%s)", out.Status, out.Message) + } +} + +func hasWarning(warnings []string, substr string) bool { + for _, w := range warnings { + if strings.Contains(w, substr) { + return true + } + } + return false +} + +// ---- WS2: Pullup sizing & Parallel JDBC ---- + +func TestPullupAutoSizingByDefault(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Automatic sizing = sparkJobSizing omitted (the public API auto-sizes when it is absent). + if out.Request.JobSettings.PullupSettings.SparkJobSizing != nil { + t.Errorf("expected automatic sizing (nil sparkJobSizing) by default, got %+v", out.Request.JobSettings.PullupSettings.SparkJobSizing) + } +} + +func TestPullupManualSizing(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.SizingMaxExecutors = 4 + in.SizingExecutorMemoryGb = "2" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + s := out.Request.JobSettings.PullupSettings.SparkJobSizing + if s == nil { + t.Fatalf("expected manual sizing object, got nil") + } + if s.NumExecutors != 4 || s.ExecutorMemoryGb != 2 { + t.Errorf("manual sizing values not applied: %+v", s) + } + // Unset manual fields fall back to wizard defaults (1 GB / 1). + if s.NumExecutorCores != 1 || s.DriverMemoryGb != 1 || s.MemoryOverheadGb != 1 { + t.Errorf("expected unset manual fields to default to 1, got %+v", s) + } +} + +func TestParallelJdbcManualRequiresPartitionCount(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.ParallelJdbcPartitionColumn = "id" // specific column implies MANUAL; a count is required + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input when a partition column is set without a count, got %q (%s)", out.Status, out.Message) + } +} + +func TestParallelJdbcManualValid(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.ParallelJdbcPartitionColumn = "id" + in.ParallelJdbcPartitionNumber = 8 + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + pj := out.Request.JobSettings.PullupSettings.LoadOptions.ParallelJdbcOptions + if pj == nil || pj.Mode != "MANUAL" || pj.PartitionColumn != "id" || pj.PartitionNumber != 8 { + t.Errorf("expected MANUAL parallel JDBC id/8, got %+v", pj) + } +} + +func TestParallelJdbcAutoColumnRequiresCount(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.ParallelJdbcMode = "AUTO_COLUMN" // count required, no column + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for AUTO_COLUMN without a partition count, got %q (%s)", out.Status, out.Message) + } +} + +// ---- WS3: Pushdown compute settings ---- + +func TestPushdownComputeSettings(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.JobType = "PUSHDOWN" + in.PushdownConnections = 20 + in.PushdownThreads = 5 + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + ps := out.Request.JobSettings.PushdownSettings + if ps == nil || ps.Connections != 20 || ps.Threads != 5 { + t.Errorf("expected pushdown compute 20/5, got %+v", ps) + } +} + +func TestPushdownComputeOutOfRange(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.JobType = "PUSHDOWN" + in.PushdownConnections = 100 // > 50 + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for connections out of range, got %q (%s)", out.Status, out.Message) + } +} + +func TestSizingRejectedOnPushdown(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.JobType = "PUSHDOWN" + in.SizingMaxExecutors = 2 // sizing is PULLUP-only + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for sizing on a PUSHDOWN job, got %q (%s)", out.Status, out.Message) + } +} + +func TestComputeRejectedOnPullup(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE")) + defer server.Close() + + in := baseInput() + in.PushdownConnections = 10 // compute is PUSHDOWN-only + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for compute on a PULLUP job, got %q (%s)", out.Status, out.Message) + } +} + +// ---- WS4: per-type notification messages ---- + +func TestPerTypeNotificationMessage(t *testing.T) { + server := httptest.NewServer(usersMux()) + defer server.Close() + + in := baseInput() + in.Notify = []string{"jobFailed", "scoreBelow"} + in.NotifyMessages = map[string]string{"jobFailed": "page me"} + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + n := out.Request.Notifications + if n == nil || !n.UseIndividualMessages { + t.Fatalf("expected useIndividualMessages=true, got %+v", n) + } + var msg string + for _, o := range n.NotificationOptions { + if o.NotificationType == "JOB_FAILED" { + msg = o.Message + } + } + if msg != "page me" { + t.Errorf("expected jobFailed per-type message 'page me', got %q", msg) + } +} + +// ---- WS6: permission preflight ---- + +func TestCreatePermissionDenied(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_RUN")) // run but not create + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusError { + t.Fatalf("expected error when lacking create permission, got %q (%s)", out.Status, out.Message) + } + if !strings.Contains(out.Message, "DATA_QUALITY_JOB_CREATE") { + t.Errorf("expected the missing permission named, got %q", out.Message) + } +} + +func TestSchedulePermissionDropsSchedule(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE", "DATA_QUALITY_JOB_RUN")) // no schedule + defer server.Close() + + in := baseInput() + in.ScheduleRepeat = "DAILY" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Request.SchedulingSettings != nil { + t.Errorf("expected schedule dropped without Schedule permission, got %+v", out.Request.SchedulingSettings) + } + if !hasWarning(out.Warnings, "Schedule") { + t.Errorf("expected a Schedule-permission warning, got %v", out.Warnings) + } +} + +func TestRunPermissionWarns(t *testing.T) { + server := httptest.NewServer(muxWithPerms("DATA_QUALITY_JOB_CREATE", "DATA_QUALITY_JOB_SCHEDULE")) // no run + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !hasWarning(out.Warnings, "Run permission") { + t.Errorf("expected a Run-permission warning, got %v", out.Warnings) + } +} + +// ---- WS1: job name handling ---- + +func TestServerGeneratedJobName(t *testing.T) { + mux := muxWithPerms("DATA_QUALITY_JOB_CREATE") + mux.HandleFunc("/rest/dq/internal/v1/jobs/name", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jobName":"sales.transactions_2"}`)) + }) + server := httptest.NewServer(mux) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.JobName != "sales.transactions_2" { + t.Errorf("expected server-generated unique name, got %q", out.JobName) + } +} + +func TestInvalidJobNameRejected(t *testing.T) { + mux := muxWithPerms("DATA_QUALITY_JOB_CREATE") + mux.HandleFunc("/rest/dq/internal/v1/jobs/", func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/validJobName") { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte("false")) + return + } + w.WriteHeader(http.StatusNotFound) + }) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.JobName = "bad name!" + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsInput { + t.Fatalf("expected needs_input for an invalid job name, got %q (%s)", out.Status, out.Message) + } + if out.AffectedStep == "" { + t.Errorf("expected an affectedStep for the name error") + } +} + +// ---- WS7: success deep link ---- + +func TestSuccessReturnsJobDetailsLink(t *testing.T) { + mux := muxWithPerms("DATA_QUALITY_JOB_CREATE", "DATA_QUALITY_JOB_SCHEDULE", "DATA_QUALITY_JOB_RUN") + mux.Handle("/rest/dq/1.0/jobs", createdJobHandler()) + server := httptest.NewServer(mux) + defer server.Close() + + in := baseInput() + in.Confirm = true + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusCreated { + t.Fatalf("expected created, got %q (%s)", out.Status, out.Message) + } + if out.JobDetailsLink != "/data-quality/jobs?jobName=sales.transactions" { + t.Errorf("expected job details deep link, got %q", out.JobDetailsLink) + } +} diff --git a/pkg/tools/prepare_create_dq_job/tool.go b/pkg/tools/prepare_create_dq_job/tool.go new file mode 100644 index 0000000..fa5dedc --- /dev/null +++ b/pkg/tools/prepare_create_dq_job/tool.go @@ -0,0 +1,494 @@ +// Package prepare_create_dq_job implements the prepare_create_dq_job MCP tool — +// a read-only companion to create_dq_job. Given whatever the agent knows so far +// (a connection, and optionally a data source / schema / table), it walks the +// same discovery chain the data-quality job-creation wizard uses: resolve the +// connection, detect the job type (PUSHDOWN/PULLUP) from its capabilities, and +// enumerate data sources, schemas, tables, and columns. It returns a status +// that tells the agent whether it has everything needed to call create_dq_job, +// what's still missing (with the options to choose from), or what couldn't be +// resolved. It performs NO mutations. +package prepare_create_dq_job + +import ( + "context" + "fmt" + "net/http" + "strings" + + "github.com/collibra/chip/pkg/chip" + "github.com/collibra/chip/pkg/clients" + "github.com/google/uuid" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// maxOptions caps how many options are returned in any one response. +const maxOptions = 200 + +// pageLimit is the per-request page size for the monitoring/edge list endpoints, +// which reject limit > 100 with a 400 VALIDATION_ERROR. +const pageLimit = 100 + +// Status is the discovery outcome that drives the next conversational turn. +type Status string + +const ( + // StatusReady means connection + data source + schema + table all resolved; + // `resolved` holds the exact inputs to pass to create_dq_job. + StatusReady Status = "ready" + // StatusIncomplete means a required selection is missing; the response + // includes the pre-fetched options for the next field to choose. + StatusIncomplete Status = "incomplete" + // StatusNeedsClarification means an input could not be resolved (unknown + // connection, ambiguous job type, etc.); options for recovery are included. + StatusNeedsClarification Status = "needs_clarification" +) + +// Input is the tool's typed input. Provide as much as is known; omit a field to +// enumerate the options for it. +type Input struct { + // TableAssetID resolves the whole data location from a catalog Table asset (the "from a table + // asset page" entry point). When set, the tool walks Table -> Schema -> Database -> System and + // matches the DQ connection by its systemAssetId, then enumerates as usual — no need to pass + // connection/dataSource/schema/table. + TableAssetID string `json:"tableAssetId,omitempty" jsonschema:"DGC catalog Table asset UUID. When provided, the connection/dataSource/schema/table are resolved from the asset (via the connection's systemAssetId mapping); the other location fields can be omitted."` + // TableAssetURL / TableAssetName are alternative ways to identify the catalog Table asset (the + // AC's URL / name identifiers). tableAssetId wins if set; then URL; then name (with optional + // tableAssetDomain to disambiguate). Name lookup uses the public assets API (no search index). + TableAssetURL string `json:"tableAssetUrl,omitempty" jsonschema:"Catalog Table asset URL (e.g. https:///asset/). The asset UUID is extracted from it."` + TableAssetName string `json:"tableAssetName,omitempty" jsonschema:"Catalog Table asset name (signifier, e.g. 'transactions'). Resolved via the public assets API; if multiple match, the options are returned for the user to pick one (optionally narrow with tableAssetDomain)."` + TableAssetDomain string `json:"tableAssetDomain,omitempty" jsonschema:"Optional domain/path substring (e.g. 'sales') to disambiguate when tableAssetName matches multiple assets."` + + Connection string `json:"connection,omitempty" jsonschema:"The data-quality edge connection — accepts a connection UUID or its name (case-insensitive, e.g. 'POSTGRES-SOURCE'). On a catalog table-asset page pass tableAssetId instead; from an LLM, omit it to list available connections and ask the user."` + DataSourceName string `json:"dataSourceName,omitempty" jsonschema:"The database/catalog within the connection (e.g. 'postgres'). Omit to list the data sources available on the resolved connection."` + SchemaName string `json:"schemaName,omitempty" jsonschema:"The schema within the data source (e.g. 'sales'). Omit to list the schemas in the resolved data source."` + TableName string `json:"tableName,omitempty" jsonschema:"The table to monitor (e.g. 'customers'). Omit to list the tables in the resolved schema."` +} + +// Output is the structured discovery response. +type Output struct { + Status Status `json:"status" jsonschema:"ready when connection+dataSource+schema+table are all resolved; incomplete when a selection is missing (options provided); needs_clarification when an input could not be resolved."` + Message string `json:"message" jsonschema:"Human-readable summary of the outcome and what to do next."` + Resolved *ResolvedPlan `json:"resolved,omitempty" jsonschema:"Present only when status=ready. Pass these fields straight to create_dq_job."` + JobType string `json:"jobType,omitempty" jsonschema:"Detected job type for the resolved connection: PUSHDOWN or PULLUP. Empty when the connection advertises more than one and the user must choose."` + ConnectionOptions []ConnectionOption `json:"connectionOptions,omitempty" jsonschema:"Connections to choose from — returned when 'connection' was omitted or could not be resolved."` + TableAssetOptions []TableAssetOption `json:"tableAssetOptions,omitempty" jsonschema:"Returned when tableAssetName matched multiple Table assets. Present these to the user; re-call with the chosen asset's id as tableAssetId."` + DataSourceOptions []string `json:"dataSourceOptions,omitempty" jsonschema:"Data source names to choose from — returned when 'dataSourceName' was omitted or unmatched."` + SchemaOptions []string `json:"schemaOptions,omitempty" jsonschema:"Schema names to choose from — returned when 'schemaName' was omitted or unmatched."` + TableOptions []string `json:"tableOptions,omitempty" jsonschema:"Table names to choose from — returned when 'tableName' was omitted or unmatched."` + Columns []ColumnInfo `json:"columns,omitempty" jsonschema:"Columns of the resolved table (the actual schema). Offer these so the user can pick a subset to monitor, then pass the chosen names to create_dq_job.selectedColumns. Omit selectedColumns to monitor all columns (the default)."` + Monitors []clients.DqMonitorInfo `json:"monitors,omitempty" jsonschema:"Present only when status=ready. The available profile monitors, each with a defaultEnabled flag. Show these to the user so they can choose a set, then pass the chosen keys to create_dq_job.monitors (omit to use the defaults). Note: enabling descriptiveStatistics unmasks sensitive data."` + AdaptiveMonitorSettings []clients.DqAdaptiveMonitorSetting `json:"adaptiveMonitorSettings,omitempty" jsonschema:"Present only when status=ready. The 'Advanced monitor settings' (adaptive behavior) the user can tune, each with its default. ALWAYS surface these together with monitors — do not omit them: tell the user data lookback and learning phase are adjustable (defaults 10 and 4) and pass overrides to create_dq_job.dataLookback / learningPhase."` + Notifications []clients.DqNotificationInfo `json:"notifications,omitempty" jsonschema:"Present only when status=ready. The available notification alerts, each with a defaultEnabled flag and whether it takes a threshold quantity. Offer these to the user; pass the chosen keys to create_dq_job.notify (+ thresholds + notifyRecipients). The invoking user is always a recipient; additional recipients are validated against active accounts."` + UnsupportedOptions []string `json:"unsupportedOptions,omitempty" jsonschema:"Present only when status=ready. The wizard options this tool does NOT set, each with the default the server will apply. Present these to the user PROACTIVELY at the ready step (do not wait to be asked) so they know what is auto-filled and what would require the DQ UI."` + OptionsTruncated bool `json:"optionsTruncated" jsonschema:"True when an options list was truncated below the instance's true total."` +} + +// ResolvedPlan is the set of inputs create_dq_job needs, fully resolved. +type ResolvedPlan struct { + JobType string `json:"jobType" jsonschema:"PUSHDOWN or PULLUP. Empty if the connection supports both and the user must pick."` + SuggestedJobName string `json:"suggestedJobName" jsonschema:"Default job name '.
'. The user may override; the server auto-increments on collision."` + EdgeSiteName string `json:"edgeSiteName" jsonschema:"dataLocation.edgeSiteName for create_dq_job."` + EdgeConnectionName string `json:"edgeConnectionName" jsonschema:"dataLocation.edgeConnectionName for create_dq_job."` + DataSourceName string `json:"dataSourceName" jsonschema:"dataLocation.dataSourceName for create_dq_job."` + SchemaName string `json:"schemaName" jsonschema:"dataLocation.schemaName for create_dq_job."` + TableName string `json:"tableName" jsonschema:"dataLocation.tableName for create_dq_job."` + DatabaseProductName string `json:"databaseProductName,omitempty" jsonschema:"dataLocation.databaseProductName for create_dq_job (e.g. POSTGRES)."` + TableAssetLink string `json:"tableAssetLink,omitempty" jsonschema:"Catalog deep-link path to the resolved Table asset (relative to the instance URL). Present only when resolved from tableAssetId."` +} + +// ConnectionOption is one selectable connection. +type ConnectionOption struct { + ConnectionID string `json:"connectionId" jsonschema:"Connection UUID."` + ConnectionName string `json:"connectionName" jsonschema:"Connection display name."` + CapabilityTypes []string `json:"capabilityTypes" jsonschema:"DQ capabilities the connection supports (PUSHDOWN/PULLUP) — i.e. the possible job types."` + DatabaseProduct string `json:"databaseProductName,omitempty" jsonschema:"Database vendor (e.g. POSTGRES)."` +} + +// TableAssetOption is one catalog Table asset candidate for disambiguation. +type TableAssetOption struct { + AssetID string `json:"assetId" jsonschema:"Catalog Table asset UUID — pass as tableAssetId to select it."` + DisplayName string `json:"displayName" jsonschema:"Table signifier (name)."` + Domain string `json:"domain,omitempty" jsonschema:"Domain/path of the asset (helps tell duplicates apart)."` + FullName string `json:"fullName,omitempty" jsonschema:"Fully-qualified asset name."` +} + +// ColumnInfo is one column of the resolved table. +type ColumnInfo struct { + Name string `json:"name" jsonschema:"Column name."` + Type string `json:"type" jsonschema:"Source column type (e.g. int4, text, numeric, timestamp)."` +} + +// NewTool returns the registered tool. +func NewTool(collibraClient *http.Client) *chip.Tool[Input, Output] { + return &chip.Tool[Input, Output]{ + Name: "prepare_create_dq_job", + Title: "Prepare to Create Data Quality Job", + Description: "Read-only companion to create_dq_job. Walks the data-quality wizard's discovery chain — resolve the edge " + + "connection, detect the job type (PUSHDOWN/PULLUP) from its capabilities, and enumerate data sources, schemas, " + + "tables, and columns — driven by whatever inputs you already have. Returns status='ready' with a fully-resolved " + + "plan to hand to create_dq_job, 'incomplete' with the options for the next field to pick, or 'needs_clarification' " + + "when something couldn't be resolved. Call this to gather/validate inputs before create_dq_job; make one selection " + + "per turn (connection -> dataSource -> schema -> table) and re-call until status='ready'.", + Handler: handler(collibraClient), + Permissions: []string{}, + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true}, + } +} + +func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { + return func(ctx context.Context, input Input) (Output, error) { + // Step 0: catalog table-asset entry point. Identify the Table asset by id, URL, or name + // (signifier), then resolve the full data location from it (Table -> Schema -> Database -> + // System -> connection via systemAssetId) and fall through to the normal discovery chain. + assetID := strings.TrimSpace(input.TableAssetID) + if assetID == "" && strings.TrimSpace(input.TableAssetURL) != "" { + assetID = extractAssetID(input.TableAssetURL) + if assetID == "" { + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("Could not extract an asset UUID from URL %q. Use a catalog asset URL like .../asset/, or pass tableAssetId.", input.TableAssetURL)}, nil + } + } + if assetID == "" && strings.TrimSpace(input.TableAssetName) != "" { + matches, err := clients.FindTableAssetsByName(ctx, collibraClient, strings.TrimSpace(input.TableAssetName), maxOptions) + if err != nil { + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("Could not look up Table assets named %q: %v", input.TableAssetName, err)}, nil + } + if d := strings.TrimSpace(input.TableAssetDomain); d != "" { + matches = filterByDomain(matches, d) + } + switch len(matches) { + case 0: + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("No Table asset named %q was found. Provide a more specific name (optionally tableAssetDomain), the table asset URL, or the table asset ID.", input.TableAssetName)}, nil + case 1: + assetID = matches[0].ID + default: + return Output{ + Status: StatusNeedsClarification, + Message: fmt.Sprintf("%d Table assets are named %q — pick one and call again with its assetId as tableAssetId (or narrow with tableAssetDomain).", len(matches), input.TableAssetName), + TableAssetOptions: toTableAssetOptions(matches), + }, nil + } + } + if assetID != "" { + loc, err := clients.ResolveDqLocationFromTableAsset(ctx, collibraClient, assetID) + if err != nil { + return Output{ + Status: StatusNeedsClarification, + Message: fmt.Sprintf("Could not resolve table asset %s to a DQ location: %v", assetID, err), + }, nil + } + input.Connection = loc.ConnectionID + input.DataSourceName = loc.DataSourceName + input.SchemaName = loc.SchemaName + input.TableName = loc.TableName + input.TableAssetID = assetID // so the ready block emits the catalog deep link + } + + // Step 1: a connection is the entry point. Without one, list options. + if strings.TrimSpace(input.Connection) == "" { + return enumerateConnections(ctx, collibraClient, "Provide a connection (or a tableAssetId) to begin. Pick one from connectionOptions and call again.") + } + + conn, err := resolveConnection(ctx, collibraClient, input.Connection) + if err != nil { + return enumerateConnections(ctx, collibraClient, + fmt.Sprintf("Could not resolve connection %q: %v. Pick one from connectionOptions and call again.", input.Connection, err)) + } + + // Job type comes straight off the connection's capabilities. + jobType, jobTypeMsg := detectJobType(conn) + out := Output{JobType: jobType} + + // Step 2: data source. + dataSources, err := clients.ListDqDataSources(ctx, collibraClient, conn.ConnectionID, pageLimit, 0) + if err != nil { + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("Resolved connection %q but failed to list data sources: %v", conn.ConnectionName, err)}, nil + } + dsNames := dataSourceNames(dataSources) + if strings.TrimSpace(input.DataSourceName) == "" { + out.Status = StatusIncomplete + out.Message = joinMsg(jobTypeMsg, fmt.Sprintf("dataSourceName is required. Pick one from dataSourceOptions on connection %q and call again.", conn.ConnectionName)) + out.DataSourceOptions, out.OptionsTruncated = capStrings(dsNames) + return out, nil + } + matchedDS, ok := matchIgnoreCase(dsNames, input.DataSourceName) + if !ok { + out.Status = StatusNeedsClarification + out.Message = fmt.Sprintf("Data source %q was not found on connection %q. Pick one from dataSourceOptions.", input.DataSourceName, conn.ConnectionName) + out.DataSourceOptions, out.OptionsTruncated = capStrings(dsNames) + return out, nil + } + + // Step 3: schema. + schemas, err := clients.ListDqSchemas(ctx, collibraClient, conn.EdgeSiteID, conn.ConnectionID, matchedDS, pageLimit, 0) + if err != nil { + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("Failed to list schemas in data source %q: %v", matchedDS, err)}, nil + } + schemaNames := schemaNames(schemas) + if strings.TrimSpace(input.SchemaName) == "" { + out.Status = StatusIncomplete + out.Message = joinMsg(jobTypeMsg, fmt.Sprintf("schemaName is required. Pick one from schemaOptions in data source %q and call again.", matchedDS)) + out.SchemaOptions, out.OptionsTruncated = capStrings(schemaNames) + return out, nil + } + matchedSchema, ok := matchIgnoreCase(schemaNames, input.SchemaName) + if !ok { + out.Status = StatusNeedsClarification + out.Message = fmt.Sprintf("Schema %q was not found in data source %q. Pick one from schemaOptions.", input.SchemaName, matchedDS) + out.SchemaOptions, out.OptionsTruncated = capStrings(schemaNames) + return out, nil + } + + // Step 4: table. + tables, err := clients.ListDqTables(ctx, collibraClient, conn.EdgeSiteID, conn.ConnectionID, matchedDS, matchedSchema, pageLimit, 0) + if err != nil { + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("Failed to list tables in schema %q: %v", matchedSchema, err)}, nil + } + tableNames := tableNames(tables) + if strings.TrimSpace(input.TableName) == "" { + out.Status = StatusIncomplete + out.Message = joinMsg(jobTypeMsg, fmt.Sprintf("tableName is required. Pick one from tableOptions in schema %q and call again.", matchedSchema)) + out.TableOptions, out.OptionsTruncated = capStrings(tableNames) + return out, nil + } + matchedTable, ok := matchIgnoreCase(tableNames, input.TableName) + if !ok { + out.Status = StatusNeedsClarification + out.Message = fmt.Sprintf("Table %q was not found in schema %q. Pick one from tableOptions.", input.TableName, matchedSchema) + out.TableOptions, out.OptionsTruncated = capStrings(tableNames) + return out, nil + } + + // Everything resolved — surface columns and the ready-to-create plan. + columns, err := clients.ListDqColumns(ctx, collibraClient, conn.EdgeSiteID, conn.ConnectionID, matchedDS, matchedSchema, matchedTable, pageLimit, 0) + if err != nil { + // Columns are advisory; don't fail the resolution if they can't be fetched. + out.Message = fmt.Sprintf("(columns could not be fetched: %v) ", err) + } else { + out.Columns = toColumnInfos(columns) + } + + out.Status = StatusReady + out.Monitors = clients.DqMonitorCatalog() + out.AdaptiveMonitorSettings = clients.DqAdaptiveMonitorSettings() + out.Notifications = clients.DqNotificationCatalog() + out.UnsupportedOptions = clients.UnsupportedWizardOptions(jobType) + + // Suggested job name: ask the server for a collision-free default (auto-increments, e.g. "..._2"). + suggestedName := matchedSchema + "." + matchedTable + if generated, gErr := clients.GenerateUniqueJobName(ctx, collibraClient, matchedSchema, matchedTable); gErr == nil && strings.TrimSpace(generated) != "" { + suggestedName = generated + } + + // Permission preflight (best-effort): Create is a hard prerequisite — if it's missing, the flow + // should not begin. Schedule/Run gaps are surfaced as a note (create_dq_job degrades gracefully). + permNote := "" + if global, resource, permErr := clients.GetDqConnectionPermissions(ctx, collibraClient, conn.ConnectionID); permErr == nil { + has := func(p string) bool { return clients.HasPermission(global, p) || clients.HasPermission(resource, p) } + manageAll := has(clients.PermResourceManageAll) + if !manageAll && !has(clients.PermDqJobCreate) { + return Output{ + Status: StatusNeedsClarification, + JobType: jobType, + Message: fmt.Sprintf("You do not have the Data Quality Job > Create permission (DATA_QUALITY_JOB_CREATE) on connection %q required to create a DQ job on %s.%s, so the flow cannot begin. Ask an administrator for the Data Quality Editor or Data Quality Manager role on this connection.", conn.ConnectionName, matchedSchema, matchedTable), + }, nil + } + var missing []string + if !manageAll && !has(clients.PermDqJobSchedule) { + missing = append(missing, "Schedule") + } + if !manageAll && !has(clients.PermDqJobRun) { + missing = append(missing, "Run") + } + if len(missing) > 0 { + permNote = fmt.Sprintf(" NOTE: you lack the Data Quality Job > %s permission(s) on this connection; create_dq_job will warn and degrade (a schedule is dropped; a run may be rejected server-side).", strings.Join(missing, " and ")) + } + } + + out.Resolved = &ResolvedPlan{ + JobType: jobType, + SuggestedJobName: suggestedName, + EdgeSiteName: conn.EdgeSiteName, + EdgeConnectionName: conn.ConnectionName, + DataSourceName: matchedDS, + SchemaName: matchedSchema, + TableName: matchedTable, + DatabaseProductName: conn.DatabaseProductName, + } + if id := strings.TrimSpace(input.TableAssetID); id != "" { + out.Resolved.TableAssetLink = clients.CatalogAssetPath(id) + } + out.Message += fmt.Sprintf("Ready. Resolved a %s job for %s.%s on connection %q (%d column(s)). Before creating, PROACTIVELY tell the user (do not wait for them to ask): (1) only connection/data source/schema/table are required — everything else is auto-filled, and the exact defaults are listed in unsupportedOptions; (2) offer them a custom job name (default %q); (3) offer column selection (a subset from `columns` → create_dq_job.selectedColumns; default = all %d) and an optional single-column row filter (create_dq_job.filterColumn/filterOperator/filterValue, e.g. amount > 100); (4) walk the Monitors step, which has THREE parts you must each mention: (a) monitor toggles — `monitors` lists them with defaultEnabled flags, pass a custom set to create_dq_job.monitors (enabling descriptiveStatistics unmasks sensitive data); (b) ADVANCED monitor settings — explicitly tell the user that Data Lookback and Learning Phase are adjustable (see `adaptiveMonitorSettings`; defaults 10 and 4) and can be overridden via create_dq_job.dataLookback/learningPhase; (c) row sampling — create_dq_job.sampleSize (default = all rows); (5) for recurring monitoring, offer a schedule (create_dq_job.scheduleRepeat = HOURLY/DAILY/WEEKLY/WEEKDAYS/MONTHLY + scheduleRunTime) and PAIR it with a timeSliceColumn (a date/timestamp column from `columns`) so each run scans only that period's slice — the tool writes a ${rd}/${rdEnd} WHERE for you; without a timeSliceColumn every run rescans the whole table; (6) offer notifications — `notifications` lists the available alerts with defaultEnabled flags; pass chosen keys to create_dq_job.notify (+ thresholds), and additional recipients via create_dq_job.notifyRecipients (the invoking user is always included; others are validated against active accounts); (7) note those listed options are not configurable through this tool and would need the DQ UI. Then call create_dq_job with `resolved` — it returns a preview before anything is created.", + displayJobType(jobType), matchedSchema, matchedTable, conn.ConnectionName, len(out.Columns), suggestedName, len(out.Columns)) + out.Message += " (8) for a PULLUP job, optionally Size Job Resources — automatic by default (recommended); manual sizing (create_dq_job.sizing*) and Parallel JDBC (create_dq_job.parallelJdbc*) are advanced; for a PUSHDOWN job, optionally set compute via create_dq_job.pushdownConnections/pushdownThreads (defaults 10/2). Per-notification messages can be set via create_dq_job.notifyMessages." + permNote + if jobTypeMsg != "" { + out.Message = joinMsg(jobTypeMsg, out.Message) + } + return out, nil + } +} + +// resolveConnection matches the input against connections by UUID or name +// (case-insensitive). Errors carry enough context for enumerateConnections. +func resolveConnection(ctx context.Context, client *http.Client, identifier string) (*clients.DqConnection, error) { + conns, err := clients.ListDqConnections(ctx, client) + if err != nil { + return nil, err + } + id := strings.TrimSpace(identifier) + var byName []clients.DqConnection + for i := range conns { + if conns[i].ConnectionID == id { + return &conns[i], nil + } + if strings.EqualFold(conns[i].ConnectionName, id) { + byName = append(byName, conns[i]) + } + } + switch len(byName) { + case 1: + return &byName[0], nil + case 0: + return nil, fmt.Errorf("no connection matched") + default: + return nil, fmt.Errorf("%d connections share that name — use the connection UUID", len(byName)) + } +} + +// detectJobType returns the single capability type, or "" plus a clarifying +// message when the connection advertises zero or multiple. +func detectJobType(conn *clients.DqConnection) (string, string) { + switch len(conn.CapabilityTypes) { + case 1: + return conn.CapabilityTypes[0], "" + case 0: + return "", fmt.Sprintf("Note: connection %q advertises no DQ capability — it may not support data-quality jobs.", conn.ConnectionName) + default: + return "", fmt.Sprintf("Note: connection %q supports multiple job types (%s); set jobType explicitly when calling create_dq_job.", + conn.ConnectionName, strings.Join(conn.CapabilityTypes, ", ")) + } +} + +func enumerateConnections(ctx context.Context, client *http.Client, message string) (Output, error) { + conns, err := clients.ListDqConnections(ctx, client) + if err != nil { + return Output{Status: StatusNeedsClarification, Message: fmt.Sprintf("Failed to list connections: %v", err)}, nil + } + opts := make([]ConnectionOption, 0, len(conns)) + for _, c := range conns { + opts = append(opts, ConnectionOption{ + ConnectionID: c.ConnectionID, + ConnectionName: c.ConnectionName, + CapabilityTypes: c.CapabilityTypes, + DatabaseProduct: c.DatabaseProductName, + }) + } + truncated := false + if len(opts) > maxOptions { + opts = opts[:maxOptions] + truncated = true + } + return Output{ + Status: StatusIncomplete, + Message: message, + ConnectionOptions: opts, + OptionsTruncated: truncated, + }, nil +} + +func dataSourceNames(in []clients.DqDataSource) []string { + out := make([]string, 0, len(in)) + for _, d := range in { + out = append(out, d.DataSourceName) + } + return out +} + +func schemaNames(in []clients.DqSchema) []string { + out := make([]string, 0, len(in)) + for _, s := range in { + out = append(out, s.Name) + } + return out +} + +func tableNames(in []clients.DqTable) []string { + out := make([]string, 0, len(in)) + for _, t := range in { + out = append(out, t.Name) + } + return out +} + +func toColumnInfos(in []clients.DqColumn) []ColumnInfo { + out := make([]ColumnInfo, 0, len(in)) + for _, c := range in { + out = append(out, ColumnInfo{Name: c.Name, Type: c.Type}) + } + return out +} + +func matchIgnoreCase(options []string, value string) (string, bool) { + for _, o := range options { + if strings.EqualFold(o, value) { + return o, true + } + } + return "", false +} + +func capStrings(in []string) ([]string, bool) { + if len(in) > maxOptions { + return in[:maxOptions], true + } + return in, false +} + +func joinMsg(a, b string) string { + if a == "" { + return b + } + return a + " " + b +} + +// extractAssetID pulls a UUID out of a catalog asset URL (e.g. https://host/asset/?tab=x). +// Returns "" if no UUID path segment is found. +func extractAssetID(raw string) string { + raw = strings.TrimSpace(raw) + if i := strings.IndexAny(raw, "?#"); i >= 0 { + raw = raw[:i] + } + raw = strings.TrimRight(raw, "/") + for _, seg := range strings.Split(raw, "/") { + if _, err := uuid.Parse(seg); err == nil { + return seg + } + } + return "" +} + +// filterByDomain keeps matches whose domain/path contains the given substring (case-insensitive). +func filterByDomain(in []clients.TableAssetMatch, domain string) []clients.TableAssetMatch { + d := strings.ToLower(strings.TrimSpace(domain)) + var out []clients.TableAssetMatch + for _, m := range in { + if strings.Contains(strings.ToLower(m.DomainName), d) { + out = append(out, m) + } + } + return out +} + +func toTableAssetOptions(in []clients.TableAssetMatch) []TableAssetOption { + out := make([]TableAssetOption, 0, len(in)) + for _, m := range in { + out = append(out, TableAssetOption{AssetID: m.ID, DisplayName: m.DisplayName, Domain: m.DomainName, FullName: m.FullName}) + } + return out +} + +func displayJobType(jobType string) string { + if jobType == "" { + return "data-quality" + } + return jobType +} diff --git a/pkg/tools/prepare_create_dq_job/tool_test.go b/pkg/tools/prepare_create_dq_job/tool_test.go new file mode 100644 index 0000000..a03b7a0 --- /dev/null +++ b/pkg/tools/prepare_create_dq_job/tool_test.go @@ -0,0 +1,204 @@ +package prepare_create_dq_job_test + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + tools "github.com/collibra/chip/pkg/tools/prepare_create_dq_job" + "github.com/collibra/chip/pkg/tools/testutil" +) + +const ( + connID = "conn-1" + siteID = "site-1" + // Catalog asset chain UUIDs for the tableAssetId resolution test. + sysAssetID = "00000000-0000-0000-0000-0000000000d1" + dbAssetID = "00000000-0000-0000-0000-0000000000c1" + schemaAssetID = "00000000-0000-0000-0000-0000000000b1" + tableAssetID = "00000000-0000-0000-0000-0000000000a1" +) + +func dqMux() *http.ServeMux { + mux := http.NewServeMux() + mux.Handle("/rest/dq/internal/v1/connections", testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + return http.StatusOK, map[string]any{ + "results": []map[string]any{{ + "connectionId": connID, + "connectionName": "POSTGRES-SOURCE", + "capabilityTypes": []string{"PULLUP"}, + "databaseProductName": "POSTGRES", + "edgeSiteId": siteID, + "edgeSiteName": "EDGE-1", + "systemAssetId": sysAssetID, + }}, + } + })) + // Catalog asset graph (GraphQL): Table -> Schema -> Database -> System via incoming relations. + mux.HandleFunc("/graphql/knowledgeGraph/v1", func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + s := string(body) + w.Header().Set("Content-Type", "application/json") + var asset string + switch { + case strings.Contains(s, tableAssetID): + asset = `{"id":"` + tableAssetID + `","displayName":"customers","type":{"name":"Table"},"incomingRelations":[{"type":{"id":"r","role":"contains"},"source":{"id":"` + schemaAssetID + `","displayName":"sales","type":{"name":"Schema"}}}]}` + case strings.Contains(s, schemaAssetID): + asset = `{"id":"` + schemaAssetID + `","displayName":"sales","type":{"name":"Schema"},"incomingRelations":[{"type":{"id":"r","role":"has"},"source":{"id":"` + dbAssetID + `","displayName":"postgres","type":{"name":"Database"}}}]}` + case strings.Contains(s, dbAssetID): + asset = `{"id":"` + dbAssetID + `","displayName":"postgres","type":{"name":"Database"},"incomingRelations":[{"type":{"id":"r","role":"groups"},"source":{"id":"` + sysAssetID + `","displayName":"postgres","type":{"name":"System"}}}]}` + } + _, _ = w.Write([]byte(`{"data":{"assets":[` + asset + `]}}`)) + }) + mux.Handle("/rest/dq/internal/v1/monitoring/edge/connections/"+connID+"/dataSources", testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + return http.StatusOK, map[string]any{"total": 1, "results": []map[string]any{{"dataSourceName": "postgres", "supportsSchemas": true}}} + })) + mux.Handle("/rest/dq/internal/v1/monitoring/edge/"+siteID+"/connections/"+connID+"/schemas", testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + return http.StatusOK, map[string]any{"total": 1, "results": []map[string]any{{"name": "sales"}}} + })) + mux.Handle("/rest/dq/internal/v1/monitoring/edge/"+siteID+"/connections/"+connID+"/tables", testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + return http.StatusOK, map[string]any{"total": 1, "results": []map[string]any{{"name": "customers", "type": "TABLE"}}} + })) + mux.Handle("/rest/dq/internal/v1/monitoring/edge/"+siteID+"/connections/"+connID+"/columns", testutil.JsonHandlerOut(func(r *http.Request) (int, map[string]any) { + return http.StatusOK, map[string]any{"results": []map[string]any{{"name": "id", "type": "int4"}, {"name": "balance", "type": "numeric"}}} + })) + return mux +} + +func TestReadyResolvesFullPlan(t *testing.T) { + server := httptest.NewServer(dqMux()) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), tools.Input{ + Connection: "postgres-source", // case-insensitive name match + DataSourceName: "postgres", + SchemaName: "sales", + TableName: "customers", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusReady { + t.Fatalf("expected status ready, got %q (%s)", out.Status, out.Message) + } + if out.Resolved == nil { + t.Fatalf("expected resolved plan, got nil") + } + if out.Resolved.JobType != "PULLUP" { + t.Errorf("expected jobType PULLUP, got %q", out.Resolved.JobType) + } + if out.Resolved.EdgeSiteName != "EDGE-1" || out.Resolved.EdgeConnectionName != "POSTGRES-SOURCE" { + t.Errorf("unexpected edge fields: %+v", out.Resolved) + } + if out.Resolved.DatabaseProductName != "POSTGRES" { + t.Errorf("expected databaseProductName POSTGRES, got %q", out.Resolved.DatabaseProductName) + } + if out.Resolved.SuggestedJobName != "sales.customers" { + t.Errorf("expected suggestedJobName sales.customers, got %q", out.Resolved.SuggestedJobName) + } + if len(out.Columns) != 2 { + t.Errorf("expected 2 columns, got %d", len(out.Columns)) + } + // Monitors are surfaced at ready with a default-enabled indicator. + if len(out.Monitors) != 9 { + t.Fatalf("expected 9 available monitors, got %d", len(out.Monitors)) + } + byKey := map[string]bool{} + for _, m := range out.Monitors { + byKey[m.Key] = m.DefaultEnabled + } + if !byKey["rowCount"] { + t.Errorf("expected rowCount to be defaultEnabled") + } + if byKey["min"] { + t.Errorf("expected min to be off by default") + } + if byKey["descriptiveStatistics"] { + t.Errorf("expected descriptiveStatistics to be off by default") + } + // Advanced monitor settings are surfaced explicitly (with defaults) at ready. + if len(out.AdaptiveMonitorSettings) != 2 { + t.Fatalf("expected 2 adaptive monitor settings, got %d", len(out.AdaptiveMonitorSettings)) + } + defByKey := map[string]int{} + for _, s := range out.AdaptiveMonitorSettings { + defByKey[s.Key] = s.Default + } + if defByKey["dataLookback"] != 10 || defByKey["learningPhase"] != 4 { + t.Errorf("expected dataLookback=10, learningPhase=4 defaults, got %+v", out.AdaptiveMonitorSettings) + } + // Notifications catalog is surfaced at ready. + if len(out.Notifications) != 7 { + t.Fatalf("expected 7 notification options, got %d", len(out.Notifications)) + } + notifDefault := map[string]bool{} + for _, n := range out.Notifications { + notifDefault[n.Key] = n.DefaultEnabled + } + if !notifDefault["jobFailed"] || notifDefault["jobCompleted"] { + t.Errorf("unexpected notification defaults: %+v", out.Notifications) + } +} + +func TestTableAssetIdResolvesLocation(t *testing.T) { + server := httptest.NewServer(dqMux()) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), tools.Input{ + TableAssetID: tableAssetID, // resolve everything from the catalog Table asset + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusReady { + t.Fatalf("expected status ready, got %q (%s)", out.Status, out.Message) + } + r := out.Resolved + if r == nil { + t.Fatalf("expected resolved plan, got nil") + } + if r.EdgeConnectionName != "POSTGRES-SOURCE" || r.DataSourceName != "postgres" || r.SchemaName != "sales" || r.TableName != "customers" { + t.Errorf("unexpected resolved location from table asset: %+v", r) + } + if r.TableAssetLink != "/asset/"+tableAssetID { + t.Errorf("expected table asset deep link, got %q", r.TableAssetLink) + } +} + +func TestNoConnectionEnumerates(t *testing.T) { + server := httptest.NewServer(dqMux()) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), tools.Input{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusIncomplete { + t.Fatalf("expected status incomplete, got %q", out.Status) + } + if len(out.ConnectionOptions) != 1 || out.ConnectionOptions[0].ConnectionName != "POSTGRES-SOURCE" { + t.Fatalf("expected one connection option POSTGRES-SOURCE, got %+v", out.ConnectionOptions) + } +} + +func TestUnknownSchemaNeedsClarification(t *testing.T) { + server := httptest.NewServer(dqMux()) + defer server.Close() + + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), tools.Input{ + Connection: connID, // resolve by UUID + DataSourceName: "postgres", + SchemaName: "nope", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if out.Status != tools.StatusNeedsClarification { + t.Fatalf("expected needs_clarification, got %q (%s)", out.Status, out.Message) + } + if len(out.SchemaOptions) != 1 || out.SchemaOptions[0] != "sales" { + t.Errorf("expected schemaOptions [sales], got %+v", out.SchemaOptions) + } +} diff --git a/pkg/tools/register.go b/pkg/tools/register.go index fee28b6..87ad993 100644 --- a/pkg/tools/register.go +++ b/pkg/tools/register.go @@ -8,6 +8,7 @@ import ( "github.com/collibra/chip/pkg/skills" "github.com/collibra/chip/pkg/tools/add_data_classification_match" "github.com/collibra/chip/pkg/tools/create_asset" + "github.com/collibra/chip/pkg/tools/create_dq_job" "github.com/collibra/chip/pkg/tools/discover_business_glossary" "github.com/collibra/chip/pkg/tools/discover_data_assets" "github.com/collibra/chip/pkg/tools/edit_asset" @@ -27,6 +28,7 @@ import ( "github.com/collibra/chip/pkg/tools/list_context_specifications" "github.com/collibra/chip/pkg/tools/list_data_contracts" "github.com/collibra/chip/pkg/tools/prepare_create_asset" + "github.com/collibra/chip/pkg/tools/prepare_create_dq_job" "github.com/collibra/chip/pkg/tools/pull_data_contract_manifest" "github.com/collibra/chip/pkg/tools/push_data_contract_manifest" "github.com/collibra/chip/pkg/tools/remove_data_classification_match" @@ -80,6 +82,8 @@ func RegisterAll(server *chip.Server, client *http.Client, toolConfig *chip.Serv toolRegister(server, toolConfig, list_context_specifications.NewTool(client)) toolRegister(server, toolConfig, get_context_specification.NewTool(client)) } + toolRegister(server, toolConfig, prepare_create_dq_job.NewTool(client)) + toolRegister(server, toolConfig, create_dq_job.NewTool(client)) if toolConfig.EnableDebugTools { toolRegister(server, toolConfig, get_debug_mcp_init_request.NewTool(client)) From deb3c62eb53b225ea0fac56594fc984f2d88111a Mon Sep 17 00:00:00 2001 From: vvaks0 Date: Wed, 15 Jul 2026 19:33:57 -0400 Subject: [PATCH 2/3] feat(dq): use public API for job naming, gate DQ tools behind --experimental, clarify descriptions Addresses review feedback on #88. - Drop the internal /jobs/name, /jobs/{name}/exists, and /jobs/{name}/validJobName calls. Naming now relies on the PUBLIC create API: omitting jobName lets the server auto-assign a collision-free name (returned in the response); a user-supplied name is validated client-side against the server's rules (charset + no leading hyphen). - Rename the tools to create_data_quality_job / prepare_create_data_quality_job and expand their descriptions with plain-language framing (job/edge/connection) plus example user requests, for better LLM comprehension and evaluations. - Gate both tools behind a new "data-quality" experimental feature (off by default), like the Data Product skills, since they write to Collibra. Co-Authored-By: Claude Opus 4.8 (1M context) --- cmd/chip/experimental.go | 3 +- pkg/clients/dgc_dq_client.go | 78 +++++------------- pkg/tools/create_dq_job/tool.go | 104 ++++++++++++++---------- pkg/tools/create_dq_job/tool_test.go | 33 ++++---- pkg/tools/prepare_create_dq_job/tool.go | 68 ++++++++-------- pkg/tools/register.go | 11 ++- pkg/tools/register_test.go | 20 +++++ 7 files changed, 163 insertions(+), 154 deletions(-) diff --git a/cmd/chip/experimental.go b/cmd/chip/experimental.go index faf0a22..dc39d01 100644 --- a/cmd/chip/experimental.go +++ b/cmd/chip/experimental.go @@ -17,8 +17,9 @@ import ( // the help text, example YAML, and validation warning all read from this // map so no other code needs to change. var knownExperimentalFeatures = map[string]string{ - skills.FeatureName: "Embedded skill catalog served via list_collibra_skills and load_collibra_skill.", + skills.FeatureName: "Embedded skill catalog served via list_collibra_skills and load_collibra_skill.", tools.ContextSpecificationsFeature: "Context specification tools: list_context_specifications, get_context_specification, and contextSpecificationId parameter on get_asset_details.", + tools.DataQualityFeatureName: "Data Quality job tools (prepare_create_data_quality_job, create_data_quality_job) that create and queue DQ jobs.", } // validateExperimental warns (without exiting) when the user enabled an diff --git a/pkg/clients/dgc_dq_client.go b/pkg/clients/dgc_dq_client.go index aa41b51..3f3f007 100644 --- a/pkg/clients/dgc_dq_client.go +++ b/pkg/clients/dgc_dq_client.go @@ -8,6 +8,7 @@ import ( "log/slog" "net/http" "net/url" + "regexp" "strconv" "strings" ) @@ -398,7 +399,7 @@ type DqProfileMonitors struct { } // DqMonitorInfo describes one profile monitor for display and selection. Key matches the -// DqProfileMonitors JSON field and the create_dq_job `monitors` input; DefaultEnabled marks +// DqProfileMonitors JSON field and the create_data_quality_job `monitors` input; DefaultEnabled marks // the monitors the wizard turns on by default. type DqMonitorInfo struct { Key string `json:"key"` @@ -500,7 +501,7 @@ func EnabledMonitorKeys(pm *DqProfileMonitors) []string { } // DqAdaptiveMonitorSetting describes one tunable "Advanced monitor setting" (the adaptive -// behavior in the Monitors step). Key matches the create_dq_job input; Default is the wizard +// behavior in the Monitors step). Key matches the create_data_quality_job input; Default is the wizard // default applied when the user doesn't override. type DqAdaptiveMonitorSetting struct { Key string `json:"key"` @@ -708,64 +709,23 @@ func GetDqDataDistribution(ctx context.Context, collibraHttpClient *http.Client, } // ===================================================================================== -// Job-name helpers (the wizard's Step 1 name handling) — all on the internal jobs surface. +// Job-name validation (client-side). +// +// The job-creation flow no longer calls the internal /jobs/name, /jobs/{name}/exists, or +// /jobs/{name}/validJobName endpoints. Collision-free naming is delegated to the PUBLIC create API: +// omitting jobName makes the server auto-assign and auto-increment it (e.g. "sales.orders_2"), and +// the assigned name is returned in the create response. A user-supplied name is validated here +// against the same rules the server enforces before submit. // ===================================================================================== -type dqJobNameRequest struct { - SchemaName string `json:"schemaName"` - TableName string `json:"tableName"` -} - -type dqJobNameResponse struct { - JobName string `json:"jobName"` -} - -// GenerateUniqueJobName asks the server for a collision-free default job name for schema.table. -// The server auto-increments (e.g. "sales.orders_2") when the base name is taken, so the user is -// never asked to resolve a name conflict manually — POST /rest/dq/internal/v1/jobs/name. -func GenerateUniqueJobName(ctx context.Context, collibraHttpClient *http.Client, schemaName, tableName string) (string, error) { - payload, err := json.Marshal(dqJobNameRequest{SchemaName: schemaName, TableName: tableName}) - if err != nil { - return "", fmt.Errorf("failed to marshal job-name request: %w", err) - } - req, err := http.NewRequestWithContext(ctx, "POST", "/rest/dq/internal/v1/jobs/name", bytes.NewBuffer(payload)) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - body, err := executeRequest(collibraHttpClient, req) - if err != nil { - return "", err - } - var resp dqJobNameResponse - if err := json.Unmarshal(body, &resp); err != nil { - return "", fmt.Errorf("failed to parse job-name response: %w", err) - } - return resp.JobName, nil -} - -// IsValidDqJobName asks the server whether a job name is syntactically valid — the wizard rejects -// special characters other than - and _. GET /rest/dq/internal/v1/jobs/{name}/validJobName. -func IsValidDqJobName(ctx context.Context, collibraHttpClient *http.Client, jobName string) (bool, error) { - return dqGetBool(ctx, collibraHttpClient, "/rest/dq/internal/v1/jobs/"+url.PathEscape(jobName)+"/validJobName") -} - -// DqJobExists reports whether a job with the given name already exists — -// GET /rest/dq/internal/v1/jobs/{name}/exists. -func DqJobExists(ctx context.Context, collibraHttpClient *http.Client, jobName string) (bool, error) { - return dqGetBool(ctx, collibraHttpClient, "/rest/dq/internal/v1/jobs/"+url.PathEscape(jobName)+"/exists") -} +// dqJobNamePattern mirrors the server's dataset-name constraint (@DatasetName, +// VALIDATION_PATTERN_DATASET_NAME default ^[a-zA-Z0-9_.-]+$). +var dqJobNamePattern = regexp.MustCompile(`^[a-zA-Z0-9_.-]+$`) -func dqGetBool(ctx context.Context, collibraHttpClient *http.Client, endpoint string) (bool, error) { - body, err := dqGet(ctx, collibraHttpClient, endpoint) - if err != nil { - return false, err - } - var b bool - if err := json.Unmarshal(body, &b); err != nil { - return false, fmt.Errorf("failed to parse boolean response from %s: %w", endpoint, err) - } - return b, nil +// IsValidDqJobName reports whether jobName satisfies the DQ server's job-name rules: the dataset-name +// charset (letters, digits, '.', '-', '_') and no leading hyphen (the server's validJobName check). +func IsValidDqJobName(jobName string) bool { + return jobName != "" && !strings.HasPrefix(jobName, "-") && dqJobNamePattern.MatchString(jobName) } // ===================================================================================== @@ -1064,8 +1024,8 @@ func dqPageQuery(values url.Values, limit, offset int) string { // UnsupportedWizardOptions lists the data-quality wizard configuration steps the // create-DQ-job tools do NOT expose yet, each paired with the default the server -// applies. Shared by prepare_create_dq_job (to disclose proactively at the ready -// step) and create_dq_job (to disclose again at preview), so the user is never +// applies. Shared by prepare_create_data_quality_job (to disclose proactively at the ready +// step) and create_data_quality_job (to disclose again at preview), so the user is never // misled about scope. jobType selects the type-specific entry (Sizing for Pullup, // Compute for Pushdown). func UnsupportedWizardOptions(jobType string) []string { diff --git a/pkg/tools/create_dq_job/tool.go b/pkg/tools/create_dq_job/tool.go index aead255..e23db89 100644 --- a/pkg/tools/create_dq_job/tool.go +++ b/pkg/tools/create_dq_job/tool.go @@ -1,11 +1,11 @@ -// Package create_dq_job implements the create_dq_job MCP tool — it creates a +// Package create_data_quality_job implements the create_data_quality_job MCP tool — it creates a // Collibra data-quality job for a table and (per the wizard) queues a run. // // IT POSTS TO THE PUBLIC DQ API: POST /rest/dq/1.0/jobs (clients.CreateDqJob / // clients.CreateDqJobRequest, contract dq/udq-app-client/oas/dq-v1-public-oas-spec.yaml). // The public create accepts the full job definition — sourceQuery, runDate window, monitors, // schedule, back-runs, notifications, and pullup/pushdown settings — so the internal BFF -// endpoint is no longer needed. (Discovery in prepare_create_dq_job still uses internal +// endpoint is no longer needed. (Discovery in prepare_create_data_quality_job still uses internal // endpoints; the public DQ API exposes no connection/edge metadata browse.) // // HOW SCHEDULING + TIME-SLICE + ${rd} FIT TOGETHER: @@ -42,10 +42,10 @@ const ( StatusError Status = "error" ) -// Input — the five dataLocation fields are required (easiest via prepare_create_dq_job's +// Input — the five dataLocation fields are required (easiest via prepare_create_data_quality_job's // `resolved` block). Time-slice / schedule / back-run fields are optional and structured. type Input struct { - EdgeSiteName string `json:"edgeSiteName" jsonschema:"Edge site name. From prepare_create_dq_job resolved.edgeSiteName."` + EdgeSiteName string `json:"edgeSiteName" jsonschema:"Edge site name. From prepare_create_data_quality_job resolved.edgeSiteName."` EdgeConnectionName string `json:"edgeConnectionName" jsonschema:"Edge connection name, e.g. 'POSTGRES-SOURCE'. The tool resolves its databaseProductName + job type from the connection for the create request."` DataSourceName string `json:"dataSourceName" jsonschema:"Data source / database, e.g. 'postgres'."` SchemaName string `json:"schemaName" jsonschema:"Schema, e.g. 'sales'."` @@ -56,12 +56,12 @@ type Input struct { // --- Column selection. Omit to monitor ALL columns (the default). Provide a subset to // scope profiling/monitoring to just those columns (each sent with selected=true). --- - SelectedColumns []string `json:"selectedColumns,omitempty" jsonschema:"Columns to monitor. Omit for ALL columns (default). Provide a subset — exact names from prepare_create_dq_job's columns — to profile only those. Does not change the source query (still SELECT *); only scopes what is profiled."` + SelectedColumns []string `json:"selectedColumns,omitempty" jsonschema:"Columns to monitor. Omit for ALL columns (default). Provide a subset — exact names from prepare_create_data_quality_job's columns — to profile only those. Does not change the source query (still SELECT *); only scopes what is profiled."` // --- Monitors. Omit to use the default set; provide an authoritative set of monitor keys // to enable (anything not listed is turned off). descriptiveStatistics unmasks sensitive // data — only include it after explicit user confirmation. --- - Monitors []string `json:"monitors,omitempty" jsonschema:"Monitor keys to enable — authoritative (anything omitted is turned OFF). Omit to use the defaults (rowCount, nullValues, emptyFields, uniqueness). Valid keys (see prepare_create_dq_job's monitors): rowCount, nullValues, emptyFields, uniqueness, min, mean, max, executionTime, descriptiveStatistics. descriptiveStatistics UNMASKS sensitive data — only include it after explicit user confirmation."` + Monitors []string `json:"monitors,omitempty" jsonschema:"Monitor keys to enable — authoritative (anything omitted is turned OFF). Omit to use the defaults (rowCount, nullValues, emptyFields, uniqueness). Valid keys (see prepare_create_data_quality_job's monitors): rowCount, nullValues, emptyFields, uniqueness, min, mean, max, executionTime, descriptiveStatistics. descriptiveStatistics UNMASKS sensitive data — only include it after explicit user confirmation."` // --- Advanced monitor settings (adaptive behavior). Omit both to use the wizard defaults // (lookback 10, learning phase 4). Setting either sends a structured `settings` object. --- @@ -71,7 +71,7 @@ type Input struct { // --- Row filter (the wizard's single-column predicate). Scopes the job to rows matching // "filterColumn filterOperator filterValue" (e.g. amount > 100). Omit to scan all rows. // This is ONE predicate — not compound AND/OR or free-form SQL (mirrors the DQ wizard). --- - FilterColumn string `json:"filterColumn,omitempty" jsonschema:"Column for a single-predicate row filter (the wizard's row filter). Use a name from prepare_create_dq_job's columns. Set filterColumn + filterOperator (+ filterValue) to scope the job to matching rows; omit to scan all rows."` + FilterColumn string `json:"filterColumn,omitempty" jsonschema:"Column for a single-predicate row filter (the wizard's row filter). Use a name from prepare_create_data_quality_job's columns. Set filterColumn + filterOperator (+ filterValue) to scope the job to matching rows; omit to scan all rows."` FilterOperator string `json:"filterOperator,omitempty" jsonschema:"Comparison operator for the row filter: = != <> > >= < <= (the wizard's set) or LIKE; the value is treated as a single literal. For IS NULL / IS NOT NULL leave filterValue empty. Required when filterColumn is set."` FilterValue string `json:"filterValue,omitempty" jsonschema:"Right-hand value for the row filter, passed BARE — do NOT add quotes. The tool wraps it in single quotes for you (matching the DQ wizard), e.g. pass US (not 'US'), 100, 2024-01-01. Leave empty for valueless operators (IS NULL / IS NOT NULL). Applied by writing it into the job's source-query WHERE."` @@ -82,7 +82,7 @@ type Input struct { // --- Time slice (incremental scanning). Pick a DATE COLUMN; when set, the tool writes a // ">= '${rd}' AND < '${rdEnd}'" predicate into the job SQL's WHERE. The scheduler // substitutes ${rd}/${rdEnd} per run, so each run scans only that slice, not the whole table. --- - TimeSliceColumn string `json:"timeSliceColumn,omitempty" jsonschema:"Date/timestamp column to slice on (e.g. 'txn_ts'). When set, the tool adds a WHERE \"\" >= '${rd}' AND \"\" < '${rdEnd}' predicate to the job SQL so each scheduled/backrun run scans only that slice. REQUIRED for incremental scheduling — without it, every scheduled run rescans the whole table. Pick a real date/timestamp column from prepare_create_dq_job's columns."` + TimeSliceColumn string `json:"timeSliceColumn,omitempty" jsonschema:"Date/timestamp column to slice on (e.g. 'txn_ts'). When set, the tool adds a WHERE \"\" >= '${rd}' AND \"\" < '${rdEnd}' predicate to the job SQL so each scheduled/backrun run scans only that slice. REQUIRED for incremental scheduling — without it, every scheduled run rescans the whole table. Pick a real date/timestamp column from prepare_create_data_quality_job's columns."` TimeSliceSize int `json:"timeSliceSize,omitempty" jsonschema:"Time-slice window width (default 1 when timeSliceColumn set). Combined with timeSliceUnit, e.g. size=1 unit=DAYS = one day per run."` TimeSliceUnit string `json:"timeSliceUnit,omitempty" jsonschema:"HOURS | DAYS | WEEKS | MONTHS (default DAYS)."` @@ -102,7 +102,7 @@ type Input struct { // --- Notifications (optional). Configured when `notify` or `notifyRecipients` is set; otherwise // no notifications. The invoking user is always the default recipient. --- - Notify []string `json:"notify,omitempty" jsonschema:"Notification keys to enable (authoritative; omit but set notifyRecipients to use the defaults jobFailed/rowsBelow/scoreBelow/runTimeAbove). Keys (see prepare_create_dq_job notifications): jobFailed, rowsBelow, scoreBelow, runTimeAbove, jobCompleted, runsWithoutData, daysWithoutData."` + Notify []string `json:"notify,omitempty" jsonschema:"Notification keys to enable (authoritative; omit but set notifyRecipients to use the defaults jobFailed/rowsBelow/scoreBelow/runTimeAbove). Keys (see prepare_create_data_quality_job notifications): jobFailed, rowsBelow, scoreBelow, runTimeAbove, jobCompleted, runsWithoutData, daysWithoutData."` NotifyRowsBelow int `json:"notifyRowsBelow,omitempty" jsonschema:"Threshold for rowsBelow — alert when row count <= this. Default 1."` NotifyScoreBelow int `json:"notifyScoreBelow,omitempty" jsonschema:"Threshold for scoreBelow — alert when score (0-100) <= this. Default 75."` NotifyRunTimeAboveMinutes int `json:"notifyRunTimeAboveMinutes,omitempty" jsonschema:"Threshold for runTimeAbove — alert when run time minutes > this. Default 60."` @@ -141,9 +141,9 @@ type Input struct { PushdownConnections int `json:"pushdownConnections,omitempty" jsonschema:"PUSHDOWN compute: number of connections (1-50). Default 10 when omitted."` PushdownThreads int `json:"pushdownThreads,omitempty" jsonschema:"PUSHDOWN compute: number of threads (1-10). Default 2 when omitted."` - // --- Catalog linkage (optional). When set (e.g. from prepare_create_dq_job's resolved table + // --- Catalog linkage (optional). When set (e.g. from prepare_create_data_quality_job's resolved table // asset), the success result includes a deep link to the catalog Table asset. --- - TableAssetID string `json:"tableAssetId,omitempty" jsonschema:"Catalog Table asset UUID this job monitors (from prepare_create_dq_job). When set, the success result includes a catalog deep link to the asset."` + TableAssetID string `json:"tableAssetId,omitempty" jsonschema:"Catalog Table asset UUID this job monitors (from prepare_create_data_quality_job). When set, the success result includes a catalog deep link to the asset."` // AcknowledgeDescriptiveStatistics must be true to enable the descriptiveStatistics monitor — it // UNMASKS sensitive values. Without it the tool refuses to proceed (the wizard's explicit-confirm). @@ -169,18 +169,28 @@ type Output struct { func NewTool(collibraClient *http.Client) *chip.Tool[Input, Output] { return &chip.Tool[Input, Output]{ - Name: "create_dq_job", + Name: "create_data_quality_job", Title: "Create Data Quality Job", - Description: "Creates a Collibra data-quality job for a table (and queues a run). Provide the table location " + - "(edgeSiteName, edgeConnectionName, dataSourceName, schemaName, tableName — easiest from prepare_create_dq_job's " + - "`resolved`). Optional structured config: column selection (selectedColumns — subset to monitor; omit for all), monitors (monitors — " + - "set of monitor keys to enable; omit for defaults) with adaptive settings (dataLookback/learningPhase), row sampling (sampleSize — sample N " + - "rows; omit for all), row filter (filterColumn/filterOperator/filterValue — single-column predicate like amount > 100, inlined into the " + - "source-query WHERE), time slice (timeSliceColumn/Size/Unit — writes a ${rd}/${rdEnd} WHERE so each run scans one slice), schedule " + - "(scheduleRepeat/RunTime/DaysOfWeek/DayOfMonth/MonthlyMode + runDateOffset — pair with a timeSliceColumn for incremental runs), " + - "back runs (backrun*), notifications (notify keys + thresholds + notifyRecipients — the invoking user is always a recipient). jobType resolves from the connection if omitted. " + - "Posts to the PUBLIC DQ API (POST /rest/dq/1.0/jobs). Defaults to a PREVIEW: " + - "confirm=false returns the exact payload without creating; call again with confirm=true after the user approves.", + Description: "Sets up automated data-quality monitoring on a database table in Collibra, and starts the first run. " + + "A 'data quality job' profiles a table and watches it over time for problems — sudden row-count drops, spikes in " + + "null or empty values, duplicates, and other anomalies. Use prepare_create_data_quality_job FIRST to discover and " + + "validate the inputs, then pass its `resolved` block here.\n\n" + + "Where the data lives (all required): edgeSiteName and edgeConnectionName identify the source — a 'connection' is a " + + "saved link to a source database, reached through a Collibra Edge site (the agent that runs the scan); dataSourceName " + + "is the database/catalog, plus schemaName and tableName.\n\n" + + "Optional tuning: selectedColumns (which columns to monitor; omit for all), monitors (which checks to run; omit for " + + "sensible defaults) with adaptive-baseline settings (dataLookback/learningPhase), sampleSize (scan a sample instead " + + "of every row), a single-column row filter (filterColumn/filterOperator/filterValue, e.g. amount > 100), a time slice " + + "(timeSliceColumn so each run only checks that period's new rows), a recurring schedule " + + "(scheduleRepeat/scheduleRunTime/… — pair it with a timeSliceColumn for incremental runs), historical back-runs " + + "(backrun*), and notifications (notify keys + thresholds + notifyRecipients — the requesting user is always notified).\n\n" + + "jobType is PUSHDOWN (the checks run inside the source database) or PULLUP (data is pulled into Spark to be checked); " + + "it is detected from the connection if you omit it. Leave jobName empty to let Collibra assign a unique name.\n\n" + + "Safety: this WRITES to Collibra, so it defaults to a PREVIEW — confirm=false returns the exact request without " + + "creating anything; review it with the user, then call again with confirm=true to create and run the job.\n\n" + + "Example user requests: \"Set up data quality monitoring on the sales.orders table\"; \"Create a DQ job for orders " + + "and alert me if the row count drops\"; \"Watch my Postgres customers table for nulls and duplicates every day\"; " + + "\"Start checking data quality on this table.\"", Handler: handler(collibraClient), Permissions: []string{}, } @@ -192,7 +202,7 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { return Output{ Status: StatusNeedsInput, Message: fmt.Sprintf("Missing required field(s): %s.", strings.Join(missing, ", ")), - Guidance: "Call prepare_create_dq_job to resolve connection/data source/schema/table, then pass its `resolved` block here.", + Guidance: "Call prepare_create_data_quality_job to resolve connection/data source/schema/table, then pass its `resolved` block here.", }, nil } @@ -200,7 +210,7 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { // valueless operators like IS NULL). Catch a half-specified filter before any API call. if hasFilterFields(input) { if strings.TrimSpace(input.FilterColumn) == "" { - return Output{Status: StatusNeedsInput, Message: "filterOperator/filterValue given without filterColumn.", Guidance: "Set filterColumn (a column from prepare_create_dq_job) to apply a row filter, or omit all filter fields."}, nil + return Output{Status: StatusNeedsInput, Message: "filterOperator/filterValue given without filterColumn.", Guidance: "Set filterColumn (a column from prepare_create_data_quality_job) to apply a row filter, or omit all filter fields."}, nil } if strings.TrimSpace(input.FilterOperator) == "" { return Output{Status: StatusNeedsInput, Message: "filterColumn set but filterOperator is missing.", Guidance: "Provide filterOperator (e.g. =, >, <, LIKE, IN, IS NULL), plus filterValue unless the operator takes none."}, nil @@ -213,7 +223,7 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { return Output{ Status: StatusNeedsInput, Message: fmt.Sprintf("Unknown monitor(s): %s.", strings.Join(unknown, ", ")), - Guidance: "Use monitor keys from prepare_create_dq_job's `monitors`: " + strings.Join(clients.MonitorKeys(), ", ") + ".", + Guidance: "Use monitor keys from prepare_create_data_quality_job's `monitors`: " + strings.Join(clients.MonitorKeys(), ", ") + ".", }, nil } } @@ -270,7 +280,7 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { return Output{ Status: StatusNeedsInput, Message: fmt.Sprintf("Unknown notification(s): %s.", strings.Join(unknown, ", ")), - Guidance: "Use notification keys from prepare_create_dq_job's `notifications`: " + strings.Join(clients.NotificationKeys(), ", ") + ".", + Guidance: "Use notification keys from prepare_create_data_quality_job's `notifications`: " + strings.Join(clients.NotificationKeys(), ", ") + ".", }, nil } // Recipients: invoking user (always) + any additional usernames/emails, de-duped. The public @@ -333,27 +343,25 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { } } if jobType == "" { - return Output{Status: StatusNeedsInput, Message: "Could not determine job type.", Guidance: "Set jobType to PUSHDOWN or PULLUP, or call prepare_create_dq_job."}, nil + return Output{Status: StatusNeedsInput, Message: "Could not determine job type.", Guidance: "Set jobType to PUSHDOWN or PULLUP, or call prepare_create_data_quality_job."}, nil } if connErr != nil || conn == nil { - return Output{Status: StatusNeedsInput, Message: fmt.Sprintf("Could not resolve connection %q (needed for connectionId/edgeSiteId).", input.EdgeConnectionName), Guidance: "Verify the connection name via prepare_create_dq_job."}, nil + return Output{Status: StatusNeedsInput, Message: fmt.Sprintf("Could not resolve connection %q (needed for connectionId/edgeSiteId).", input.EdgeConnectionName), Guidance: "Verify the connection name via prepare_create_data_quality_job."}, nil } - // Job name: default from the server's collision-free generator (auto-increments, e.g. "..._2"); - // validate a user-provided name against the server's rules (rejects special chars). + // Job name: when the user supplies one, validate it client-side against the DQ server's rules; + // when omitted, leave it empty so the PUBLIC create auto-assigns a collision-free name + // (auto-incremented, e.g. "sales.orders_2"). The assigned name is read back from the response. jobName := strings.TrimSpace(input.JobName) - if jobName == "" { - if generated, err := clients.GenerateUniqueJobName(ctx, collibraClient, input.SchemaName, input.TableName); err == nil && strings.TrimSpace(generated) != "" { - jobName = generated - } else { - jobName = input.SchemaName + "." + input.TableName - } - } else if ok, err := clients.IsValidDqJobName(ctx, collibraClient, jobName); err == nil && !ok { - guidance := "Use only letters, numbers, '-' and '_'." - if suggestion, sErr := clients.GenerateUniqueJobName(ctx, collibraClient, input.SchemaName, input.TableName); sErr == nil && strings.TrimSpace(suggestion) != "" { - guidance += fmt.Sprintf(" A valid default is %q.", suggestion) - } - return Output{Status: StatusNeedsInput, AffectedStep: stepSelectData, Message: fmt.Sprintf("Job name %q is invalid (special characters other than - and _ are not allowed).", jobName), Guidance: guidance}, nil + if jobName != "" && !clients.IsValidDqJobName(jobName) { + return Output{ + Status: StatusNeedsInput, + AffectedStep: stepSelectData, + Message: fmt.Sprintf("Job name %q is invalid.", jobName), + Guidance: fmt.Sprintf( + "Use only letters, numbers, '.', '-' and '_', and don't start with '-'. A valid default is %q.", + input.SchemaName+"."+input.TableName), + }, nil } // Type-specific config must match the job type (sizing/Parallel JDBC are PULLUP-only; compute is PUSHDOWN-only). @@ -440,16 +448,22 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { if len(warnings) > 0 { warnNote = " WARNINGS: " + strings.Join(warnings, " ") } + previewName := jobName + nameClause := "job name is overridable via jobName" + if jobName == "" { + previewName = input.SchemaName + "." + input.TableName + nameClause = fmt.Sprintf("job name will be auto-assigned (default %q, with a numeric suffix if it already exists) and returned on create; set jobName to override", previewName) + } return Output{ Status: StatusPreview, JobType: jobType, - JobName: jobName, + JobName: previewName, Request: req, Warnings: warnings, UnsupportedOptions: clients.UnsupportedWizardOptions(jobType), Message: fmt.Sprintf("Preview only — nothing created. Will create a %s job %q for %s.%s (columns=%s, rows=%s, filter=%q, monitors=[%s], adaptive=%s, timeSlice=%v, schedule=%s, backrun=%v, notifications=%s, %s). "+ - "Review with the user (job name is overridable via jobName), then call again with confirm=true.%s%s%s", - jobType, jobName, input.SchemaName, input.TableName, columnsDesc, sampleDesc, filterDesc, strings.Join(enabledMonitors, ", "), adaptiveDesc, hasSlice, describeSchedule(req.SchedulingSettings), req.Backrun != nil, notifyDesc, describeCompute(req.JobSettings, jobType), sensitiveWarning, rdNote, warnNote), + "Review with the user (%s), then call again with confirm=true.%s%s%s", + jobType, previewName, input.SchemaName, input.TableName, columnsDesc, sampleDesc, filterDesc, strings.Join(enabledMonitors, ", "), adaptiveDesc, hasSlice, describeSchedule(req.SchedulingSettings), req.Backrun != nil, notifyDesc, describeCompute(req.JobSettings, jobType), nameClause, sensitiveWarning, rdNote, warnNote), }, nil } diff --git a/pkg/tools/create_dq_job/tool_test.go b/pkg/tools/create_dq_job/tool_test.go index 8ba74e7..a132644 100644 --- a/pkg/tools/create_dq_job/tool_test.go +++ b/pkg/tools/create_dq_job/tool_test.go @@ -91,8 +91,13 @@ func TestPreviewDoesNotCreate(t *testing.T) { if dl.EdgeConnectionName != "POSTGRES-SOURCE" || dl.EdgeSiteName != "EDGE-1" || dl.DatabaseProductName != "POSTGRES" { t.Errorf("dataLocation not resolved into request: %+v", dl) } - if out.Request.JobType != "PULLUP" || out.Request.JobName != "sales.transactions" { - t.Errorf("unexpected request: jobType=%s jobName=%s", out.Request.JobType, out.Request.JobName) + // jobName is omitted from the request so the PUBLIC create auto-assigns a collision-free name; + // the preview surfaces the default display name. + if out.Request.JobType != "PULLUP" || out.Request.JobName != "" { + t.Errorf("unexpected request: jobType=%s jobName=%q (expected empty so the server auto-assigns)", out.Request.JobType, out.Request.JobName) + } + if out.JobName != "sales.transactions" { + t.Errorf("expected preview display name sales.transactions, got %q", out.JobName) } // PULLUP defaults to automatic sizing: pullupSettings present, sparkJobSizing omitted (nil). if out.Request.JobSettings == nil || out.Request.JobSettings.PullupSettings == nil { @@ -944,34 +949,32 @@ func TestRunPermissionWarns(t *testing.T) { // ---- WS1: job name handling ---- -func TestServerGeneratedJobName(t *testing.T) { +// When the user omits a name, the create request omits jobName and the server-assigned +// (auto-incremented) name from the create response is surfaced. +func TestServerAssignsJobNameWhenOmitted(t *testing.T) { mux := muxWithPerms("DATA_QUALITY_JOB_CREATE") - mux.HandleFunc("/rest/dq/internal/v1/jobs/name", func(w http.ResponseWriter, _ *http.Request) { + mux.HandleFunc("/rest/dq/1.0/jobs", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"jobName":"sales.transactions_2"}`)) + w.WriteHeader(http.StatusCreated) + _, _ = w.Write([]byte(`{"jobName":"sales.transactions_2","jobType":"PULLUP","jobRunId":"run-1"}`)) }) server := httptest.NewServer(mux) defer server.Close() - out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), baseInput()) + in := baseInput() + in.Confirm = true + out, err := tools.NewTool(testutil.NewClient(server)).Handler(t.Context(), in) if err != nil { t.Fatalf("unexpected error: %v", err) } if out.JobName != "sales.transactions_2" { - t.Errorf("expected server-generated unique name, got %q", out.JobName) + t.Errorf("expected the server-assigned job name to be surfaced, got %q", out.JobName) } } func TestInvalidJobNameRejected(t *testing.T) { + // Name validation is client-side now — no server endpoint is consulted. mux := muxWithPerms("DATA_QUALITY_JOB_CREATE") - mux.HandleFunc("/rest/dq/internal/v1/jobs/", func(w http.ResponseWriter, r *http.Request) { - if strings.HasSuffix(r.URL.Path, "/validJobName") { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte("false")) - return - } - w.WriteHeader(http.StatusNotFound) - }) server := httptest.NewServer(mux) defer server.Close() diff --git a/pkg/tools/prepare_create_dq_job/tool.go b/pkg/tools/prepare_create_dq_job/tool.go index fa5dedc..f401d27 100644 --- a/pkg/tools/prepare_create_dq_job/tool.go +++ b/pkg/tools/prepare_create_dq_job/tool.go @@ -1,10 +1,10 @@ -// Package prepare_create_dq_job implements the prepare_create_dq_job MCP tool — -// a read-only companion to create_dq_job. Given whatever the agent knows so far +// Package prepare_create_data_quality_job implements the prepare_create_data_quality_job MCP tool — +// a read-only companion to create_data_quality_job. Given whatever the agent knows so far // (a connection, and optionally a data source / schema / table), it walks the // same discovery chain the data-quality job-creation wizard uses: resolve the // connection, detect the job type (PUSHDOWN/PULLUP) from its capabilities, and // enumerate data sources, schemas, tables, and columns. It returns a status -// that tells the agent whether it has everything needed to call create_dq_job, +// that tells the agent whether it has everything needed to call create_data_quality_job, // what's still missing (with the options to choose from), or what couldn't be // resolved. It performs NO mutations. package prepare_create_dq_job @@ -33,7 +33,7 @@ type Status string const ( // StatusReady means connection + data source + schema + table all resolved; - // `resolved` holds the exact inputs to pass to create_dq_job. + // `resolved` holds the exact inputs to pass to create_data_quality_job. StatusReady Status = "ready" // StatusIncomplete means a required selection is missing; the response // includes the pre-fetched options for the next field to choose. @@ -68,31 +68,31 @@ type Input struct { type Output struct { Status Status `json:"status" jsonschema:"ready when connection+dataSource+schema+table are all resolved; incomplete when a selection is missing (options provided); needs_clarification when an input could not be resolved."` Message string `json:"message" jsonschema:"Human-readable summary of the outcome and what to do next."` - Resolved *ResolvedPlan `json:"resolved,omitempty" jsonschema:"Present only when status=ready. Pass these fields straight to create_dq_job."` + Resolved *ResolvedPlan `json:"resolved,omitempty" jsonschema:"Present only when status=ready. Pass these fields straight to create_data_quality_job."` JobType string `json:"jobType,omitempty" jsonschema:"Detected job type for the resolved connection: PUSHDOWN or PULLUP. Empty when the connection advertises more than one and the user must choose."` ConnectionOptions []ConnectionOption `json:"connectionOptions,omitempty" jsonschema:"Connections to choose from — returned when 'connection' was omitted or could not be resolved."` TableAssetOptions []TableAssetOption `json:"tableAssetOptions,omitempty" jsonschema:"Returned when tableAssetName matched multiple Table assets. Present these to the user; re-call with the chosen asset's id as tableAssetId."` DataSourceOptions []string `json:"dataSourceOptions,omitempty" jsonschema:"Data source names to choose from — returned when 'dataSourceName' was omitted or unmatched."` SchemaOptions []string `json:"schemaOptions,omitempty" jsonschema:"Schema names to choose from — returned when 'schemaName' was omitted or unmatched."` TableOptions []string `json:"tableOptions,omitempty" jsonschema:"Table names to choose from — returned when 'tableName' was omitted or unmatched."` - Columns []ColumnInfo `json:"columns,omitempty" jsonschema:"Columns of the resolved table (the actual schema). Offer these so the user can pick a subset to monitor, then pass the chosen names to create_dq_job.selectedColumns. Omit selectedColumns to monitor all columns (the default)."` - Monitors []clients.DqMonitorInfo `json:"monitors,omitempty" jsonschema:"Present only when status=ready. The available profile monitors, each with a defaultEnabled flag. Show these to the user so they can choose a set, then pass the chosen keys to create_dq_job.monitors (omit to use the defaults). Note: enabling descriptiveStatistics unmasks sensitive data."` - AdaptiveMonitorSettings []clients.DqAdaptiveMonitorSetting `json:"adaptiveMonitorSettings,omitempty" jsonschema:"Present only when status=ready. The 'Advanced monitor settings' (adaptive behavior) the user can tune, each with its default. ALWAYS surface these together with monitors — do not omit them: tell the user data lookback and learning phase are adjustable (defaults 10 and 4) and pass overrides to create_dq_job.dataLookback / learningPhase."` - Notifications []clients.DqNotificationInfo `json:"notifications,omitempty" jsonschema:"Present only when status=ready. The available notification alerts, each with a defaultEnabled flag and whether it takes a threshold quantity. Offer these to the user; pass the chosen keys to create_dq_job.notify (+ thresholds + notifyRecipients). The invoking user is always a recipient; additional recipients are validated against active accounts."` + Columns []ColumnInfo `json:"columns,omitempty" jsonschema:"Columns of the resolved table (the actual schema). Offer these so the user can pick a subset to monitor, then pass the chosen names to create_data_quality_job.selectedColumns. Omit selectedColumns to monitor all columns (the default)."` + Monitors []clients.DqMonitorInfo `json:"monitors,omitempty" jsonschema:"Present only when status=ready. The available profile monitors, each with a defaultEnabled flag. Show these to the user so they can choose a set, then pass the chosen keys to create_data_quality_job.monitors (omit to use the defaults). Note: enabling descriptiveStatistics unmasks sensitive data."` + AdaptiveMonitorSettings []clients.DqAdaptiveMonitorSetting `json:"adaptiveMonitorSettings,omitempty" jsonschema:"Present only when status=ready. The 'Advanced monitor settings' (adaptive behavior) the user can tune, each with its default. ALWAYS surface these together with monitors — do not omit them: tell the user data lookback and learning phase are adjustable (defaults 10 and 4) and pass overrides to create_data_quality_job.dataLookback / learningPhase."` + Notifications []clients.DqNotificationInfo `json:"notifications,omitempty" jsonschema:"Present only when status=ready. The available notification alerts, each with a defaultEnabled flag and whether it takes a threshold quantity. Offer these to the user; pass the chosen keys to create_data_quality_job.notify (+ thresholds + notifyRecipients). The invoking user is always a recipient; additional recipients are validated against active accounts."` UnsupportedOptions []string `json:"unsupportedOptions,omitempty" jsonschema:"Present only when status=ready. The wizard options this tool does NOT set, each with the default the server will apply. Present these to the user PROACTIVELY at the ready step (do not wait to be asked) so they know what is auto-filled and what would require the DQ UI."` OptionsTruncated bool `json:"optionsTruncated" jsonschema:"True when an options list was truncated below the instance's true total."` } -// ResolvedPlan is the set of inputs create_dq_job needs, fully resolved. +// ResolvedPlan is the set of inputs create_data_quality_job needs, fully resolved. type ResolvedPlan struct { JobType string `json:"jobType" jsonschema:"PUSHDOWN or PULLUP. Empty if the connection supports both and the user must pick."` SuggestedJobName string `json:"suggestedJobName" jsonschema:"Default job name '.
'. The user may override; the server auto-increments on collision."` - EdgeSiteName string `json:"edgeSiteName" jsonschema:"dataLocation.edgeSiteName for create_dq_job."` - EdgeConnectionName string `json:"edgeConnectionName" jsonschema:"dataLocation.edgeConnectionName for create_dq_job."` - DataSourceName string `json:"dataSourceName" jsonschema:"dataLocation.dataSourceName for create_dq_job."` - SchemaName string `json:"schemaName" jsonschema:"dataLocation.schemaName for create_dq_job."` - TableName string `json:"tableName" jsonschema:"dataLocation.tableName for create_dq_job."` - DatabaseProductName string `json:"databaseProductName,omitempty" jsonschema:"dataLocation.databaseProductName for create_dq_job (e.g. POSTGRES)."` + EdgeSiteName string `json:"edgeSiteName" jsonschema:"dataLocation.edgeSiteName for create_data_quality_job."` + EdgeConnectionName string `json:"edgeConnectionName" jsonschema:"dataLocation.edgeConnectionName for create_data_quality_job."` + DataSourceName string `json:"dataSourceName" jsonschema:"dataLocation.dataSourceName for create_data_quality_job."` + SchemaName string `json:"schemaName" jsonschema:"dataLocation.schemaName for create_data_quality_job."` + TableName string `json:"tableName" jsonschema:"dataLocation.tableName for create_data_quality_job."` + DatabaseProductName string `json:"databaseProductName,omitempty" jsonschema:"dataLocation.databaseProductName for create_data_quality_job (e.g. POSTGRES)."` TableAssetLink string `json:"tableAssetLink,omitempty" jsonschema:"Catalog deep-link path to the resolved Table asset (relative to the instance URL). Present only when resolved from tableAssetId."` } @@ -121,14 +121,20 @@ type ColumnInfo struct { // NewTool returns the registered tool. func NewTool(collibraClient *http.Client) *chip.Tool[Input, Output] { return &chip.Tool[Input, Output]{ - Name: "prepare_create_dq_job", + Name: "prepare_create_data_quality_job", Title: "Prepare to Create Data Quality Job", - Description: "Read-only companion to create_dq_job. Walks the data-quality wizard's discovery chain — resolve the edge " + - "connection, detect the job type (PUSHDOWN/PULLUP) from its capabilities, and enumerate data sources, schemas, " + - "tables, and columns — driven by whatever inputs you already have. Returns status='ready' with a fully-resolved " + - "plan to hand to create_dq_job, 'incomplete' with the options for the next field to pick, or 'needs_clarification' " + - "when something couldn't be resolved. Call this to gather/validate inputs before create_dq_job; make one selection " + - "per turn (connection -> dataSource -> schema -> table) and re-call until status='ready'.", + Description: "Read-only companion to create_data_quality_job — use it FIRST to gather and validate everything that tool " + + "needs to set up data-quality monitoring on a database table. It walks the discovery chain: find the source " + + "connection (a saved link to a source database, reached through a Collibra Edge site — the agent that runs the scan), " + + "detect the job type (PUSHDOWN vs PULLUP) from what the connection supports, and list the available databases, " + + "schemas, tables, and columns — using whatever you already know (a connection name, a catalog table, etc.).\n\n" + + "Returns status='ready' with a fully-resolved plan to hand straight to create_data_quality_job, 'incomplete' with " + + "the choices for the next field to pick, or 'needs_clarification' when something can't be resolved. Make one " + + "selection per turn (connection -> dataSource -> schema -> table) and call again until status='ready'. This is " + + "read-only — it never creates anything.\n\n" + + "Example user requests: \"I want to monitor data quality on a table\"; \"Set up a DQ check but I'm not sure which " + + "connection to use\"; \"What tables can I run data quality on in my warehouse?\"; \"Prepare a data quality job for " + + "the customers table.\"", Handler: handler(collibraClient), Permissions: []string{}, Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true}, @@ -273,14 +279,12 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { out.Notifications = clients.DqNotificationCatalog() out.UnsupportedOptions = clients.UnsupportedWizardOptions(jobType) - // Suggested job name: ask the server for a collision-free default (auto-increments, e.g. "..._2"). + // Suggested job name: the create step defaults to ".
" and lets the PUBLIC create + // API auto-assign a collision-free name (a numeric suffix is added if it already exists). suggestedName := matchedSchema + "." + matchedTable - if generated, gErr := clients.GenerateUniqueJobName(ctx, collibraClient, matchedSchema, matchedTable); gErr == nil && strings.TrimSpace(generated) != "" { - suggestedName = generated - } // Permission preflight (best-effort): Create is a hard prerequisite — if it's missing, the flow - // should not begin. Schedule/Run gaps are surfaced as a note (create_dq_job degrades gracefully). + // should not begin. Schedule/Run gaps are surfaced as a note (create_data_quality_job degrades gracefully). permNote := "" if global, resource, permErr := clients.GetDqConnectionPermissions(ctx, collibraClient, conn.ConnectionID); permErr == nil { has := func(p string) bool { return clients.HasPermission(global, p) || clients.HasPermission(resource, p) } @@ -300,7 +304,7 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { missing = append(missing, "Run") } if len(missing) > 0 { - permNote = fmt.Sprintf(" NOTE: you lack the Data Quality Job > %s permission(s) on this connection; create_dq_job will warn and degrade (a schedule is dropped; a run may be rejected server-side).", strings.Join(missing, " and ")) + permNote = fmt.Sprintf(" NOTE: you lack the Data Quality Job > %s permission(s) on this connection; create_data_quality_job will warn and degrade (a schedule is dropped; a run may be rejected server-side).", strings.Join(missing, " and ")) } } @@ -317,9 +321,9 @@ func handler(collibraClient *http.Client) chip.ToolHandlerFunc[Input, Output] { if id := strings.TrimSpace(input.TableAssetID); id != "" { out.Resolved.TableAssetLink = clients.CatalogAssetPath(id) } - out.Message += fmt.Sprintf("Ready. Resolved a %s job for %s.%s on connection %q (%d column(s)). Before creating, PROACTIVELY tell the user (do not wait for them to ask): (1) only connection/data source/schema/table are required — everything else is auto-filled, and the exact defaults are listed in unsupportedOptions; (2) offer them a custom job name (default %q); (3) offer column selection (a subset from `columns` → create_dq_job.selectedColumns; default = all %d) and an optional single-column row filter (create_dq_job.filterColumn/filterOperator/filterValue, e.g. amount > 100); (4) walk the Monitors step, which has THREE parts you must each mention: (a) monitor toggles — `monitors` lists them with defaultEnabled flags, pass a custom set to create_dq_job.monitors (enabling descriptiveStatistics unmasks sensitive data); (b) ADVANCED monitor settings — explicitly tell the user that Data Lookback and Learning Phase are adjustable (see `adaptiveMonitorSettings`; defaults 10 and 4) and can be overridden via create_dq_job.dataLookback/learningPhase; (c) row sampling — create_dq_job.sampleSize (default = all rows); (5) for recurring monitoring, offer a schedule (create_dq_job.scheduleRepeat = HOURLY/DAILY/WEEKLY/WEEKDAYS/MONTHLY + scheduleRunTime) and PAIR it with a timeSliceColumn (a date/timestamp column from `columns`) so each run scans only that period's slice — the tool writes a ${rd}/${rdEnd} WHERE for you; without a timeSliceColumn every run rescans the whole table; (6) offer notifications — `notifications` lists the available alerts with defaultEnabled flags; pass chosen keys to create_dq_job.notify (+ thresholds), and additional recipients via create_dq_job.notifyRecipients (the invoking user is always included; others are validated against active accounts); (7) note those listed options are not configurable through this tool and would need the DQ UI. Then call create_dq_job with `resolved` — it returns a preview before anything is created.", + out.Message += fmt.Sprintf("Ready. Resolved a %s job for %s.%s on connection %q (%d column(s)). Before creating, PROACTIVELY tell the user (do not wait for them to ask): (1) only connection/data source/schema/table are required — everything else is auto-filled, and the exact defaults are listed in unsupportedOptions; (2) offer them a custom job name (default %q); (3) offer column selection (a subset from `columns` → create_data_quality_job.selectedColumns; default = all %d) and an optional single-column row filter (create_data_quality_job.filterColumn/filterOperator/filterValue, e.g. amount > 100); (4) walk the Monitors step, which has THREE parts you must each mention: (a) monitor toggles — `monitors` lists them with defaultEnabled flags, pass a custom set to create_data_quality_job.monitors (enabling descriptiveStatistics unmasks sensitive data); (b) ADVANCED monitor settings — explicitly tell the user that Data Lookback and Learning Phase are adjustable (see `adaptiveMonitorSettings`; defaults 10 and 4) and can be overridden via create_data_quality_job.dataLookback/learningPhase; (c) row sampling — create_data_quality_job.sampleSize (default = all rows); (5) for recurring monitoring, offer a schedule (create_data_quality_job.scheduleRepeat = HOURLY/DAILY/WEEKLY/WEEKDAYS/MONTHLY + scheduleRunTime) and PAIR it with a timeSliceColumn (a date/timestamp column from `columns`) so each run scans only that period's slice — the tool writes a ${rd}/${rdEnd} WHERE for you; without a timeSliceColumn every run rescans the whole table; (6) offer notifications — `notifications` lists the available alerts with defaultEnabled flags; pass chosen keys to create_data_quality_job.notify (+ thresholds), and additional recipients via create_data_quality_job.notifyRecipients (the invoking user is always included; others are validated against active accounts); (7) note those listed options are not configurable through this tool and would need the DQ UI. Then call create_data_quality_job with `resolved` — it returns a preview before anything is created.", displayJobType(jobType), matchedSchema, matchedTable, conn.ConnectionName, len(out.Columns), suggestedName, len(out.Columns)) - out.Message += " (8) for a PULLUP job, optionally Size Job Resources — automatic by default (recommended); manual sizing (create_dq_job.sizing*) and Parallel JDBC (create_dq_job.parallelJdbc*) are advanced; for a PUSHDOWN job, optionally set compute via create_dq_job.pushdownConnections/pushdownThreads (defaults 10/2). Per-notification messages can be set via create_dq_job.notifyMessages." + permNote + out.Message += " (8) for a PULLUP job, optionally Size Job Resources — automatic by default (recommended); manual sizing (create_data_quality_job.sizing*) and Parallel JDBC (create_data_quality_job.parallelJdbc*) are advanced; for a PUSHDOWN job, optionally set compute via create_data_quality_job.pushdownConnections/pushdownThreads (defaults 10/2). Per-notification messages can be set via create_data_quality_job.notifyMessages." + permNote if jobTypeMsg != "" { out.Message = joinMsg(jobTypeMsg, out.Message) } @@ -363,7 +367,7 @@ func detectJobType(conn *clients.DqConnection) (string, string) { case 0: return "", fmt.Sprintf("Note: connection %q advertises no DQ capability — it may not support data-quality jobs.", conn.ConnectionName) default: - return "", fmt.Sprintf("Note: connection %q supports multiple job types (%s); set jobType explicitly when calling create_dq_job.", + return "", fmt.Sprintf("Note: connection %q supports multiple job types (%s); set jobType explicitly when calling create_data_quality_job.", conn.ConnectionName, strings.Join(conn.CapabilityTypes, ", ")) } } diff --git a/pkg/tools/register.go b/pkg/tools/register.go index 87ad993..f1c0b77 100644 --- a/pkg/tools/register.go +++ b/pkg/tools/register.go @@ -51,6 +51,11 @@ var CopilotToolNames = []string{ "discover_business_glossary", } +// DataQualityFeatureName gates the data-quality job-creation tools (create_data_quality_job and its +// prepare_ companion) behind --experimental. They WRITE to Collibra (create and queue jobs), so they +// stay opt-in — like the Data Product skills — until they graduate. Off by default. +const DataQualityFeatureName = "data-quality" + func RegisterAll(server *chip.Server, client *http.Client, toolConfig *chip.ServerToolConfig) error { toolRegister(server, toolConfig, discover_data_assets.NewTool(client)) toolRegister(server, toolConfig, discover_business_glossary.NewTool(client)) @@ -82,8 +87,10 @@ func RegisterAll(server *chip.Server, client *http.Client, toolConfig *chip.Serv toolRegister(server, toolConfig, list_context_specifications.NewTool(client)) toolRegister(server, toolConfig, get_context_specification.NewTool(client)) } - toolRegister(server, toolConfig, prepare_create_dq_job.NewTool(client)) - toolRegister(server, toolConfig, create_dq_job.NewTool(client)) + if toolConfig.IsExperimentalEnabled(DataQualityFeatureName) { + toolRegister(server, toolConfig, prepare_create_dq_job.NewTool(client)) + toolRegister(server, toolConfig, create_dq_job.NewTool(client)) + } if toolConfig.EnableDebugTools { toolRegister(server, toolConfig, get_debug_mcp_init_request.NewTool(client)) diff --git a/pkg/tools/register_test.go b/pkg/tools/register_test.go index cd9e89a..176116a 100644 --- a/pkg/tools/register_test.go +++ b/pkg/tools/register_test.go @@ -28,6 +28,26 @@ func TestRegisterAll_DebugToolVisibleWhenEnabled(t *testing.T) { } } +var dataQualityToolNames = []string{"prepare_create_data_quality_job", "create_data_quality_job"} + +func TestRegisterAll_DataQualityToolsHiddenByDefault(t *testing.T) { + names := listToolNames(t, &chip.ServerToolConfig{}) + for _, name := range dataQualityToolNames { + if slices.Contains(names, name) { + t.Fatalf("expected %q to be absent without the %q experimental feature; got tools=%v", name, tools.DataQualityFeatureName, names) + } + } +} + +func TestRegisterAll_DataQualityToolsVisibleWhenEnabled(t *testing.T) { + names := listToolNames(t, &chip.ServerToolConfig{Experimental: []string{tools.DataQualityFeatureName}}) + for _, name := range dataQualityToolNames { + if !slices.Contains(names, name) { + t.Fatalf("expected %q to be present with the %q experimental feature; got tools=%v", name, tools.DataQualityFeatureName, names) + } + } +} + func listToolNames(t *testing.T, cfg *chip.ServerToolConfig) []string { t.Helper() server := chip.NewServer() From 96e56583dfe9b91601be77e5f530879195b82c61 Mon Sep 17 00:00:00 2001 From: vvaks0 Date: Sun, 19 Jul 2026 22:09:09 -0400 Subject: [PATCH 3/3] docs(dq): correct selectedColumns description to match query builder selectedColumns projects the chosen columns into the source query's SELECT list (falling back to SELECT * only when omitted), per dqColumnsPart in pkg/clients/dq_source_query.go. The schema description wrongly claimed it left the query as SELECT *. Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/tools/create_dq_job/tool.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/tools/create_dq_job/tool.go b/pkg/tools/create_dq_job/tool.go index e23db89..9f4cf49 100644 --- a/pkg/tools/create_dq_job/tool.go +++ b/pkg/tools/create_dq_job/tool.go @@ -56,7 +56,7 @@ type Input struct { // --- Column selection. Omit to monitor ALL columns (the default). Provide a subset to // scope profiling/monitoring to just those columns (each sent with selected=true). --- - SelectedColumns []string `json:"selectedColumns,omitempty" jsonschema:"Columns to monitor. Omit for ALL columns (default). Provide a subset — exact names from prepare_create_data_quality_job's columns — to profile only those. Does not change the source query (still SELECT *); only scopes what is profiled."` + SelectedColumns []string `json:"selectedColumns,omitempty" jsonschema:"Columns to monitor. Omit for ALL columns (default → the source query scans SELECT *). Provide a subset — exact names from prepare_create_data_quality_job's columns — to profile only those; the selected columns are projected into the source query (SELECT \"col1\", \"col2\", ...), so only they are scanned and profiled."` // --- Monitors. Omit to use the default set; provide an authoritative set of monitor keys // to enable (anything not listed is turned off). descriptiveStatistics unmasks sensitive