refactor: partition tasks by range and by status in OLAP repository (#6)

* refactor: timescale partition by statuses

* pub concurrency 2 -> 1

Author: abelanger5

internal/msgqueue/mq_pub_buffer.go

@@ -9,7 +9,7 @@ import (
 
 const PUB_FLUSH_INTERVAL = 10 * time.Millisecond
 const PUB_BUFFER_SIZE = 1000
-const PUB_MAX_CONCURRENCY = 2
+const PUB_MAX_CONCURRENCY = 1
 
 type PubFunc func(m *Message) error
 

pkg/repository/olap.go

@@ -99,7 +99,7 @@ func NewOLAPEventRepository(l *zerolog.Logger) OLAPEventRepository {
 	queries := timescalev2.New()
 
 	// create partitions of the events OLAP table
-	partitionCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	partitionCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
 	err = queries.CreateOLAPPartitions(partitionCtx, timescalePool, NUM_PARTITIONS)
@@ -108,6 +108,57 @@ func NewOLAPEventRepository(l *zerolog.Logger) OLAPEventRepository {
 		log.Fatal(err)
 	}
 
+	today := time.Now().UTC()
+	tomorrow := today.AddDate(0, 0, 1)
+	sevenDaysAgo := today.AddDate(0, 0, -7)
+
+	err = queries.CreateOLAPTaskPartition(partitionCtx, timescalePool, pgtype.Date{
+		Time:  today,
+		Valid: true,
+	})
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	err = queries.CreateOLAPTaskPartition(partitionCtx, timescalePool, pgtype.Date{
+		Time:  tomorrow,
+		Valid: true,
+	})
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	partitions, err := queries.ListOLAPTaskPartitionsBeforeDate(partitionCtx, timescalePool, pgtype.Date{
+		Time:  sevenDaysAgo,
+		Valid: true,
+	})
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	for _, partition := range partitions {
+		_, err := timescalePool.Exec(
+			partitionCtx,
+			fmt.Sprintf("ALTER TABLE v2_task DETACH PARTITION %s CONCURRENTLY", partition),
+		)
+
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		_, err = timescalePool.Exec(
+			partitionCtx,
+			fmt.Sprintf("DROP TABLE %s", partition),
+		)
+
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+
 	return &olapEventRepository{
 		pool:       timescalePool,
 		l:          l,

pkg/repository/v2/sqlcv2/dags.sql

@@ -1,3 +1,16 @@
+-- name: CreateDAGPartition :exec
+SELECT create_v2_dag_partition(
+    @date::date
+);
+
+-- name: ListDAGPartitionsBeforeDate :many
+SELECT
+    p::text AS partition_name
+FROM
+    get_v2_dag_partitions_before(
+        @date::date
+    ) AS p;
+
 -- name: GetDAGData :many
 WITH input AS (
     SELECT

pkg/repository/v2/sqlcv2/dags.sql.go

@@ -1,9 +1,9 @@
--- name: CreateTablePartition :exec
+-- name: CreateTaskPartition :exec
 SELECT create_v2_task_partition(
     @date::date
 );
 
--- name: ListTablePartitionsBeforeDate :many
+-- name: ListTaskPartitionsBeforeDate :many
 SELECT
     p::text AS partition_name
 FROM

pkg/repository/v2/sqlcv2/tasks.sql

@@ -81,7 +81,7 @@ func (r *TaskRepositoryImpl) UpdateTablePartitions(ctx context.Context) error {
 	tomorrow := today.AddDate(0, 0, 1)
 	sevenDaysAgo := today.AddDate(0, 0, -7)
 
-	err := r.queries.CreateTablePartition(ctx, r.pool, pgtype.Date{
+	err := r.queries.CreateTaskPartition(ctx, r.pool, pgtype.Date{
 		Time:  today,
 		Valid: true,
 	})
@@ -90,7 +90,7 @@ func (r *TaskRepositoryImpl) UpdateTablePartitions(ctx context.Context) error {
 		return err
 	}
 
-	err = r.queries.CreateTablePartition(ctx, r.pool, pgtype.Date{
+	err = r.queries.CreateTaskPartition(ctx, r.pool, pgtype.Date{
 		Time:  tomorrow,
 		Valid: true,
 	})
@@ -99,7 +99,7 @@ func (r *TaskRepositoryImpl) UpdateTablePartitions(ctx context.Context) error {
 		return err
 	}
 
-	partitions, err := r.queries.ListTablePartitionsBeforeDate(ctx, r.pool, pgtype.Date{
+	partitions, err := r.queries.ListTaskPartitionsBeforeDate(ctx, r.pool, pgtype.Date{
 		Time:  sevenDaysAgo,
 		Valid: true,
 	})
@@ -128,6 +128,53 @@ func (r *TaskRepositoryImpl) UpdateTablePartitions(ctx context.Context) error {
 		}
 	}
 
+	err = r.queries.CreateDAGPartition(ctx, r.pool, pgtype.Date{
+		Time:  today,
+		Valid: true,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	err = r.queries.CreateDAGPartition(ctx, r.pool, pgtype.Date{
+		Time:  tomorrow,
+		Valid: true,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	dagPartitions, err := r.queries.ListDAGPartitionsBeforeDate(ctx, r.pool, pgtype.Date{
+		Time:  sevenDaysAgo,
+		Valid: true,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	for _, partition := range dagPartitions {
+		_, err := r.pool.Exec(
+			ctx,
+			fmt.Sprintf("ALTER TABLE v2_task DETACH PARTITION %s CONCURRENTLY", partition),
+		)
+
+		if err != nil {
+			return err
+		}
+
+		_, err = r.pool.Exec(
+			ctx,
+			fmt.Sprintf("DROP TABLE %s", partition),
+		)
+
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 

pkg/repository/v2/sqlcv2/tasks.sql.go

@@ -3,6 +3,19 @@ SELECT create_v2_task_events_partitions(
     @partitions::int
 );
 
+-- name: CreateOLAPTaskPartition :exec
+SELECT create_v2_tasks_olap_partition(
+    @date::date
+);
+
+-- name: ListOLAPTaskPartitionsBeforeDate :many
+SELECT
+    p::text AS partition_name
+FROM
+    get_v2_tasks_olap_partitions_before(
+        @date::date
+    ) AS p;
+
 -- name: CreateTasksOLAP :copyfrom
 INSERT INTO v2_tasks_olap (
     tenant_id,
@@ -480,4 +493,4 @@ WITH locked_events AS (
 SELECT
     COUNT(*)
 FROM
-    locked_events;
+    locked_events;