refactor(display): added row count to partition list command

turbot · Jan 3, 2025 · 3bf8cfc · 3bf8cfc
1 parent b73df0e
commit 3bf8cfc
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 22 deletions.
diff --git a/internal/database/tables.go b/internal/database/tables.go
@@ -139,7 +139,7 @@ func getDirNames(folderPath string) ([]string, error) {
 	return dirNames, nil
 }
 
-func GetRowCount(ctx context.Context, tableName string) (int64, error) {
+func GetRowCount(ctx context.Context, tableName string, partitionName *string) (int64, error) {
 	// Open a DuckDB connection
 	db, err := sql.Open("duckdb", filepaths.TailpipeDbFilePath())
 	if err != nil {
@@ -152,6 +152,9 @@ func GetRowCount(ctx context.Context, tableName string) (int64, error) {
 		return 0, fmt.Errorf("invalid table name")
 	}
 	query := fmt.Sprintf("SELECT COUNT(*) FROM %s", tableName) // #nosec G201 // this is a controlled query tableName must match a regex
+	if partitionName != nil {
+		query = fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE tp_partition = '%s'", tableName, *partitionName) // #nosec G201 // this is a controlled query tableName must match a regex
+	}
 	rows, err := db.QueryContext(ctx, query)
 	if err != nil {
 		return 0, fmt.Errorf("failed to get row count: %w", err)

diff --git a/internal/display/partition.go b/internal/display/partition.go
@@ -6,16 +6,19 @@ import (
 	"path"
 	"strings"
 
+	"github.com/dustin/go-humanize"
+
 	"github.com/turbot/pipe-fittings/printers"
 	"github.com/turbot/tailpipe/internal/config"
 	"github.com/turbot/tailpipe/internal/constants"
+	"github.com/turbot/tailpipe/internal/database"
 )
 
 type PartitionResource struct {
-	Name        string       `json:"name"`
-	Description *string      `json:"description,omitempty"`
-	Plugin      string       `json:"plugin"`
-	Local       FileMetadata `json:"local,omitempty"`
+	Name        string             `json:"name"`
+	Description *string            `json:"description,omitempty"`
+	Plugin      string             `json:"plugin"`
+	Local       TableResourceFiles `json:"local,omitempty"`
 }
 
 // GetShowData implements the printers.Showable interface
@@ -37,6 +40,7 @@ func (r *PartitionResource) GetListData() *printers.RowData {
 		printers.NewFieldValue("PLUGIN", r.Plugin),
 		printers.NewFieldValue("LOCAL SIZE", r.Local.HumanizeSize()),
 		printers.NewFieldValue("FILES", r.Local.HumanizeCount()),
+		printers.NewFieldValue("ROWS", humanize.Comma(r.Local.RowCount)),
 	)
 	return res
 }
@@ -99,7 +103,16 @@ func (r *PartitionResource) setFileInformation() error {
 		return err
 	}
 
-	r.Local = metadata
+	r.Local.FileMetadata = metadata
+
+	if metadata.FileCount > 0 {
+		var rc int64
+		rc, err = database.GetRowCount(context.Background(), nameParts[0], &nameParts[1])
+		if err != nil {
+			return fmt.Errorf("unable to obtain row count: %w", err)
+		}
+		r.Local.RowCount = rc
+	}
 
 	return nil
 }
diff --git a/internal/display/table.go b/internal/display/table.go
@@ -91,14 +91,6 @@ func ListTableResources(ctx context.Context) ([]*TableResource, error) {
 				return nil, fmt.Errorf("unable to obtain file information: %w", err)
 			}
 
-			if table.Local.FileCount > 0 {
-				rc, err := database.GetRowCount(ctx, t)
-				if err != nil {
-					return nil, fmt.Errorf("unable to obtain row count: %w", err)
-				}
-				table.Local.RowCount = rc
-			}
-
 			res = append(res, table)
 		}
 	}
@@ -142,14 +134,6 @@ func GetTableResource(ctx context.Context, tableName string) (*TableResource, er
 			return nil, fmt.Errorf("unable to obtain file information: %w", err)
 		}
 
-		if table.Local.FileCount > 0 {
-			rc, err := database.GetRowCount(ctx, tableName)
-			if err != nil {
-				return nil, fmt.Errorf("unable to obtain row count: %w", err)
-			}
-			table.Local.RowCount = rc
-		}
-
 		return table, nil
 	} else {
 		return nil, fmt.Errorf("table %s not found", tableName)
@@ -174,6 +158,15 @@ func (r *TableResource) setFileInformation() error {
 
 	r.Local.FileMetadata = metadata
 
+	if metadata.FileCount > 0 {
+		var rc int64
+		rc, err = database.GetRowCount(context.Background(), r.Name, nil)
+		if err != nil {
+			return fmt.Errorf("unable to obtain row count: %w", err)
+		}
+		r.Local.RowCount = rc
+	}
+
 	return nil
 }