zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-14 20:54:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add MySQL and PostgreSQL connectors to Go syncer (#18162)

Browse Source 
As title

---------

Co-authored-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
mkaaad
2026-08-13 20:16:44 +08:00
committed by GitHub
parent c5ff2bced5
commit 3ce1b9be3d
13 changed files with 2111 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
go.mod
View File

@@ -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

8
go.sum
View File

@@ -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=

7
internal/dao/connector.go
View File

@@ -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,

4
internal/dao/sync_task.go
View File

@@ -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,

30
internal/entity/connector.go
View File

@@ -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
}

158
internal/entity/flexible_time.go Normal file
View File

@@ -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
}

594
internal/syncer/connector/mysql.go Normal file
View File

@@ -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:<database>:<id value>" 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
}
}

351
internal/syncer/connector/mysql_test.go Normal file
View File

@@ -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)
}
}
}

621
internal/syncer/connector/postgresql.go Normal file
View File

@@ -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:<database>:<id value>" 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
}
}

342
internal/syncer/connector/postgresql_test.go Normal file
View File

@@ -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)
}
}

6
internal/syncer/sync_runner.go
View File

@@ -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 {

2
internal/syncer/syncer.go
View File

@@ -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)) {

8
internal/syncer/syncer_test.go
View File

@@ -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)
}
}