Add EstimatedTableRowsCount to avoid count in large tables

pkg/client/client.go

@@ -221,12 +221,36 @@ func (client *Client) TableRows(table string, opts RowsOptions) (*Result, error)
 	return client.query(sql)
 }
 
+func (client *Client) EstimatedTableRowsCount(table string, opts RowsOptions) (*Result, error) {
+	schema, table := getSchemaAndTable(table)
+	sql := fmt.Sprintf(`SELECT reltuples FROM pg_class WHERE oid = '%s.%s'::regclass;`, schema, table)
+
+	result, err := client.query(sql)
+	if err != nil {
+		return nil, err
+	}
+	// float64 to int64 conversion
+	estimatedRowsCount := result.Rows[0][0].(float64)
+	result.Rows[0] = Row{int64(estimatedRowsCount)}
+
+	return result, nil
+}
+
 func (client *Client) TableRowsCount(table string, opts RowsOptions) (*Result, error) {
 	schema, table := getSchemaAndTable(table)
 	sql := fmt.Sprintf(`SELECT COUNT(1) FROM "%s"."%s"`, schema, table)
 
 	if opts.Where != "" {
 		sql += fmt.Sprintf(" WHERE %s", opts.Where)
+	} else if client.serverType == postgresType {
+		tableInfo, err := client.TableInfo(table)
+		if err != nil {
+			return nil, err
+		}
+		estimatedRowsCount := tableInfo.Rows[0][3].(float64)
+		if estimatedRowsCount > 100000 {
+			return client.EstimatedTableRowsCount(table, opts)
+		}
 	}
 
 	return client.query(sql)

pkg/client/client_test.go

@@ -278,6 +278,35 @@ func test_TableInfo(t *testing.T) {
 	assert.Equal(t, 1, len(res.Rows))
 }
 
+func test_EstimatedTableRowsCount(t *testing.T) {
+	var count int64 = 15
+	res, err := testClient.EstimatedTableRowsCount("books", RowsOptions{})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, []string{"reltuples"}, res.Columns)
+	assert.Equal(t, []Row{Row{count}}, res.Rows)
+}
+
+func test_TableRowsCount(t *testing.T) {
+	var count int64 = 15
+	res, err := testClient.TableRowsCount("books", RowsOptions{})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, []string{"count"}, res.Columns)
+	assert.Equal(t, []Row{Row{count}}, res.Rows)
+}
+
+func test_TableRowsCountWithLargeTable(t *testing.T) {
+	var count int64 = 100010
+	testClient.db.MustExec(`create table large_table as select s from generate_Series(1,100010) s;`)
+	testClient.db.MustExec(`VACUUM large_table;`)
+	res, err := testClient.TableRowsCount("large_table", RowsOptions{})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, []string{"reltuples"}, res.Columns)
+	assert.Equal(t, []Row{Row{count}}, res.Rows)
+}
+
 func test_TableIndexes(t *testing.T) {
 	res, err := testClient.TableIndexes("books")
 
@@ -400,6 +429,9 @@ func TestAll(t *testing.T) {
 	test_Table(t)
 	test_TableRows(t)
 	test_TableInfo(t)
+	test_EstimatedTableRowsCount(t)
+	test_TableRowsCount(t)
+	test_TableRowsCountWithLargeTable(t)
 	test_TableIndexes(t)
 	test_TableConstraints(t)
 	test_Query(t)