diff --git a/pkg/client/client.go b/pkg/client/client.go index 0e826ff..5e7350a 100644 --- a/pkg/client/client.go +++ b/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) diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go index 22f3416..2f5014d 100644 --- a/pkg/client/client_test.go +++ b/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)