diff --git a/internal/infra/postgresql/column.go b/internal/infra/postgresql/column.go index 448dc21..f2f91c7 100644 --- a/internal/infra/postgresql/column.go +++ b/internal/infra/postgresql/column.go @@ -11,9 +11,10 @@ var ErrUnsupportedPartitionKeyType = errors.New("unsupported partition key colum type ColumnType string const ( - Date ColumnType = "date" - DateTime ColumnType = "timestamp" - UUID ColumnType = "uuid" + Date ColumnType = "date" + DateTime ColumnType = "timestamp" + DateTimeWithTZ ColumnType = "timestamp with time zone" + UUID ColumnType = "uuid" ) func (p Postgres) GetColumnDataType(schema, table, column string) (ColumnType, error) { @@ -40,7 +41,7 @@ func (p Postgres) GetColumnDataType(schema, table, column string) (ColumnType, e case "timestamp without time zone": return DateTime, nil case "timestamp with time zone": - return DateTime, nil + return DateTimeWithTZ, nil case "uuid": return UUID, nil default: diff --git a/internal/infra/postgresql/column_internal_test.go b/internal/infra/postgresql/column_internal_test.go index 19b3fd9..3f65f98 100644 --- a/internal/infra/postgresql/column_internal_test.go +++ b/internal/infra/postgresql/column_internal_test.go @@ -50,7 +50,7 @@ func TestGetColumn(t *testing.T) { { "Date time with time zone", "timestamp with time zone", - DateTime, + DateTimeWithTZ, }, { "UUID", diff --git a/pkg/ppm/bounds.go b/pkg/ppm/bounds.go index 5adeeda..acb8add 100644 --- a/pkg/ppm/bounds.go +++ b/pkg/ppm/bounds.go @@ -30,6 +30,11 @@ func parseBounds(partition postgresql.PartitionResult) (lowerBound time.Time, up return lowerBound, upperBound, nil } + lowerBound, upperBound, err = parseBoundAsDateTimeWithTimezone(partition) + if err == nil { + return lowerBound, upperBound, nil + } + lowerBound, upperBound, err = parseBoundAsUUIDv7(partition) if err == nil { return lowerBound, upperBound, nil @@ -70,6 +75,23 @@ func parseBoundAsDateTime(partition postgresql.PartitionResult) (lowerBound, upp return lowerBound, upperBound, nil } +func parseBoundAsDateTimeWithTimezone(partition postgresql.PartitionResult) (lowerBound, upperBound time.Time, err error) { + lowerBound, err = time.Parse("2006-01-02 15:04:05Z07", partition.LowerBound) + if err != nil { + return time.Time{}, time.Time{}, fmt.Errorf("can't parse lowerbound as datetime with timezone: %w", err) + } + + upperBound, err = time.Parse("2006-01-02 15:04:05Z07", partition.UpperBound) + if err != nil { + return time.Time{}, time.Time{}, fmt.Errorf("can't parse upperbound as datetime with timezone: %w", err) + } + + lowerBound = convertToDateTimeWithoutTimezone(lowerBound) + upperBound = convertToDateTimeWithoutTimezone(upperBound) + + return lowerBound, upperBound, nil +} + func parseBoundAsUUIDv7(partition postgresql.PartitionResult) (lowerBound, upperBound time.Time, err error) { lowerBoundUUID, err := uuid.Parse(partition.LowerBound) if err != nil { @@ -90,3 +112,11 @@ func parseBoundAsUUIDv7(partition postgresql.PartitionResult) (lowerBound, upper return lowerBound, upperBound, nil } + +func convertToDateTimeWithoutTimezone(bound time.Time) time.Time { + parsedTime, err := time.Parse("2006-01-02 15:04:05", bound.UTC().Format("2006-01-02 15:04:05")) + if err != nil { + return time.Time{} + } + return parsedTime +} diff --git a/pkg/ppm/bounds_internal_test.go b/pkg/ppm/bounds_internal_test.go index ad103c0..afdda04 100644 --- a/pkg/ppm/bounds_internal_test.go +++ b/pkg/ppm/bounds_internal_test.go @@ -37,6 +37,17 @@ func TestParseBounds(t *testing.T) { "2024-01-01T10:00:00Z", "2025-02-03T12:53:00Z", }, + { + "Datetime with timezone bounds", + postgresql.PartitionResult{ + Schema: "public", + Name: "my_table", + LowerBound: "2024-01-01 23:30:00-01", + UpperBound: "2025-02-03 00:30:00+01", + }, + "2024-01-02T00:30:00Z", + "2025-02-02T23:30:00Z", + }, { "UUIDv7 bounds", postgresql.PartitionResult{ diff --git a/pkg/ppm/checkpartition.go b/pkg/ppm/checkpartition.go index 30184c1..563aa0b 100644 --- a/pkg/ppm/checkpartition.go +++ b/pkg/ppm/checkpartition.go @@ -21,6 +21,7 @@ var ( var SupportedPartitionKeyDataType = []postgresql.ColumnType{ postgresql.Date, postgresql.DateTime, + postgresql.DateTimeWithTZ, postgresql.UUID, } diff --git a/pkg/ppm/checkpartition_test.go b/pkg/ppm/checkpartition_test.go index a7da574..2bfdea9 100644 --- a/pkg/ppm/checkpartition_test.go +++ b/pkg/ppm/checkpartition_test.go @@ -52,48 +52,49 @@ func TestCheckPartitions(t *testing.T) { // Build mock for each partitions for _, p := range partitions { var tables []partition.Partition - - var table partition.Partition - - for i := 0; i <= p.Retention; i++ { - switch p.Interval { - case partition.Daily: - table, _ = p.GeneratePartition(time.Now().AddDate(0, 0, -i)) - case partition.Weekly: - table, _ = p.GeneratePartition(time.Now().AddDate(0, 0, -i*7)) - case partition.Quarterly: - table, _ = p.GeneratePartition(time.Now().AddDate(0, i*-3, 0)) - case partition.Monthly: - table, _ = p.GeneratePartition(time.Now().AddDate(0, -i, 0)) - case partition.Yearly: - table, _ = p.GeneratePartition(time.Now().AddDate(-i, 0, 0)) - default: - t.Errorf("unuspported partition interval in retention table mock") - } - - postgreSQLMock.On("GetColumnDataType", table.Schema, table.ParentTable, p.PartitionKey).Return(postgresql.Date, nil).Once() - tables = append(tables, table) + var retentionTables []partition.Partition + var preprovisionedTables []partition.Partition + + // Create retention partitions + forDate := time.Now() + switch p.Interval { + case partition.Daily: + retentionTables, _ = p.GetRetentionPartitions(forDate) + case partition.Weekly: + retentionTables, _ = p.GetRetentionPartitions(forDate) + case partition.Quarterly: + retentionTables, _ = p.GetRetentionPartitions(forDate) + case partition.Monthly: + retentionTables, _ = p.GetRetentionPartitions(forDate) + case partition.Yearly: + retentionTables, _ = p.GetRetentionPartitions(forDate) + default: + t.Errorf("unuspported partition interval in retention table mock") } - - for i := 0; i <= p.PreProvisioned; i++ { - switch p.Interval { - case partition.Daily: - table, _ = p.GeneratePartition(time.Now().AddDate(0, 0, i)) - case partition.Weekly: - table, _ = p.GeneratePartition(time.Now().AddDate(0, 0, i*7)) - case partition.Monthly: - table, _ = p.GeneratePartition(time.Now().AddDate(0, i, 0)) - case partition.Quarterly: - table, _ = p.GeneratePartition(time.Now().AddDate(0, i*3, 0)) - case partition.Yearly: - table, _ = p.GeneratePartition(time.Now().AddDate(i, 0, 0)) - default: - t.Errorf("unuspported partition interval in preprovisonned table mock") - } - - postgreSQLMock.On("GetColumnDataType", table.Schema, table.ParentTable, p.PartitionKey).Return(postgresql.Date, nil).Once() - tables = append(tables, table) + tables = append(tables, retentionTables...) + + // Create current partition + currentPartition, _ := p.GeneratePartition(forDate) + tables = append(tables, currentPartition) + + // Create preprovisioned partitions + switch p.Interval { + case partition.Daily: + preprovisionedTables, _ = p.GetPreProvisionedPartitions(forDate) + case partition.Weekly: + preprovisionedTables, _ = p.GetPreProvisionedPartitions(forDate) + case partition.Monthly: + preprovisionedTables, _ = p.GetPreProvisionedPartitions(forDate) + case partition.Quarterly: + preprovisionedTables, _ = p.GetPreProvisionedPartitions(forDate) + case partition.Yearly: + preprovisionedTables, _ = p.GetPreProvisionedPartitions(forDate) + default: + t.Errorf("unuspported partition interval in preprovisonned table mock") } + tables = append(tables, preprovisionedTables...) + + postgreSQLMock.On("GetColumnDataType", p.Schema, p.Table, p.PartitionKey).Return(postgresql.Date, nil).Once() postgreSQLMock.On("GetPartitionSettings", p.Schema, p.Table).Return(string(partition.Range), p.PartitionKey, nil).Once() diff --git a/pkg/ppm/provisioning.go b/pkg/ppm/provisioning.go index d90a72a..85525ca 100644 --- a/pkg/ppm/provisioning.go +++ b/pkg/ppm/provisioning.go @@ -103,7 +103,7 @@ func (p PPM) CreatePartition(partitionConfiguration partition.Configuration, par case postgresql.Date: lowerBound = partition.LowerBound.Format("2006-01-02") upperBound = partition.UpperBound.Format("2006-01-02") - case postgresql.DateTime: + case postgresql.DateTime, postgresql.DateTimeWithTZ: lowerBound = partition.LowerBound.Format("2006-01-02 00:00:00") upperBound = partition.UpperBound.Format("2006-01-02 00:00:00") case postgresql.UUID: