diff --git a/go.mod b/go.mod index 795e1bf498..77393b10f6 100644 --- a/go.mod +++ b/go.mod @@ -148,6 +148,10 @@ require ( github.com/googleapis/gax-go/v2 v2.23.0 // indirect github.com/goph/emperror v0.17.2 // indirect github.com/hashicorp/hcl v1.0.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/pgx/v5 v5.10.0 + github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/klauspost/compress v1.18.6 // indirect diff --git a/go.sum b/go.sum index 355378b132..dc57fd0fbc 100644 --- a/go.sum +++ b/go.sum @@ -311,6 +311,14 @@ github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/infiniflow/infinity/go v0.0.0-20260806040857-d755c5ad25d9 h1:mqPwCOTktwl0V1f3tkfCRFonGQ3XrQ43gvH+TOWec6o= github.com/infiniflow/infinity/go v0.0.0-20260806040857-d755c5ad25d9/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= diff --git a/internal/dao/connector.go b/internal/dao/connector.go index 3a8647b473..ab186dd38c 100644 --- a/internal/dao/connector.go +++ b/internal/dao/connector.go @@ -460,8 +460,11 @@ func scheduleConnectorTask(ctx context.Context, tx *gorm.DB, connectorID, kbID, return "", err } if err == nil { - pollRangeStart = latest.PollRangeEnd totalDocsIndexed = latest.TotalDocsIndexed + if latest.PollRangeEnd != nil { + pollRangeEnd := latest.PollRangeEnd.Time() + pollRangeStart = &pollRangeEnd + } } } @@ -478,7 +481,7 @@ func scheduleConnectorTask(ctx context.Context, tx *gorm.DB, connectorID, kbID, TaskType: taskType, Status: string(entity.TaskStatusSchedule), FromBeginning: &fromBeginning, - PollRangeStart: pollRangeStart, + PollRangeStart: entity.NewFlexibleTime(pollRangeStart), TimeStarted: &now, ErrorMsg: "", TotalDocsIndexed: totalDocsIndexed, diff --git a/internal/dao/sync_task.go b/internal/dao/sync_task.go index 2d2bc953bc..5bfee83981 100644 --- a/internal/dao/sync_task.go +++ b/internal/dao/sync_task.go @@ -331,7 +331,7 @@ func (d *SyncTaskDAO) CompleteSyncTask(ctx context.Context, taskContext SyncTask err := d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { result := tx.Model(&entity.SyncLogs{}).Where("id = ? AND status = ?", taskContext.Task.ID, SyncStatusRunning).Updates(map[string]any{ "status": SyncStatusDone, - "poll_range_end": pollRangeEnd, + "poll_range_end": entity.FlexibleTime(pollRangeEnd), "new_docs_indexed": newDocs, "total_docs_indexed": gorm.Expr("total_docs_indexed + ?", totalDocs), "error_msg": errorMsg, @@ -518,7 +518,7 @@ func createScheduledTask(ctx context.Context, tx *gorm.DB, connectorID, kbID, ta TaskType: taskType, Status: SyncStatusSchedule, FromBeginning: &reindex, - PollRangeStart: pollRangeStart, + PollRangeStart: entity.NewFlexibleTime(pollRangeStart), TimeStarted: &now, ErrorMsg: "", TotalDocsIndexed: totalDocsIndexed, diff --git a/internal/entity/connector.go b/internal/entity/connector.go index a8dd73c121..c3b0b8c3f4 100644 --- a/internal/entity/connector.go +++ b/internal/entity/connector.go @@ -105,21 +105,21 @@ func (Connector2Kb) TableName() string { // SyncLogs sync logs model type SyncLogs struct { - ID string `gorm:"column:id;primaryKey;size:32" json:"id"` - ConnectorID string `gorm:"column:connector_id;size:32;index" json:"connector_id"` - TaskType string `gorm:"column:task_type;size:32;not null;default:sync;index" json:"task_type"` - Status string `gorm:"column:status;size:128;not null;index" json:"status"` - FromBeginning *string `gorm:"column:from_beginning;size:1" json:"from_beginning,omitempty"` - NewDocsIndexed int64 `gorm:"column:new_docs_indexed;default:0" json:"new_docs_indexed"` - TotalDocsIndexed int64 `gorm:"column:total_docs_indexed;default:0" json:"total_docs_indexed"` - DocsRemovedFromIndex int64 `gorm:"column:docs_removed_from_index;default:0" json:"docs_removed_from_index"` - ErrorMsg string `gorm:"column:error_msg;type:longtext;not null" json:"error_msg"` - ErrorCount int64 `gorm:"column:error_count;default:0" json:"error_count"` - FullExceptionTrace *string `gorm:"column:full_exception_trace;type:longtext" json:"full_exception_trace,omitempty"` - TimeStarted *time.Time `gorm:"column:time_started;index" json:"time_started,omitempty"` - PollRangeStart *time.Time `gorm:"column:poll_range_start;index" json:"poll_range_start,omitempty"` - PollRangeEnd *time.Time `gorm:"column:poll_range_end;index" json:"poll_range_end,omitempty"` - KbID string `gorm:"column:kb_id;size:32;not null;index" json:"kb_id"` + ID string `gorm:"column:id;primaryKey;size:32" json:"id"` + ConnectorID string `gorm:"column:connector_id;size:32;index" json:"connector_id"` + TaskType string `gorm:"column:task_type;size:32;not null;default:sync;index" json:"task_type"` + Status string `gorm:"column:status;size:128;not null;index" json:"status"` + FromBeginning *string `gorm:"column:from_beginning;size:1" json:"from_beginning,omitempty"` + NewDocsIndexed int64 `gorm:"column:new_docs_indexed;default:0" json:"new_docs_indexed"` + TotalDocsIndexed int64 `gorm:"column:total_docs_indexed;default:0" json:"total_docs_indexed"` + DocsRemovedFromIndex int64 `gorm:"column:docs_removed_from_index;default:0" json:"docs_removed_from_index"` + ErrorMsg string `gorm:"column:error_msg;type:longtext;not null" json:"error_msg"` + ErrorCount int64 `gorm:"column:error_count;default:0" json:"error_count"` + FullExceptionTrace *string `gorm:"column:full_exception_trace;type:longtext" json:"full_exception_trace,omitempty"` + TimeStarted *time.Time `gorm:"column:time_started;index" json:"time_started,omitempty"` + PollRangeStart *FlexibleTime `gorm:"column:poll_range_start;type:varchar(255);index" json:"poll_range_start,omitempty"` + PollRangeEnd *FlexibleTime `gorm:"column:poll_range_end;type:varchar(255);index" json:"poll_range_end,omitempty"` + KbID string `gorm:"column:kb_id;size:32;not null;index" json:"kb_id"` BaseModel } diff --git a/internal/entity/flexible_time.go b/internal/entity/flexible_time.go new file mode 100644 index 0000000000..e556a4464e --- /dev/null +++ b/internal/entity/flexible_time.go @@ -0,0 +1,158 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package entity + +import ( + "database/sql/driver" + "encoding/json" + "fmt" + "log" + "net/mail" + "regexp" + "strings" + "time" +) + +// compactOffsetPattern expands the legacy "+0000" style offsets to "+00:00". +var compactOffsetPattern = regexp.MustCompile(`([+-]\d{2})(\d{2})$`) + +// flexibleTimeLayouts covers every combination of the timestamp spellings +// persisted by the Python backend (ISO strings in varchar columns) and the Go +// backend (native time.Time and the UTC ISO strings written by Value): T or +// space separator, optional fractional seconds, and optional zone. +var flexibleTimeLayouts = []string{ + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z07:00", + "2006-01-02 15:04:05.999999999Z07:00", + "2006-01-02 15:04:05Z07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999", + "2006-01-02T15:04:05", + "2006-01-02 15:04:05", + time.RFC3339Nano, + time.RFC3339, +} + +// FlexibleTime scans time values from both native time.Time columns and the +// varchar timestamp strings the Python backend writes, while still +// serializing to JSON like time.Time. +type FlexibleTime time.Time + +// Scan implements sql.Scanner. +func (f *FlexibleTime) Scan(value any) error { + if f == nil { + return fmt.Errorf("cannot scan into nil FlexibleTime") + } + if value == nil { + *f = FlexibleTime{} + return nil + } + switch typed := value.(type) { + case time.Time: + *f = FlexibleTime(typed) + return nil + case []byte: + f.parse(string(typed)) + return nil + case string: + f.parse(typed) + return nil + } + return fmt.Errorf("cannot scan %T into FlexibleTime", value) +} + +// parse parses a persisted timestamp string. A value that matches no known +// format falls back to the zero time instead of an error: the sync_logs poll +// waterline is a varchar shared with the Python backend, and one +// legacy/unexpected timestamp must not wedge the syncer by making ClaimTask +// or ListDueTasks fail forever. The next successful sync rewrites the +// waterline. +func (f *FlexibleTime) parse(value string) { + original := value + value = strings.TrimSpace(value) + // RFC 5322 headers (for example raw Gmail Date values) keep the compact + // "+0000" zone spelling, so try the standard library parser first. + if value != "" { + if parsed, err := mail.ParseDate(value); err == nil { + *f = FlexibleTime(parsed) + return + } + } + if value != "" { + last := value[len(value)-1] + if last == 'Z' || last == 'z' { + value = value[:len(value)-1] + "+00:00" + } + } + value = compactOffsetPattern.ReplaceAllString(value, "$1:$2") + for _, layout := range flexibleTimeLayouts { + if parsed, err := time.Parse(layout, value); err == nil { + *f = FlexibleTime(parsed) + return + } + } + if strings.TrimSpace(original) != "" { + log.Printf("flexible_time: cannot parse %q as time, falling back to zero time", original) + } + *f = FlexibleTime{} +} + +// Value implements driver.Valuer. The shared sync_logs columns are varchar +// and the metadata DSN renders time.Time in the local zone, so an explicit +// UTC ISO string is written to keep the stored instant unambiguous. The +// layout mirrors Python's DateTimeTzField.db_value (datetime.isoformat on a +// UTC value): 6-digit microseconds and a "+00:00" offset. +func (f FlexibleTime) Value() (driver.Value, error) { + value := time.Time(f) + if value.IsZero() { + return nil, nil + } + return value.UTC().Format("2006-01-02T15:04:05.999999-07:00"), nil +} + +// Time returns the underlying time value. +func (f FlexibleTime) Time() time.Time { + return time.Time(f) +} + +// NewFlexibleTime wraps a time.Time pointer, returning nil for nil input. +func NewFlexibleTime(value *time.Time) *FlexibleTime { + if value == nil { + return nil + } + converted := FlexibleTime(*value) + return &converted +} + +// MarshalJSON serializes like time.Time (RFC3339). +func (f FlexibleTime) MarshalJSON() ([]byte, error) { + return json.Marshal(time.Time(f)) +} + +// UnmarshalJSON parses an RFC3339 timestamp. +func (f *FlexibleTime) UnmarshalJSON(data []byte) error { + var value *time.Time + if err := json.Unmarshal(data, &value); err != nil { + return err + } + if value == nil { + *f = FlexibleTime{} + return nil + } + *f = FlexibleTime(*value) + return nil +} diff --git a/internal/syncer/connector/mysql.go b/internal/syncer/connector/mysql.go new file mode 100644 index 0000000000..e9172e2fae --- /dev/null +++ b/internal/syncer/connector/mysql.go @@ -0,0 +1,594 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package connector + +import ( + "context" + "crypto/md5" + "database/sql" + "database/sql/driver" + "encoding/hex" + "fmt" + "io" + "regexp" + "strings" + "time" + + "github.com/go-sql-driver/mysql" +) + +const defaultMySQLBatchSize = 32 + +// MySQLConnector imports MySQL rows as documents. +// +// It mirrors the Python RDBMSConnector's MySQL dialect: a custom SQL query +// runs verbatim, otherwise every table is loaded. Rows become documents whose +// content is built from the configured content columns (or every column), the +// id column (or an MD5 of the content) forms the stable document id, and the +// timestamp column drives incremental sync and the document update time. +type MySQLConnector struct { + host string + port int + database string + query string + contentColumns []string + metadataColumns []string + idColumn string + timestampColumn string + batchSize int + username string + password string + + openDB func(dsn string) (*sql.DB, error) +} + +// NewMySQLConnector creates a MySQL connector from Python-compatible config. +func NewMySQLConnector(config map[string]any) (*MySQLConnector, error) { + credentials, _ := config["credentials"].(map[string]any) + connector := &MySQLConnector{ + host: strings.TrimSpace(stringConfig(config["host"])), + port: configInt(config["port"], 3306), + database: strings.TrimSpace(stringConfig(config["database"])), + idColumn: strings.TrimSpace(stringConfig(config["id_column"])), + timestampColumn: strings.TrimSpace(stringConfig(config["timestamp_column"])), + batchSize: configInt(config["batch_size"], defaultMySQLBatchSize), + username: strings.TrimSpace(stringConfig(credentials["username"])), + password: stringConfig(credentials["password"]), + openDB: func(dsn string) (*sql.DB, error) { + return sql.Open("mysql", dsn) + }, + } + connector.query = connector.sanitizeQuery(stringConfig(config["query"])) + connector.contentColumns = connector.splitColumns(config["content_columns"]) + connector.metadataColumns = connector.splitColumns(config["metadata_columns"]) + return connector, nil +} + +// Validate validates MySQL connector settings and credentials. +func (c *MySQLConnector) Validate(ctx context.Context) error { + if c == nil { + return fmt.Errorf("mysql connector is nil") + } + if c.username == "" { + return fmt.Errorf("RDBMS (mysql): missing username") + } + if c.host == "" { + return fmt.Errorf("Database host is required") + } + if c.database == "" { + return fmt.Errorf("Database name is required") + } + if c.batchSize <= 0 { + return fmt.Errorf("batch_size must be a positive integer") + } + db, err := c.open() + if err != nil { + return fmt.Errorf("Failed to connect to MySQL: %w", err) + } + defer db.Close() + var one int + if err := db.QueryRowContext(ctx, "SELECT 1").Scan(&one); err != nil { + return fmt.Errorf("Failed to connect to MySQL: %w", err) + } + return nil +} + +// OpenSync opens one MySQL sync session. +func (c *MySQLConnector) OpenSync(ctx context.Context, request SyncRequest) (SyncSession, error) { + db, err := c.open() + if err != nil { + return nil, err + } + bases, err := c.baseQueries(ctx, db) + if err != nil { + db.Close() + return nil, err + } + queries := c.buildSyncQueries(bases, request) + return &mysqlSyncSession{connector: c, db: db, queries: queries, batchSize: c.batchSize}, nil +} + +// OpenPrune opens one complete MySQL prune snapshot session. +func (c *MySQLConnector) OpenPrune(ctx context.Context, request PruneRequest) (PruneSession, error) { + db, err := c.open() + if err != nil { + return nil, err + } + bases, err := c.baseQueries(ctx, db) + if err != nil { + db.Close() + return nil, err + } + queries := make([]string, 0, len(bases)) + for _, base := range bases { + queries = append(queries, c.buildSlimQuery(base)) + } + return &mysqlPruneSession{connector: c, db: db, queries: queries, batchSize: c.batchSize}, nil +} + +// open builds a MySQL connection with Python-compatible settings. +func (c *MySQLConnector) open() (*sql.DB, error) { + cfg := mysql.NewConfig() + cfg.User = c.username + cfg.Passwd = c.password + cfg.Net = "tcp" + cfg.Addr = fmt.Sprintf("%s:%d", c.host, c.port) + cfg.DBName = c.database + cfg.Params = map[string]string{"charset": "utf8mb4"} + cfg.ParseTime = true + cfg.Loc = time.UTC + return c.openDB(cfg.FormatDSN()) +} + +// baseQueries returns the configured query or a SELECT per table. +func (c *MySQLConnector) baseQueries(ctx context.Context, db *sql.DB) ([]string, error) { + if c.query != "" { + return []string{c.query}, nil + } + rows, err := db.QueryContext(ctx, "SHOW TABLES") + if err != nil { + return nil, err + } + defer rows.Close() + var tables []string + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return nil, err + } + tables = append(tables, name) + } + if err := rows.Err(); err != nil { + return nil, err + } + queries := make([]string, 0, len(tables)) + for _, table := range tables { + queries = append(queries, fmt.Sprintf("SELECT * FROM %s", table)) + } + return queries, nil +} + +// buildSyncQueries applies the incremental window when a timestamp column exists. +func (c *MySQLConnector) buildSyncQueries(bases []string, request SyncRequest) []string { + var start, end *time.Time + if !request.FromBeginning { + start = request.WindowStart + end = &request.WindowEnd + } + if c.timestampColumn == "" || (start == nil && end == nil) { + return bases + } + queries := make([]string, 0, len(bases)) + for _, base := range bases { + queries = append(queries, c.buildTimeFilteredQuery(base, start, end)) + } + return queries +} + +// buildTimeFilteredQuery wraps the base query and appends timestamp bounds. +func (c *MySQLConnector) buildTimeFilteredQuery(base string, start, end *time.Time) string { + conditions := []string{} + if start != nil { + conditions = append(conditions, fmt.Sprintf("ragflow_src.%s >= %s", c.timestampColumn, c.formatDatetime(*start))) + } + if end != nil { + conditions = append(conditions, fmt.Sprintf("ragflow_src.%s <= %s", c.timestampColumn, c.formatDatetime(*end))) + } + query := c.wrapQuery(base) + if len(conditions) > 0 { + query = query + " WHERE " + strings.Join(conditions, " AND ") + } + return query +} + +// buildSlimQuery selects only the columns needed to identify documents. +func (c *MySQLConnector) buildSlimQuery(base string) string { + columns := []string{} + if c.idColumn != "" { + columns = []string{c.idColumn} + } else { + columns = c.contentColumns + } + if len(columns) == 0 { + return c.wrapQuery(base) + } + selects := make([]string, 0, len(columns)) + for _, column := range columns { + selects = append(selects, fmt.Sprintf("ragflow_src.%s", column)) + } + return fmt.Sprintf("SELECT %s FROM (%s) AS ragflow_src", strings.Join(selects, ", "), c.stripOrderBy(base)) +} + +// wrapQuery wraps the base query as a derived table named ragflow_src. +func (c *MySQLConnector) wrapQuery(base string) string { + return fmt.Sprintf("SELECT * FROM (%s) AS ragflow_src", c.stripOrderBy(base)) +} + +// stripOrderBy removes a trailing top-level ORDER BY clause. +func (c *MySQLConnector) stripOrderBy(query string) string { + pattern := regexp.MustCompile(`(?i)\border\s+by\b`) + cleaned := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(query), ";")) + matches := pattern.FindAllStringIndex(cleaned, -1) + for i := len(matches) - 1; i >= 0; i-- { + prefix := cleaned[:matches[i][0]] + if strings.Count(prefix, "(") == strings.Count(prefix, ")") { + return strings.TrimSpace(prefix) + } + } + return cleaned +} + +// formatDatetime renders a UTC time as a MySQL datetime literal. +func (c *MySQLConnector) formatDatetime(value time.Time) string { + return "'" + value.UTC().Format("2006-01-02 15:04:05") + "'" +} + +// scanRow scans the current row into an ordered column map. +func (c *MySQLConnector) scanRow(rows *sql.Rows) (map[string]any, []string, error) { + columns, err := rows.Columns() + if err != nil { + return nil, nil, err + } + values := make([]any, len(columns)) + pointers := make([]any, len(columns)) + for i := range values { + pointers[i] = &values[i] + } + if err := rows.Scan(pointers...); err != nil { + return nil, nil, err + } + row := make(map[string]any, len(columns)) + for i, column := range columns { + row[column] = c.normalizeValue(values[i]) + } + return row, columns, nil +} + +// normalizeValue converts driver byte slices to strings. +func (c *MySQLConnector) normalizeValue(value any) any { + if bytes, ok := value.([]byte); ok { + return string(bytes) + } + if _, ok := value.(time.Time); ok { + return value + } + if valuer, ok := value.(driver.Valuer); ok { + if converted, err := valuer.Value(); err == nil { + return c.normalizeValue(converted) + } + } + return value +} + +// contentColumnsForRow resolves the content columns for a row, excluding the +// structural id and timestamp columns when no content columns are configured. +func (c *MySQLConnector) contentColumnsForRow(row map[string]any, orderedColumns []string) []string { + if len(c.contentColumns) > 0 { + return c.contentColumns + } + excluded := map[string]bool{} + if c.idColumn != "" { + excluded[c.idColumn] = true + } + if c.timestampColumn != "" { + excluded[c.timestampColumn] = true + } + columns := make([]string, 0, len(orderedColumns)) + for _, column := range orderedColumns { + if _, ok := row[column]; ok && !excluded[column] { + columns = append(columns, column) + } + } + return columns +} + +// buildContent renders the document content from the resolved content columns. +func (c *MySQLConnector) buildContent(row map[string]any, columns []string) string { + parts := []string{} + for _, column := range columns { + value, ok := row[column] + if !ok || value == nil { + continue + } + parts = append(parts, fmt.Sprintf("【%s】:\n%s", column, c.renderValue(value))) + } + return strings.Join(parts, "\n\n") +} + +// buildDocumentID derives the stable document id, matching the Python format +// "mysql::" with an MD5 content fallback. +func (c *MySQLConnector) buildDocumentID(row map[string]any, orderedColumns []string) string { + if c.idColumn != "" { + if value, ok := row[c.idColumn]; ok && value != nil { + return fmt.Sprintf("mysql:%s:%s", c.database, fmt.Sprint(value)) + } + } + content := c.buildContent(row, c.contentColumnsForRow(row, orderedColumns)) + sum := md5.Sum([]byte(content)) + return fmt.Sprintf("mysql:%s:%s", c.database, hex.EncodeToString(sum[:])) +} + +// rowToSourceDocument converts a database row into the syncer model. +func (c *MySQLConnector) rowToSourceDocument(row map[string]any, orderedColumns []string) (SourceDocument, bool) { + contentColumns := c.contentColumnsForRow(row, orderedColumns) + content := c.buildContent(row, contentColumns) + + metadata := map[string]any{} + for _, column := range c.metadataColumns { + value, ok := row[column] + if !ok || value == nil { + continue + } + metadata[column] = c.formatMetadataValue(value) + } + + updatedAt := time.Now().UTC() + if c.timestampColumn != "" { + if ts, ok := row[c.timestampColumn].(time.Time); ok { + updatedAt = ts.UTC() + } + } + + semanticID := "database_record" + if len(contentColumns) > 0 { + if value, ok := row[contentColumns[0]]; ok && value != nil { + semanticID = strings.TrimSpace(strings.ReplaceAll(strings.ReplaceAll(fmt.Sprint(value), "\n", " "), "\r", " ")) + if semanticID == "" { + semanticID = "database_record" + } else if len(semanticID) > 100 { + semanticID = semanticID[:100] + } + } + } + + sourceID := c.buildDocumentID(row, orderedColumns) + blob := []byte(content) + return SourceDocument{ + SourceID: sourceID, + SemanticIdentifier: semanticID, + Extension: ".txt", + Blob: blob, + UpdatedAt: updatedAt, + SizeBytes: int64(len(blob)), + Metadata: metadata, + Fingerprint: stableFingerprint(map[string]any{ + "id": sourceID, + "content": content, + "metadata": metadata, + }), + }, true +} + +// renderValue formats a row value for document content. +func (c *MySQLConnector) renderValue(value any) string { + if typed, ok := value.(time.Time); ok { + return typed.Format("2006-01-02 15:04:05") + } + return fmt.Sprint(value) +} + +// formatMetadataValue formats a row value for metadata, mirroring Python's +// isoformat for datetimes and string rendering otherwise. +func (c *MySQLConnector) formatMetadataValue(value any) string { + if typed, ok := value.(time.Time); ok { + return typed.Format(time.RFC3339) + } + return fmt.Sprint(value) +} + +// sanitizeQuery tolerates queries pasted from a markdown code fence. +func (c *MySQLConnector) sanitizeQuery(raw string) string { + fenceLanguages := map[string]bool{"sql": true, "tsql": true, "t-sql": true, "mssql": true, "mysql": true, "postgresql": true, "psql": true} + query := strings.TrimSpace(raw) + if query == "" { + return "" + } + if strings.HasPrefix(query, "```") { + query = query[3:] + if strings.HasSuffix(query, "```") { + query = query[:len(query)-3] + } + query = strings.TrimSpace(query) + } + if head, tail, found := strings.Cut(query, "\n"); found { + if fenceLanguages[strings.ToLower(strings.TrimSpace(head))] { + query = strings.TrimSpace(tail) + } + } + return query +} + +// splitColumns parses a comma-separated string or list column config. +func (c *MySQLConnector) splitColumns(value any) []string { + switch typed := value.(type) { + case string: + parts := strings.Split(typed, ",") + columns := make([]string, 0, len(parts)) + for _, part := range parts { + if column := strings.TrimSpace(part); column != "" { + columns = append(columns, column) + } + } + return columns + case []any: + columns := make([]string, 0, len(typed)) + for _, item := range typed { + if column := strings.TrimSpace(stringConfig(item)); column != "" { + columns = append(columns, column) + } + } + return columns + } + return nil +} + +type mysqlSyncSession struct { + connector *MySQLConnector + db *sql.DB + queries []string + queryIndex int + rows *sql.Rows + batchSize int +} + +// NextBatch returns the next MySQL document batch. +func (s *mysqlSyncSession) NextBatch(ctx context.Context) (SyncBatch, error) { + documents := make([]SourceDocument, 0, s.batchSize) + for len(documents) < s.batchSize { + if s.rows == nil { + if s.queryIndex >= len(s.queries) { + if len(documents) == 0 { + return SyncBatch{}, io.EOF + } + break + } + if err := s.openNextQuery(ctx); err != nil { + return SyncBatch{}, err + } + } + if !s.rows.Next() { + if err := s.rows.Err(); err != nil { + s.closeRows() + return SyncBatch{}, err + } + s.closeRows() + continue + } + row, columns, err := s.connector.scanRow(s.rows) + if err != nil { + // Skip rows that fail to convert (mirrors Python). + continue + } + if doc, ok := s.connector.rowToSourceDocument(row, columns); ok { + documents = append(documents, doc) + } + } + return SyncBatch{Documents: documents}, nil +} + +// Close closes the MySQL sync session. +func (s *mysqlSyncSession) Close() error { + s.closeRows() + return s.db.Close() +} + +// openNextQuery runs the next base query. +func (s *mysqlSyncSession) openNextQuery(ctx context.Context) error { + query := s.queries[s.queryIndex] + s.queryIndex++ + rows, err := s.db.QueryContext(ctx, query) + if err != nil { + return fmt.Errorf("MySQL query failed: %w", err) + } + s.rows = rows + return nil +} + +// closeRows releases the current result set. +func (s *mysqlSyncSession) closeRows() { + if s.rows != nil { + s.rows.Close() + s.rows = nil + } +} + +type mysqlPruneSession struct { + connector *MySQLConnector + db *sql.DB + queries []string + queryIndex int + rows *sql.Rows + batchSize int +} + +// NextBatch returns the next MySQL prune snapshot batch. +func (s *mysqlPruneSession) NextBatch(ctx context.Context) (PruneBatch, error) { + documents := make([]SlimDocument, 0, s.batchSize) + for len(documents) < s.batchSize { + if s.rows == nil { + if s.queryIndex >= len(s.queries) { + if len(documents) == 0 { + return PruneBatch{}, io.EOF + } + break + } + if err := s.openNextQuery(ctx); err != nil { + return PruneBatch{}, err + } + } + if !s.rows.Next() { + if err := s.rows.Err(); err != nil { + s.closeRows() + return PruneBatch{}, err + } + s.closeRows() + continue + } + row, columns, err := s.connector.scanRow(s.rows) + if err != nil { + continue + } + documents = append(documents, SlimDocument{SourceID: s.connector.buildDocumentID(row, columns)}) + } + return PruneBatch{Documents: documents}, nil +} + +// Close closes the MySQL prune session. +func (s *mysqlPruneSession) Close() error { + s.closeRows() + return s.db.Close() +} + +// openNextQuery runs the next slim query. +func (s *mysqlPruneSession) openNextQuery(ctx context.Context) error { + query := s.queries[s.queryIndex] + s.queryIndex++ + rows, err := s.db.QueryContext(ctx, query) + if err != nil { + return fmt.Errorf("MySQL query failed: %w", err) + } + s.rows = rows + return nil +} + +// closeRows releases the current result set. +func (s *mysqlPruneSession) closeRows() { + if s.rows != nil { + s.rows.Close() + s.rows = nil + } +} diff --git a/internal/syncer/connector/mysql_test.go b/internal/syncer/connector/mysql_test.go new file mode 100644 index 0000000000..c1c25dd30d --- /dev/null +++ b/internal/syncer/connector/mysql_test.go @@ -0,0 +1,351 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package connector + +import ( + "context" + "crypto/md5" + "database/sql" + "encoding/hex" + "errors" + "io" + "regexp" + "strings" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" +) + +func newFixtureMySQLConnector(t *testing.T, config map[string]any, expect func(mock sqlmock.Sqlmock)) *MySQLConnector { + t.Helper() + if config == nil { + config = map[string]any{ + "host": "127.0.0.1", + "port": "3306", + "database": "mydb", + "credentials": map[string]any{ + "username": "root", + "password": "secret", + }, + } + } + connector, err := NewMySQLConnector(config) + if err != nil { + t.Fatalf("NewMySQLConnector failed: %v", err) + } + db, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New failed: %v", err) + } + t.Cleanup(func() { + db.Close() + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } + }) + connector.openDB = func(dsn string) (*sql.DB, error) { + return db, nil + } + if expect != nil { + expect(mock) + } + return connector +} + +// TestMySQLConnectorOpenSyncCustomQuery verifies a custom query produces documents. +func TestMySQLConnectorOpenSyncCustomQuery(t *testing.T) { + query := "SELECT * FROM products WHERE status = 'active'" + updatedAt := mustTime(t, "2026-01-02T03:04:05Z") + connector := newFixtureMySQLConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "3306", + "database": "mydb", + "query": query, + "content_columns": "title,description", + "metadata_columns": "id,category,updated_at", + "id_column": "id", + "timestamp_column": "updated_at", + "credentials": map[string]any{ + "username": "root", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta(query)).WillReturnRows( + sqlmock.NewRows([]string{"id", "title", "description", "category", "updated_at"}). + AddRow(7, "Hello/World", "Some body", "news", updatedAt), + ) + }) + + session, err := connector.OpenSync(context.Background(), SyncRequest{FromBeginning: true}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 { + t.Fatalf("documents len = %d, want 1", len(batch.Documents)) + } + doc := batch.Documents[0] + if doc.SourceID != "mysql:mydb:7" { + t.Fatalf("source id = %q", doc.SourceID) + } + if doc.SemanticIdentifier != "Hello/World" { + t.Fatalf("semantic identifier = %q", doc.SemanticIdentifier) + } + if doc.Extension != ".txt" { + t.Fatalf("extension = %q", doc.Extension) + } + blob := string(doc.Blob) + if !strings.Contains(blob, "【title】:\nHello/World") || !strings.Contains(blob, "【description】:\nSome body") { + t.Fatalf("blob = %q", blob) + } + if !doc.UpdatedAt.Equal(updatedAt) { + t.Fatalf("updated at = %s", doc.UpdatedAt) + } + if doc.Metadata["category"] != "news" || doc.Metadata["id"] != "7" { + t.Fatalf("metadata = %v", doc.Metadata) + } + if doc.Metadata["updated_at"] != updatedAt.Format(time.RFC3339) { + t.Fatalf("metadata updated_at = %v", doc.Metadata["updated_at"]) + } + if doc.Fingerprint == "" { + t.Fatalf("fingerprint is empty") + } + if _, err = session.NextBatch(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("NextBatch EOF = %v", err) + } +} + +// TestMySQLConnectorOpenSyncIncrementalWindow verifies the timestamp filter SQL. +func TestMySQLConnectorOpenSyncIncrementalWindow(t *testing.T) { + connector := newFixtureMySQLConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": 3306, + "database": "mydb", + "query": "SELECT * FROM products", + "timestamp_column": "updated_at", + "credentials": map[string]any{ + "username": "root", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + expected := "SELECT * FROM (SELECT * FROM products) AS ragflow_src " + + "WHERE ragflow_src.updated_at >= '2026-01-01 00:00:00' AND ragflow_src.updated_at <= '2026-01-02 00:00:00'" + mock.ExpectQuery(regexp.QuoteMeta(expected)).WillReturnRows( + sqlmock.NewRows([]string{"updated_at"}).AddRow(mustTime(t, "2026-01-01T12:00:00Z")), + ) + }) + + start := mustTime(t, "2026-01-01T00:00:00Z") + end := mustTime(t, "2026-01-02T00:00:00Z") + session, err := connector.OpenSync(context.Background(), SyncRequest{WindowStart: &start, WindowEnd: end}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 { + t.Fatalf("documents len = %d, want 1", len(batch.Documents)) + } +} + +// TestMySQLConnectorOpenSyncAllTables verifies SHOW TABLES expands to per-table queries. +func TestMySQLConnectorOpenSyncAllTables(t *testing.T) { + connector := newFixtureMySQLConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": 3306, + "database": "mydb", + "id_column": "id", + "credentials": map[string]any{ + "username": "root", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta("SHOW TABLES")).WillReturnRows( + sqlmock.NewRows([]string{"Tables_in_mydb"}).AddRow("products").AddRow("orders"), + ) + mock.ExpectQuery(regexp.QuoteMeta("SELECT * FROM products")).WillReturnRows( + sqlmock.NewRows([]string{"id", "name"}).AddRow(1, "Product"), + ) + mock.ExpectQuery(regexp.QuoteMeta("SELECT * FROM orders")).WillReturnRows( + sqlmock.NewRows([]string{"id", "name"}).AddRow(2, "Order"), + ) + }) + + session, err := connector.OpenSync(context.Background(), SyncRequest{FromBeginning: true}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + var ids []string + for { + batch, err := session.NextBatch(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + for _, doc := range batch.Documents { + ids = append(ids, doc.SourceID) + } + } + if len(ids) != 2 || ids[0] != "mysql:mydb:1" || ids[1] != "mysql:mydb:2" { + t.Fatalf("source ids = %v", ids) + } +} + +// TestMySQLConnectorOpenPrune verifies the slim query and Python-compatible IDs. +func TestMySQLConnectorOpenPrune(t *testing.T) { + connector := newFixtureMySQLConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": 3306, + "database": "mydb", + "query": "SELECT * FROM products", + "content_columns": "title,description", + "id_column": "id", + "credentials": map[string]any{ + "username": "root", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + expected := "SELECT ragflow_src.id FROM (SELECT * FROM products) AS ragflow_src" + mock.ExpectQuery(regexp.QuoteMeta(expected)).WillReturnRows( + sqlmock.NewRows([]string{"id"}).AddRow(3).AddRow(4), + ) + }) + + session, err := connector.OpenPrune(context.Background(), PruneRequest{}) + if err != nil { + t.Fatalf("OpenPrune failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 2 || + batch.Documents[0].SourceID != "mysql:mydb:3" || + batch.Documents[1].SourceID != "mysql:mydb:4" { + t.Fatalf("slim documents = %+v", batch.Documents) + } + if _, err = session.NextBatch(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("NextBatch EOF = %v", err) + } +} + +// TestMySQLConnectorMD5FallbackID verifies the content-hash document id. +func TestMySQLConnectorMD5FallbackID(t *testing.T) { + connector := newFixtureMySQLConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": 3306, + "database": "mydb", + "query": "SELECT * FROM products", + "content_columns": "title,description", + "credentials": map[string]any{ + "username": "root", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta("SELECT * FROM products")).WillReturnRows( + sqlmock.NewRows([]string{"title", "description"}).AddRow("Hello", "World"), + ) + }) + + session, err := connector.OpenSync(context.Background(), SyncRequest{FromBeginning: true}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 { + t.Fatalf("documents len = %d, want 1", len(batch.Documents)) + } + content := "【title】:\nHello\n\n【description】:\nWorld" + sum := md5.Sum([]byte(content)) + want := "mysql:mydb:" + hex.EncodeToString(sum[:]) + if batch.Documents[0].SourceID != want { + t.Fatalf("source id = %q, want %q", batch.Documents[0].SourceID, want) + } +} + +// TestMySQLConnectorValidate verifies the SELECT 1 probe and failure paths. +func TestMySQLConnectorValidate(t *testing.T) { + connector := newFixtureMySQLConnector(t, nil, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta("SELECT 1")).WillReturnRows( + sqlmock.NewRows([]string{"1"}).AddRow(1), + ) + }) + if err := connector.Validate(context.Background()); err != nil { + t.Fatalf("Validate failed: %v", err) + } + + missing := newFixtureMySQLConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": 3306, + "database": "mydb", + "credentials": map[string]any{ + "username": "", + "password": "secret", + }, + }, nil) + if err := missing.Validate(context.Background()); err == nil || !strings.Contains(err.Error(), "username") { + t.Fatalf("Validate error = %v", err) + } +} + +// TestMySQLConnectorSanitizeQuery verifies markdown fence tolerance. +func TestMySQLConnectorSanitizeQuery(t *testing.T) { + connector := &MySQLConnector{} + cases := []struct { + in string + want string + }{ + {in: "SELECT * FROM t", want: "SELECT * FROM t"}, + {in: "```sql\nSELECT * FROM t\n```", want: "SELECT * FROM t"}, + {in: "sql\nSELECT * FROM t", want: "SELECT * FROM t"}, + {in: " ", want: ""}, + } + for _, tc := range cases { + if got := connector.sanitizeQuery(tc.in); got != tc.want { + t.Fatalf("sanitizeQuery(%q) = %q, want %q", tc.in, got, tc.want) + } + } +} + +// TestMySQLConnectorStripOrderBy verifies trailing top-level ORDER BY removal. +func TestMySQLConnectorStripOrderBy(t *testing.T) { + connector := &MySQLConnector{} + cases := []struct { + in string + want string + }{ + {in: "SELECT * FROM t ORDER BY id", want: "SELECT * FROM t"}, + {in: "SELECT ROW_NUMBER() OVER (ORDER BY id) FROM t", want: "SELECT ROW_NUMBER() OVER (ORDER BY id) FROM t"}, + {in: "SELECT * FROM t", want: "SELECT * FROM t"}, + } + for _, tc := range cases { + if got := connector.stripOrderBy(tc.in); got != tc.want { + t.Fatalf("stripOrderBy(%q) = %q, want %q", tc.in, got, tc.want) + } + } +} diff --git a/internal/syncer/connector/postgresql.go b/internal/syncer/connector/postgresql.go new file mode 100644 index 0000000000..a0dd356a0a --- /dev/null +++ b/internal/syncer/connector/postgresql.go @@ -0,0 +1,621 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package connector + +import ( + "context" + "crypto/md5" + "database/sql" + "database/sql/driver" + "encoding/hex" + "fmt" + "io" + "net/url" + "regexp" + "strconv" + "strings" + "time" + + _ "github.com/jackc/pgx/v5/stdlib" +) + +const ( + defaultPostgresBatchSize = 32 + defaultPostgresConnectTimeout = 30 +) + +// PostgreSQLConnector imports PostgreSQL rows as documents. +// +// It mirrors the Python RDBMSConnector's PostgreSQL dialect: a custom SQL +// query runs verbatim, otherwise every table in the public schema is loaded. +// Rows become documents whose content is built from the configured content +// columns (or every column), the id column (or an MD5 of the content) forms +// the stable document id, and the timestamp column drives incremental sync +// and the document update time. +type PostgreSQLConnector struct { + host string + port int + database string + query string + contentColumns []string + metadataColumns []string + idColumn string + timestampColumn string + batchSize int + username string + password string + sslmode string + connectTimeout int + + openDB func(dsn string) (*sql.DB, error) +} + +// NewPostgreSQLConnector creates a PostgreSQL connector from Python-compatible config. +func NewPostgreSQLConnector(config map[string]any) (*PostgreSQLConnector, error) { + credentials, _ := config["credentials"].(map[string]any) + connector := &PostgreSQLConnector{ + host: strings.TrimSpace(stringConfig(config["host"])), + port: configInt(config["port"], 5432), + database: strings.TrimSpace(stringConfig(config["database"])), + idColumn: strings.TrimSpace(stringConfig(config["id_column"])), + timestampColumn: strings.TrimSpace(stringConfig(config["timestamp_column"])), + batchSize: configInt(config["batch_size"], defaultPostgresBatchSize), + username: strings.TrimSpace(stringConfig(credentials["username"])), + password: stringConfig(credentials["password"]), + sslmode: strings.TrimSpace(stringConfig(config["sslmode"])), + connectTimeout: configInt(config["connect_timeout"], defaultPostgresConnectTimeout), + openDB: func(dsn string) (*sql.DB, error) { + return sql.Open("pgx", dsn) + }, + } + if connector.sslmode == "" { + connector.sslmode = "prefer" + } + connector.query = connector.sanitizeQuery(stringConfig(config["query"])) + connector.contentColumns = connector.splitColumns(config["content_columns"]) + connector.metadataColumns = connector.splitColumns(config["metadata_columns"]) + return connector, nil +} + +// Validate validates PostgreSQL connector settings and credentials. +func (c *PostgreSQLConnector) Validate(ctx context.Context) error { + if c == nil { + return fmt.Errorf("postgresql connector is nil") + } + if c.username == "" { + return fmt.Errorf("RDBMS (postgresql): missing username") + } + if c.host == "" { + return fmt.Errorf("Database host is required") + } + if c.database == "" { + return fmt.Errorf("Database name is required") + } + if c.batchSize <= 0 { + return fmt.Errorf("batch_size must be a positive integer") + } + db, err := c.open() + if err != nil { + return fmt.Errorf("Failed to connect to PostgreSQL: %w", err) + } + defer db.Close() + var one int + if err := db.QueryRowContext(ctx, "SELECT 1").Scan(&one); err != nil { + return fmt.Errorf("Failed to connect to PostgreSQL: %w", err) + } + return nil +} + +// OpenSync opens one PostgreSQL sync session. +func (c *PostgreSQLConnector) OpenSync(ctx context.Context, request SyncRequest) (SyncSession, error) { + db, err := c.open() + if err != nil { + return nil, err + } + bases, err := c.baseQueries(ctx, db) + if err != nil { + db.Close() + return nil, err + } + queries := c.buildSyncQueries(bases, request) + return &postgresSyncSession{connector: c, db: db, queries: queries, batchSize: c.batchSize}, nil +} + +// OpenPrune opens one complete PostgreSQL prune snapshot session. +func (c *PostgreSQLConnector) OpenPrune(ctx context.Context, request PruneRequest) (PruneSession, error) { + db, err := c.open() + if err != nil { + return nil, err + } + bases, err := c.baseQueries(ctx, db) + if err != nil { + db.Close() + return nil, err + } + queries := make([]string, 0, len(bases)) + for _, base := range bases { + queries = append(queries, c.buildSlimQuery(base)) + } + return &postgresPruneSession{connector: c, db: db, queries: queries, batchSize: c.batchSize}, nil +} + +// open builds a PostgreSQL connection from the connector settings. The DSN +// carries connector-controlled sslmode (default prefer, matching Python's +// psycopg2) and a finite connect_timeout so an unreachable host cannot hang +// a sync worker. +func (c *PostgreSQLConnector) open() (*sql.DB, error) { + dsn := url.URL{ + Scheme: "postgres", + User: url.UserPassword(c.username, c.password), + Host: fmt.Sprintf("%s:%d", c.host, c.port), + Path: "/" + url.PathEscape(c.database), + } + query := dsn.Query() + query.Set("sslmode", c.sslmode) + query.Set("connect_timeout", strconv.Itoa(c.connectTimeout)) + dsn.RawQuery = query.Encode() + return c.openDB(dsn.String()) +} + +// baseQueries returns the configured query or a SELECT per table. +func (c *PostgreSQLConnector) baseQueries(ctx context.Context, db *sql.DB) ([]string, error) { + if c.query != "" { + return []string{c.query}, nil + } + rows, err := db.QueryContext(ctx, "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE'") + if err != nil { + return nil, err + } + defer rows.Close() + var tables []string + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return nil, err + } + tables = append(tables, name) + } + if err := rows.Err(); err != nil { + return nil, err + } + queries := make([]string, 0, len(tables)) + for _, table := range tables { + queries = append(queries, fmt.Sprintf("SELECT * FROM \"public\".%s", quotePostgresIdentifier(table))) + } + return queries, nil +} + +// quotePostgresIdentifier double-quotes an identifier for PostgreSQL, escaping +// any embedded double quotes, so catalog-discovered names with mixed case or +// special characters survive the unquoted lowercase folding. +func quotePostgresIdentifier(name string) string { + return `"` + strings.ReplaceAll(name, `"`, `""`) + `"` +} + +// buildSyncQueries applies the incremental window when a timestamp column exists. +func (c *PostgreSQLConnector) buildSyncQueries(bases []string, request SyncRequest) []string { + var start, end *time.Time + if !request.FromBeginning { + start = request.WindowStart + end = &request.WindowEnd + } + if c.timestampColumn == "" || (start == nil && end == nil) { + return bases + } + queries := make([]string, 0, len(bases)) + for _, base := range bases { + queries = append(queries, c.buildTimeFilteredQuery(base, start, end)) + } + return queries +} + +// buildTimeFilteredQuery wraps the base query and appends timestamp bounds. +func (c *PostgreSQLConnector) buildTimeFilteredQuery(base string, start, end *time.Time) string { + conditions := []string{} + if start != nil { + conditions = append(conditions, fmt.Sprintf("ragflow_src.%s >= %s", c.timestampColumn, c.formatDatetime(*start))) + } + if end != nil { + conditions = append(conditions, fmt.Sprintf("ragflow_src.%s <= %s", c.timestampColumn, c.formatDatetime(*end))) + } + query := c.wrapQuery(base) + if len(conditions) > 0 { + query = query + " WHERE " + strings.Join(conditions, " AND ") + } + return query +} + +// buildSlimQuery selects only the columns needed to identify documents. +func (c *PostgreSQLConnector) buildSlimQuery(base string) string { + columns := []string{} + if c.idColumn != "" { + columns = []string{c.idColumn} + } else { + columns = c.contentColumns + } + if len(columns) == 0 { + return c.wrapQuery(base) + } + selects := make([]string, 0, len(columns)) + for _, column := range columns { + selects = append(selects, fmt.Sprintf("ragflow_src.%s", column)) + } + return fmt.Sprintf("SELECT %s FROM (%s) AS ragflow_src", strings.Join(selects, ", "), c.stripOrderBy(base)) +} + +// wrapQuery wraps the base query as a derived table named ragflow_src. +func (c *PostgreSQLConnector) wrapQuery(base string) string { + return fmt.Sprintf("SELECT * FROM (%s) AS ragflow_src", c.stripOrderBy(base)) +} + +// stripOrderBy removes a trailing top-level ORDER BY clause. +func (c *PostgreSQLConnector) stripOrderBy(query string) string { + pattern := regexp.MustCompile(`(?i)\border\s+by\b`) + cleaned := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(query), ";")) + matches := pattern.FindAllStringIndex(cleaned, -1) + for i := len(matches) - 1; i >= 0; i-- { + prefix := cleaned[:matches[i][0]] + if strings.Count(prefix, "(") == strings.Count(prefix, ")") { + return strings.TrimSpace(prefix) + } + } + return cleaned +} + +// formatDatetime renders a UTC time as an ISO-8601 PostgreSQL literal. +func (c *PostgreSQLConnector) formatDatetime(value time.Time) string { + return "'" + value.UTC().Format(time.RFC3339Nano) + "'" +} + +// scanRow scans the current row into an ordered column map. +func (c *PostgreSQLConnector) scanRow(rows *sql.Rows) (map[string]any, []string, error) { + columns, err := rows.Columns() + if err != nil { + return nil, nil, err + } + values := make([]any, len(columns)) + pointers := make([]any, len(columns)) + for i := range values { + pointers[i] = &values[i] + } + if err := rows.Scan(pointers...); err != nil { + return nil, nil, err + } + row := make(map[string]any, len(columns)) + for i, column := range columns { + row[column] = c.normalizeValue(values[i]) + } + return row, columns, nil +} + +// normalizeValue converts driver-specific values into plain strings so +// content and metadata rendering stay dialect-agnostic. Byte slices (jsonb) +// and driver value types (numeric) become their text form; time.Time passes +// through untouched. +func (c *PostgreSQLConnector) normalizeValue(value any) any { + if bytes, ok := value.([]byte); ok { + return string(bytes) + } + if _, ok := value.(time.Time); ok { + return value + } + if valuer, ok := value.(driver.Valuer); ok { + if converted, err := valuer.Value(); err == nil { + return c.normalizeValue(converted) + } + } + return value +} + +// contentColumnsForRow resolves the content columns for a row, excluding the +// structural id and timestamp columns when no content columns are configured. +func (c *PostgreSQLConnector) contentColumnsForRow(row map[string]any, orderedColumns []string) []string { + if len(c.contentColumns) > 0 { + return c.contentColumns + } + excluded := map[string]bool{} + if c.idColumn != "" { + excluded[c.idColumn] = true + } + if c.timestampColumn != "" { + excluded[c.timestampColumn] = true + } + columns := make([]string, 0, len(orderedColumns)) + for _, column := range orderedColumns { + if _, ok := row[column]; ok && !excluded[column] { + columns = append(columns, column) + } + } + return columns +} + +// buildContent renders the document content from the resolved content columns. +func (c *PostgreSQLConnector) buildContent(row map[string]any, columns []string) string { + parts := []string{} + for _, column := range columns { + value, ok := row[column] + if !ok || value == nil { + continue + } + parts = append(parts, fmt.Sprintf("【%s】:\n%s", column, c.renderValue(value))) + } + return strings.Join(parts, "\n\n") +} + +// buildDocumentID derives the stable document id, matching the Python format +// "postgresql::" with an MD5 content fallback. +func (c *PostgreSQLConnector) buildDocumentID(row map[string]any, orderedColumns []string) string { + if c.idColumn != "" { + if value, ok := row[c.idColumn]; ok && value != nil { + return fmt.Sprintf("postgresql:%s:%s", c.database, fmt.Sprint(value)) + } + } + content := c.buildContent(row, c.contentColumnsForRow(row, orderedColumns)) + sum := md5.Sum([]byte(content)) + return fmt.Sprintf("postgresql:%s:%s", c.database, hex.EncodeToString(sum[:])) +} + +// rowToSourceDocument converts a database row into the syncer model. +func (c *PostgreSQLConnector) rowToSourceDocument(row map[string]any, orderedColumns []string) (SourceDocument, bool) { + contentColumns := c.contentColumnsForRow(row, orderedColumns) + content := c.buildContent(row, contentColumns) + + metadata := map[string]any{} + for _, column := range c.metadataColumns { + value, ok := row[column] + if !ok || value == nil { + continue + } + metadata[column] = c.formatMetadataValue(value) + } + + updatedAt := time.Now().UTC() + if c.timestampColumn != "" { + if ts, ok := row[c.timestampColumn].(time.Time); ok { + updatedAt = ts.UTC() + } + } + + semanticID := "database_record" + if len(contentColumns) > 0 { + if value, ok := row[contentColumns[0]]; ok && value != nil { + semanticID = strings.TrimSpace(strings.ReplaceAll(strings.ReplaceAll(fmt.Sprint(value), "\n", " "), "\r", " ")) + if semanticID == "" { + semanticID = "database_record" + } else if len(semanticID) > 100 { + semanticID = semanticID[:100] + } + } + } + + sourceID := c.buildDocumentID(row, orderedColumns) + blob := []byte(content) + return SourceDocument{ + SourceID: sourceID, + SemanticIdentifier: semanticID, + Extension: ".txt", + Blob: blob, + UpdatedAt: updatedAt, + SizeBytes: int64(len(blob)), + Metadata: metadata, + Fingerprint: stableFingerprint(map[string]any{ + "id": sourceID, + "content": content, + "metadata": metadata, + }), + }, true +} + +// renderValue formats a row value for document content. +func (c *PostgreSQLConnector) renderValue(value any) string { + if typed, ok := value.(time.Time); ok { + return typed.Format("2006-01-02 15:04:05") + } + return fmt.Sprint(value) +} + +// formatMetadataValue formats a row value for metadata, mirroring Python's +// isoformat for datetimes and string rendering otherwise. +func (c *PostgreSQLConnector) formatMetadataValue(value any) string { + if typed, ok := value.(time.Time); ok { + return typed.Format(time.RFC3339) + } + return fmt.Sprint(value) +} + +// sanitizeQuery tolerates queries pasted from a markdown code fence. +func (c *PostgreSQLConnector) sanitizeQuery(raw string) string { + fenceLanguages := map[string]bool{"sql": true, "tsql": true, "t-sql": true, "mssql": true, "mysql": true, "postgresql": true, "psql": true} + query := strings.TrimSpace(raw) + if query == "" { + return "" + } + if strings.HasPrefix(query, "```") { + query = query[3:] + if strings.HasSuffix(query, "```") { + query = query[:len(query)-3] + } + query = strings.TrimSpace(query) + } + if head, tail, found := strings.Cut(query, "\n"); found { + if fenceLanguages[strings.ToLower(strings.TrimSpace(head))] { + query = strings.TrimSpace(tail) + } + } + return query +} + +// splitColumns parses a comma-separated string or list column config. +func (c *PostgreSQLConnector) splitColumns(value any) []string { + switch typed := value.(type) { + case string: + parts := strings.Split(typed, ",") + columns := make([]string, 0, len(parts)) + for _, part := range parts { + if column := strings.TrimSpace(part); column != "" { + columns = append(columns, column) + } + } + return columns + case []any: + columns := make([]string, 0, len(typed)) + for _, item := range typed { + if column := strings.TrimSpace(stringConfig(item)); column != "" { + columns = append(columns, column) + } + } + return columns + } + return nil +} + +type postgresSyncSession struct { + connector *PostgreSQLConnector + db *sql.DB + queries []string + queryIndex int + rows *sql.Rows + batchSize int +} + +// NextBatch returns the next PostgreSQL document batch. +func (s *postgresSyncSession) NextBatch(ctx context.Context) (SyncBatch, error) { + documents := make([]SourceDocument, 0, s.batchSize) + for len(documents) < s.batchSize { + if s.rows == nil { + if s.queryIndex >= len(s.queries) { + if len(documents) == 0 { + return SyncBatch{}, io.EOF + } + break + } + if err := s.openNextQuery(ctx); err != nil { + return SyncBatch{}, err + } + } + if !s.rows.Next() { + if err := s.rows.Err(); err != nil { + s.closeRows() + return SyncBatch{}, err + } + s.closeRows() + continue + } + row, columns, err := s.connector.scanRow(s.rows) + if err != nil { + // Skip rows that fail to convert (mirrors Python). + continue + } + if doc, ok := s.connector.rowToSourceDocument(row, columns); ok { + documents = append(documents, doc) + } + } + return SyncBatch{Documents: documents}, nil +} + +// Close closes the PostgreSQL sync session. +func (s *postgresSyncSession) Close() error { + s.closeRows() + return s.db.Close() +} + +// openNextQuery runs the next base query. +func (s *postgresSyncSession) openNextQuery(ctx context.Context) error { + query := s.queries[s.queryIndex] + s.queryIndex++ + rows, err := s.db.QueryContext(ctx, query) + if err != nil { + return fmt.Errorf("PostgreSQL query failed: %w", err) + } + s.rows = rows + return nil +} + +// closeRows releases the current result set. +func (s *postgresSyncSession) closeRows() { + if s.rows != nil { + s.rows.Close() + s.rows = nil + } +} + +type postgresPruneSession struct { + connector *PostgreSQLConnector + db *sql.DB + queries []string + queryIndex int + rows *sql.Rows + batchSize int +} + +// NextBatch returns the next PostgreSQL prune snapshot batch. +func (s *postgresPruneSession) NextBatch(ctx context.Context) (PruneBatch, error) { + documents := make([]SlimDocument, 0, s.batchSize) + for len(documents) < s.batchSize { + if s.rows == nil { + if s.queryIndex >= len(s.queries) { + if len(documents) == 0 { + return PruneBatch{}, io.EOF + } + break + } + if err := s.openNextQuery(ctx); err != nil { + return PruneBatch{}, err + } + } + if !s.rows.Next() { + if err := s.rows.Err(); err != nil { + s.closeRows() + return PruneBatch{}, err + } + s.closeRows() + continue + } + row, columns, err := s.connector.scanRow(s.rows) + if err != nil { + continue + } + documents = append(documents, SlimDocument{SourceID: s.connector.buildDocumentID(row, columns)}) + } + return PruneBatch{Documents: documents}, nil +} + +// Close closes the PostgreSQL prune session. +func (s *postgresPruneSession) Close() error { + s.closeRows() + return s.db.Close() +} + +// openNextQuery runs the next slim query. +func (s *postgresPruneSession) openNextQuery(ctx context.Context) error { + query := s.queries[s.queryIndex] + s.queryIndex++ + rows, err := s.db.QueryContext(ctx, query) + if err != nil { + return fmt.Errorf("PostgreSQL query failed: %w", err) + } + s.rows = rows + return nil +} + +// closeRows releases the current result set. +func (s *postgresPruneSession) closeRows() { + if s.rows != nil { + s.rows.Close() + s.rows = nil + } +} diff --git a/internal/syncer/connector/postgresql_test.go b/internal/syncer/connector/postgresql_test.go new file mode 100644 index 0000000000..1996f91de2 --- /dev/null +++ b/internal/syncer/connector/postgresql_test.go @@ -0,0 +1,342 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package connector + +import ( + "context" + "database/sql" + "errors" + "io" + "net/url" + "regexp" + "strings" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" +) + +// newFixturePostgresConnector builds a PostgreSQL connector backed by sqlmock. +func newFixturePostgresConnector(t *testing.T, config map[string]any, expect func(mock sqlmock.Sqlmock)) *PostgreSQLConnector { + t.Helper() + if config == nil { + config = map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "credentials": map[string]any{ + "username": "postgres", + "password": "secret", + }, + } + } + connector, err := NewPostgreSQLConnector(config) + if err != nil { + t.Fatalf("NewPostgreSQLConnector failed: %v", err) + } + db, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New failed: %v", err) + } + t.Cleanup(func() { + db.Close() + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } + }) + connector.openDB = func(dsn string) (*sql.DB, error) { + return db, nil + } + if expect != nil { + expect(mock) + } + return connector +} + +// TestPostgreSQLConnectorOpenSyncCustomQuery verifies a custom query produces documents. +func TestPostgreSQLConnectorOpenSyncCustomQuery(t *testing.T) { + query := "SELECT * FROM products WHERE status = 'active'" + updatedAt := mustTime(t, "2026-01-02T03:04:05Z") + connector := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "query": query, + "content_columns": "title,description", + "metadata_columns": "id,category,updated_at", + "id_column": "id", + "timestamp_column": "updated_at", + "credentials": map[string]any{ + "username": "postgres", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta(query)).WillReturnRows( + sqlmock.NewRows([]string{"id", "title", "description", "category", "updated_at"}). + AddRow(7, "Hello/World", "Some body", "news", updatedAt), + ) + }) + + session, err := connector.OpenSync(context.Background(), SyncRequest{FromBeginning: true}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 { + t.Fatalf("documents len = %d, want 1", len(batch.Documents)) + } + doc := batch.Documents[0] + if doc.SourceID != "postgresql:mydb:7" { + t.Fatalf("source id = %q", doc.SourceID) + } + if doc.SemanticIdentifier != "Hello/World" { + t.Fatalf("semantic identifier = %q", doc.SemanticIdentifier) + } + blob := string(doc.Blob) + if !strings.Contains(blob, "【title】:\nHello/World") || !strings.Contains(blob, "【description】:\nSome body") { + t.Fatalf("blob = %q", blob) + } + if !doc.UpdatedAt.Equal(updatedAt) { + t.Fatalf("updated at = %s", doc.UpdatedAt) + } + if doc.Metadata["category"] != "news" || doc.Metadata["id"] != "7" { + t.Fatalf("metadata = %v", doc.Metadata) + } + if doc.Metadata["updated_at"] != updatedAt.Format(time.RFC3339) { + t.Fatalf("metadata updated_at = %v", doc.Metadata["updated_at"]) + } + if _, err = session.NextBatch(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("NextBatch EOF = %v", err) + } +} + +// TestPostgreSQLConnectorOpenSyncIncrementalWindow verifies the ISO-8601 timestamp filter SQL. +func TestPostgreSQLConnectorOpenSyncIncrementalWindow(t *testing.T) { + connector := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "query": "SELECT * FROM products", + "timestamp_column": "updated_at", + "credentials": map[string]any{ + "username": "postgres", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + expected := "SELECT * FROM (SELECT * FROM products) AS ragflow_src " + + "WHERE ragflow_src.updated_at >= '2026-01-01T00:00:00Z' AND ragflow_src.updated_at <= '2026-01-02T00:00:00Z'" + mock.ExpectQuery(regexp.QuoteMeta(expected)).WillReturnRows( + sqlmock.NewRows([]string{"updated_at"}).AddRow(mustTime(t, "2026-01-01T12:00:00Z")), + ) + }) + + start := mustTime(t, "2026-01-01T00:00:00Z") + end := mustTime(t, "2026-01-02T00:00:00Z") + session, err := connector.OpenSync(context.Background(), SyncRequest{WindowStart: &start, WindowEnd: end}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 { + t.Fatalf("documents len = %d, want 1", len(batch.Documents)) + } +} + +// TestPostgreSQLConnectorOpenSyncAllTables verifies the information_schema table listing. +func TestPostgreSQLConnectorOpenSyncAllTables(t *testing.T) { + connector := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "id_column": "id", + "credentials": map[string]any{ + "username": "postgres", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE'")).WillReturnRows( + sqlmock.NewRows([]string{"table_name"}).AddRow("products").AddRow("orders"), + ) + mock.ExpectQuery(regexp.QuoteMeta(`SELECT * FROM "public"."products"`)).WillReturnRows( + sqlmock.NewRows([]string{"id", "name"}).AddRow(1, "Product"), + ) + mock.ExpectQuery(regexp.QuoteMeta(`SELECT * FROM "public"."orders"`)).WillReturnRows( + sqlmock.NewRows([]string{"id", "name"}).AddRow(2, "Order"), + ) + }) + + session, err := connector.OpenSync(context.Background(), SyncRequest{FromBeginning: true}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + var ids []string + for { + batch, err := session.NextBatch(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + for _, doc := range batch.Documents { + ids = append(ids, doc.SourceID) + } + } + if len(ids) != 2 || ids[0] != "postgresql:mydb:1" || ids[1] != "postgresql:mydb:2" { + t.Fatalf("source ids = %v", ids) + } +} + +// TestPostgreSQLConnectorOpenPrune verifies Python-compatible slim IDs. +func TestPostgreSQLConnectorOpenPrune(t *testing.T) { + connector := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "query": "SELECT * FROM products", + "content_columns": "title,description", + "id_column": "id", + "credentials": map[string]any{ + "username": "postgres", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + expected := "SELECT ragflow_src.id FROM (SELECT * FROM products) AS ragflow_src" + mock.ExpectQuery(regexp.QuoteMeta(expected)).WillReturnRows( + sqlmock.NewRows([]string{"id"}).AddRow(3).AddRow(4), + ) + }) + + session, err := connector.OpenPrune(context.Background(), PruneRequest{}) + if err != nil { + t.Fatalf("OpenPrune failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 2 || + batch.Documents[0].SourceID != "postgresql:mydb:3" || + batch.Documents[1].SourceID != "postgresql:mydb:4" { + t.Fatalf("slim documents = %+v", batch.Documents) + } +} + +// TestPostgreSQLConnectorDSN verifies connector-controlled sslmode and connect_timeout. +func TestPostgreSQLConnectorDSN(t *testing.T) { + var captured string + connector := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "credentials": map[string]any{ + "username": "postgres", + "password": "p@ss:word", + }, + }, nil) + connector.openDB = func(dsn string) (*sql.DB, error) { + captured = dsn + return nil, nil + } + if _, err := connector.open(); err != nil { + t.Fatalf("open failed: %v", err) + } + parsed, err := url.Parse(captured) + if err != nil { + t.Fatalf("parse dsn %q: %v", captured, err) + } + if got := parsed.Query().Get("sslmode"); got != "prefer" { + t.Fatalf("sslmode = %q", got) + } + if got := parsed.Query().Get("connect_timeout"); got != "30" { + t.Fatalf("connect_timeout = %q", got) + } + if parsed.User.Username() != "postgres" { + t.Fatalf("username = %q", parsed.User.Username()) + } + if pass, _ := parsed.User.Password(); pass != "p@ss:word" { + t.Fatalf("password = %q", pass) + } + if parsed.Host != "127.0.0.1:5432" || parsed.Path != "/mydb" { + t.Fatalf("host/path = %q %q", parsed.Host, parsed.Path) + } +} + +// TestPostgreSQLConnectorOpenSyncMixedCaseTable verifies schema-qualified +// quoting of a catalog-discovered mixed-case table name. +func TestPostgreSQLConnectorOpenSyncMixedCaseTable(t *testing.T) { + connector := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "id_column": "id", + "credentials": map[string]any{ + "username": "postgres", + "password": "secret", + }, + }, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE'")).WillReturnRows( + sqlmock.NewRows([]string{"table_name"}).AddRow("MixedCase"), + ) + mock.ExpectQuery(regexp.QuoteMeta(`SELECT * FROM "public"."MixedCase"`)).WillReturnRows( + sqlmock.NewRows([]string{"id", "name"}).AddRow(1, "Item"), + ) + }) + + session, err := connector.OpenSync(context.Background(), SyncRequest{FromBeginning: true}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 || batch.Documents[0].SourceID != "postgresql:mydb:1" { + t.Fatalf("documents = %+v", batch.Documents) + } +} + +// TestPostgreSQLConnectorValidate verifies the probe and dialect-specific error message. +func TestPostgreSQLConnectorValidate(t *testing.T) { + connector := newFixturePostgresConnector(t, nil, func(mock sqlmock.Sqlmock) { + mock.ExpectQuery(regexp.QuoteMeta("SELECT 1")).WillReturnRows( + sqlmock.NewRows([]string{"1"}).AddRow(1), + ) + }) + if err := connector.Validate(context.Background()); err != nil { + t.Fatalf("Validate failed: %v", err) + } + + missing := newFixturePostgresConnector(t, map[string]any{ + "host": "127.0.0.1", + "port": "5432", + "database": "mydb", + "credentials": map[string]any{ + "username": "", + "password": "secret", + }, + }, nil) + if err := missing.Validate(context.Background()); err == nil || !strings.Contains(err.Error(), "postgresql") { + t.Fatalf("Validate error = %v", err) + } +} diff --git a/internal/syncer/sync_runner.go b/internal/syncer/sync_runner.go index 279d68fcf1..f7e913ec24 100644 --- a/internal/syncer/sync_runner.go +++ b/internal/syncer/sync_runner.go @@ -60,7 +60,11 @@ func (r *SyncRunner) Run(ctx context.Context, taskContext service.SyncTaskContex } var windowStart *time.Time // = nil if it is `Full synchronisation` if !service.IsFromBeginning(taskContext.Task.FromBeginning) { - windowStart = taskContext.Task.PollRangeStart + if pollRangeStart := taskContext.Task.PollRangeStart; pollRangeStart != nil { + if start := pollRangeStart.Time(); !start.IsZero() { + windowStart = &start + } + } } checkpointState, err := r.prepareCheckpoint(ctx, taskContext, windowStart, time.Now().UTC()) if err != nil { diff --git a/internal/syncer/syncer.go b/internal/syncer/syncer.go index 51ae09ce07..4a7be97404 100644 --- a/internal/syncer/syncer.go +++ b/internal/syncer/syncer.go @@ -175,6 +175,8 @@ func registerBuiltInConnectors(registry *syncerconnector.Registry) { registerDAOConnector(registry, "gmail", syncerconnector.NewGmailConnector) registerDAOConnector(registry, "google-drive", syncerconnector.NewGoogleDriveConnector) registerDAOConnector(registry, "google_drive", syncerconnector.NewGoogleDriveConnector) + registerDAOConnector(registry, "mysql", syncerconnector.NewMySQLConnector) + registerDAOConnector(registry, "postgresql", syncerconnector.NewPostgreSQLConnector) } func registerDAOConnector[T syncerconnector.Connector](registry *syncerconnector.Registry, source string, factory func(map[string]any) (T, error)) { diff --git a/internal/syncer/syncer_test.go b/internal/syncer/syncer_test.go index 59136280e4..e87e0e96c4 100644 --- a/internal/syncer/syncer_test.go +++ b/internal/syncer/syncer_test.go @@ -1209,8 +1209,8 @@ func TestSyncRunnerClampsWaterlineToWindowEnd(t *testing.T) { if task.PollRangeEnd == nil { t.Fatalf("poll_range_end is nil") } - if task.PollRangeEnd.Equal(future) || task.PollRangeEnd.After(after) || task.PollRangeEnd.Before(before.Add(-time.Second)) { - t.Fatalf("poll_range_end = %s, want clamped to run window around %s - %s", task.PollRangeEnd, before, after) + if task.PollRangeEnd.Time().Equal(future) || task.PollRangeEnd.Time().After(after) || task.PollRangeEnd.Time().Before(before.Add(-time.Second)) { + t.Fatalf("poll_range_end = %s, want clamped to run window around %s - %s", task.PollRangeEnd.Time(), before, after) } } @@ -1243,8 +1243,8 @@ func TestFullSyncWaterlineUsesWindowEnd(t *testing.T) { if task.PollRangeEnd == nil { t.Fatalf("poll_range_end is nil") } - if task.PollRangeEnd.Equal(oldSourceTime) || task.PollRangeEnd.Before(before.Add(-time.Second)) || task.PollRangeEnd.After(after) { - t.Fatalf("poll_range_end = %s, want run window around %s - %s", task.PollRangeEnd, before, after) + if task.PollRangeEnd.Time().Equal(oldSourceTime) || task.PollRangeEnd.Time().Before(before.Add(-time.Second)) || task.PollRangeEnd.Time().After(after) { + t.Fatalf("poll_range_end = %s, want run window around %s - %s", task.PollRangeEnd.Time(), before, after) } }