Add correct pagination for filters

core/graph/schema.graphqls

@@ -17,6 +17,7 @@ type Column {
 type RowsResult {
   Columns: [Column!]!
   Rows: [[String!]!]!
+  TotalCount: Int!
   DisableUpdate: Boolean!
 }
 

core/src/engine/plugin.go

@@ -41,6 +41,7 @@ type Column struct {
 type GetRowsResult struct {
 	Columns       []Column
 	Rows          [][]string
+	TotalCount    int
 	DisableUpdate bool
 }
 

core/src/plugins/clickhouse/chat.go

@@ -0,0 +1,61 @@
+package clickhouse
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/clidey/whodb/core/src/common"
+	"github.com/clidey/whodb/core/src/engine"
+	"github.com/clidey/whodb/core/src/llm"
+)
+
+func (p *ClickHousePlugin) Chat(config *engine.PluginConfig, schema string, model string, previousConversation string, query string) ([]*engine.ChatMessage, error) {
+	db, err := DB(config)
+	if err != nil {
+		return nil, err
+	}
+
+	tableFields, err := getAllTableSchema(db, schema)
+	if err != nil {
+		return nil, err
+	}
+
+	tableDetails := strings.Builder{}
+	for tableName, fields := range tableFields {
+		tableDetails.WriteString(fmt.Sprintf("table: %v\n", tableName))
+		for _, field := range fields {
+			tableDetails.WriteString(fmt.Sprintf("- %v (%v)\n", field.Key, field.Value))
+		}
+	}
+
+	context := tableDetails.String()
+
+	completeQuery := fmt.Sprintf(common.RawSQLQueryPrompt, "Postgres", schema, context, previousConversation, query, "Postgres")
+
+	response, err := llm.Instance(config).Complete(completeQuery, llm.LLMModel(model), nil)
+	if err != nil {
+		return nil, err
+	}
+
+	chats := common.ExtractCodeFromResponse(*response)
+	chatMessages := []*engine.ChatMessage{}
+	for _, chat := range chats {
+		var result *engine.GetRowsResult
+		chatType := "message"
+		if chat.Type == "sql" {
+			rowResult, err := p.RawExecute(config, chat.Text)
+			if err != nil {
+				return nil, err
+			}
+			chatType = "sql"
+			result = rowResult
+		}
+		chatMessages = append(chatMessages, &engine.ChatMessage{
+			Type:   chatType,
+			Result: result,
+			Text:   chat.Text,
+		})
+	}
+
+	return chatMessages, nil
+}

core/src/plugins/clickhouse/clickhouse.go

@@ -102,6 +102,35 @@
 	return storageUnits, nil
 }
 
+func getAllTableSchema(conn *sql.DB, schema string) (map[string][]engine.Record, error) {
+	query := fmt.Sprintf(`
+		SELECT 
+			table,
+			name,
+			type
+		FROM system.columns
+		WHERE database = '%s'
+		ORDER BY table, position
+	`, schema)
+
+	rows, err := conn.QueryContext(context.Background(), query)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	tableColumnsMap := make(map[string][]engine.Record)
+	for rows.Next() {
+		var tableName, columnName, dataType string
+		if err := rows.Scan(&tableName, &columnName, &dataType); err != nil {
+			return nil, err
+		}
+		tableColumnsMap[tableName] = append(tableColumnsMap[tableName], engine.Record{Key: columnName, Value: dataType})
+	}
+
+	return tableColumnsMap, nil
+}
+
 func getTableSchema(conn *sql.DB, schema string, tableName string) ([]engine.Record, error) {
 	query := fmt.Sprintf(`
 		SELECT 
@@ -130,12 +159,6 @@
 	return result, nil
 }
 
-func (p *ClickHousePlugin) Chat(config *engine.PluginConfig, schema string, model string, previousConversation string, query string) ([]*engine.ChatMessage, error) {
-	// Implement chat functionality similar to MySQL implementation
-	// You may need to adapt this based on ClickHouse specifics
-	return nil, fmt.Errorf("chat functionality not implemented for ClickHouse")
-}
-
 func NewClickHousePlugin() *engine.Plugin {
 	return &engine.Plugin{
 		Type:            engine.DatabaseType_ClickHouse,

core/src/plugins/clickhouse/query.go

@@ -4,17 +4,23 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+
 	"github.com/clidey/whodb/core/src/engine"
 )
 
 func (p *ClickHousePlugin) GetRows(config *engine.PluginConfig, schema string, storageUnit string, where string, pageSize int, pageOffset int) (*engine.GetRowsResult, error) {
-	query := fmt.Sprintf("SELECT * FROM %s.%s", schema, storageUnit)
+	baseQuery := fmt.Sprintf("FROM %s.%s", schema, storageUnit)
 	if where != "" {
-		query += " WHERE " + where
+		baseQuery += " WHERE " + where
 	}
-	query += fmt.Sprintf(" LIMIT %d OFFSET %d", pageSize, pageOffset)
+	query := fmt.Sprintf("SELECT * %s LIMIT %d OFFSET %d", baseQuery, pageSize, pageOffset)
 
-	return p.executeQuery(config, query)
+	result, err := p.executeQuery(config, query)
+	if err != nil {
+		return nil, err
+	}
+
+	return result, nil
 }
 
 func (p *ClickHousePlugin) RawExecute(config *engine.PluginConfig, query string) (*engine.GetRowsResult, error) {

core/src/plugins/elasticsearch/elasticsearch.go

@@ -97,15 +97,15 @@ func (p *ElasticSearchPlugin) GetRows(config *engine.PluginConfig, database, col
 		query[key] = value
 	}
 
-	var buf bytes.Buffer
-	if err := json.NewEncoder(&buf).Encode(query); err != nil {
+	var searchBuf bytes.Buffer
+	if err := json.NewEncoder(&searchBuf).Encode(query); err != nil {
 		return nil, err
 	}
 
 	res, err := client.Search(
 		client.Search.WithContext(context.Background()),
 		client.Search.WithIndex(collection),
-		client.Search.WithBody(&buf),
+		client.Search.WithBody(&searchBuf),
 		client.Search.WithTrackTotalHits(true),
 	)
 	if err != nil {
@@ -122,7 +122,10 @@ func (p *ElasticSearchPlugin) GetRows(config *engine.PluginConfig, database, col
 		return nil, err
 	}
 
-	hits := searchResult["hits"].(map[string]interface{})["hits"].([]interface{})
+	hitsInfo := searchResult["hits"].(map[string]interface{})
+	totalHits := int(hitsInfo["total"].(map[string]interface{})["value"].(float64))
+
+	hits := hitsInfo["hits"].([]interface{})
 	result := &engine.GetRowsResult{
 		Columns: []engine.Column{
 			{Name: "document", Type: "Document"},
@@ -142,6 +145,7 @@ func (p *ElasticSearchPlugin) GetRows(config *engine.PluginConfig, database, col
 		result.Rows = append(result.Rows, []string{string(jsonBytes)})
 	}
 
+	result.TotalCount = totalHits
 	return result, nil
 }
 

core/src/plugins/mongodb/mongodb.go

@@ -89,6 +89,11 @@ func (p *MongoDBPlugin) GetRows(config *engine.PluginConfig, database, collectio
 		}
 	}
 
+	totalCount, err := coll.CountDocuments(context.TODO(), bsonFilter)
+	if err != nil {
+		return nil, err
+	}
+
 	findOptions := options.Find()
 	findOptions.SetLimit(int64(pageSize))
 	findOptions.SetSkip(int64(pageOffset))
@@ -124,6 +129,7 @@ func (p *MongoDBPlugin) GetRows(config *engine.PluginConfig, database, collectio
 		})
 	}
 
+	result.TotalCount = int(totalCount)
 	return result, nil
 }
 

core/src/plugins/mysql/mysql.go

@@ -137,12 +137,32 @@
 }
 
 func (p *MySQLPlugin) GetRows(config *engine.PluginConfig, schema string, storageUnit string, where string, pageSize int, pageOffset int) (*engine.GetRowsResult, error) {
-	query := fmt.Sprintf("SELECT * FROM `%v`.`%s`", schema, storageUnit)
+	db, err := DB(config)
+	if err != nil {
+		return nil, err
+	}
+
+	baseQuery := fmt.Sprintf("FROM `%v`.`%s`", schema, storageUnit)
 	if len(where) > 0 {
-		query = fmt.Sprintf("%v WHERE %v", query, where)
+		baseQuery = fmt.Sprintf("%v WHERE %v", baseQuery, where)
+	}
+
+	countQuery := fmt.Sprintf("SELECT COUNT(*) %v", baseQuery)
+	var totalCount int
+	err = db.Raw(countQuery).Scan(&totalCount).Error
+	if err != nil {
+		return nil, err
+	}
+
+	query := fmt.Sprintf("SELECT * %v LIMIT ? OFFSET ?", baseQuery)
+
+	result, err := p.executeRawSQL(config, query, pageSize, pageOffset)
+	if err != nil {
+		return nil, err
 	}
-	query = fmt.Sprintf("%v LIMIT ? OFFSET ?", query)
-	return p.executeRawSQL(config, query, pageSize, pageOffset)
+
+	result.TotalCount = totalCount
+	return result, nil
 }
 
 func (p *MySQLPlugin) executeRawSQL(config *engine.PluginConfig, query string, params ...interface{}) (*engine.GetRowsResult, error) {