WIP.

artie-labs · Aug 29, 2024 · 9dd64bc · 9dd64bc
1 parent 19fa197
commit 9dd64bc
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/config/postgres.go b/config/postgres.go
@@ -39,6 +39,7 @@ type PostgreSQLTable struct {
 
 	// Optional settings
 	BatchSize                  uint     `yaml:"batchSize,omitempty"`
+	PrimaryKeysOverride        []string `yaml:"primaryKeysOverride,omitempty"`
 	OptionalPrimaryKeyValStart string   `yaml:"optionalPrimaryKeyValStart,omitempty"`
 	OptionalPrimaryKeyValEnd   string   `yaml:"optionalPrimaryKeyValEnd,omitempty"`
 	ExcludeColumns             []string `yaml:"excludeColumns,omitempty"`

diff --git a/lib/postgres/table.go b/lib/postgres/table.go
@@ -18,7 +18,7 @@ type Table struct {
 	PrimaryKeys []string
 }
 
-func LoadTable(db *sql.DB, _schema string, name string) (*Table, error) {
+func LoadTable(db *sql.DB, _schema string, name string, primaryKeysOverride []string) (*Table, error) {
 	tbl := &Table{
 		Name:   name,
 		Schema: _schema,
@@ -29,8 +29,12 @@ func LoadTable(db *sql.DB, _schema string, name string) (*Table, error) {
 		return nil, fmt.Errorf("failed to describe table %s.%s: %w", tbl.Schema, tbl.Name, err)
 	}
 
-	if tbl.PrimaryKeys, err = schema.FetchPrimaryKeys(db, tbl.Schema, tbl.Name); err != nil {
-		return nil, fmt.Errorf("failed to retrieve primary keys: %w", err)
+	if len(primaryKeysOverride) > 0 {
+		tbl.PrimaryKeys = primaryKeysOverride
+	} else {
+		if tbl.PrimaryKeys, err = schema.FetchPrimaryKeys(db, tbl.Schema, tbl.Name); err != nil {
+			return nil, fmt.Errorf("failed to retrieve primary keys: %w", err)
+		}
 	}
 
 	return tbl, nil

diff --git a/sources/postgres/adapter/adapter.go b/sources/postgres/adapter/adapter.go
@@ -26,7 +26,7 @@ type PostgresAdapter struct {
 
 func NewPostgresAdapter(db *sql.DB, tableCfg config.PostgreSQLTable) (PostgresAdapter, error) {
 	slog.Info("Loading metadata for table")
-	table, err := postgres.LoadTable(db, tableCfg.Schema, tableCfg.Name)
+	table, err := postgres.LoadTable(db, tableCfg.Schema, tableCfg.Name, tableCfg.PrimaryKeysOverride)
 	if err != nil {
 		return PostgresAdapter{}, fmt.Errorf("failed to load metadata for table %s.%s: %w", tableCfg.Schema, tableCfg.Name, err)
 	}