Merge pull request #52 from tzzzoz/add-datetime-with-tz

Add datetime with timezone

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:

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

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

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{

pkg/ppm/checkpartition.go

@@ -21,6 +21,7 @@ var (
 var SupportedPartitionKeyDataType = []postgresql.ColumnType{
 	postgresql.Date,
 	postgresql.DateTime,
+	postgresql.DateTimeWithTZ,
 	postgresql.UUID,
 }
 

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()
 

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: