Merge pull request #44 from gatewayd-io/detect-all-changes-to-db

Detect all changes to db

Author: mostafa

README.md

@@ -10,6 +10,7 @@ GatewayD plugin for caching query results.
 - Support for setting expiry time on cached data
 - Support for caching responses from multiple databases on multiple servers
 - Support nested queries: joins, unions, multi-table selects, and the like
+- Support for invalidating cache based on truncate, drop and alter statements
 - Detect client's chosen database from the client's startup message
 - Metrics for quantifying cache hits, misses, gets, sets and deletes
 - Logging at various levels

plugin/utils.go

@@ -34,6 +34,39 @@ func GetQueryFromRequest(req string) (string, error) {
 	return string(requestDecoded[MinPgSQLMessageLength:size]), nil
 }
 
+// isMulti checks if the query is a union, intersect, or except.
+func isMulti(stmt *pgQuery.SelectStmt) bool {
+	return stmt.GetOp() == pgQuery.SetOperation_SETOP_UNION ||
+		stmt.GetOp() == pgQuery.SetOperation_SETOP_INTERSECT ||
+		stmt.GetOp() == pgQuery.SetOperation_SETOP_EXCEPT
+}
+
+// getSingleTable returns the tables used in a query.
+func getSingleTable(stmt *pgQuery.SelectStmt) []string {
+	tables := []string{}
+	for _, from := range stmt.FromClause {
+		rangeVar := from.GetRangeVar()
+		if rangeVar != nil {
+			tables = append(tables, rangeVar.Relname)
+		}
+	}
+
+	return tables
+}
+
+// getMultiTable returns the tables used in a union, intersect, or except query.
+func getMultiTable(stmt *pgQuery.SelectStmt) []string {
+	tables := []string{}
+	// Get the tables from the left side.
+	left := stmt.GetLarg()
+	tables = append(tables, getSingleTable(left)...)
+	// Get the tables from the right side.
+	right := stmt.GetRarg()
+	tables = append(tables, getSingleTable(right)...)
+
+	return tables
+}
+
 // GetTablesFromQuery returns the tables used in a query.
 func GetTablesFromQuery(query string) ([]string, error) {
 	stmt, err := pgQuery.Parse(query)
@@ -47,59 +80,27 @@ func GetTablesFromQuery(query string) ([]string, error) {
 
 	tables := []string{}
 
-	isMulti := func(stmt *pgQuery.SelectStmt) bool {
-		return stmt.GetOp() == pgQuery.SetOperation_SETOP_UNION ||
-			stmt.GetOp() == pgQuery.SetOperation_SETOP_INTERSECT ||
-			stmt.GetOp() == pgQuery.SetOperation_SETOP_EXCEPT
-	}
-
-	// If the query is a union, we need to get the tables from both sides.
-	getMultiTable := func(stmt *pgQuery.SelectStmt) {
-		// Get the tables from the left side.
-		for _, fromClause := range stmt.GetLarg().FromClause {
-			rangeVar := fromClause.GetRangeVar()
-			if rangeVar != nil {
-				tables = append(tables, rangeVar.Relname)
-			}
-		}
-
-		// Get the tables from the right side.
-		for _, fromClause := range stmt.GetRarg().FromClause {
-			rangeVar := fromClause.GetRangeVar()
-			if rangeVar != nil {
-				tables = append(tables, rangeVar.Relname)
-			}
-		}
-	}
-
-	getSingleTable := func(fromClause []*pgQuery.Node) {
-		for _, from := range fromClause {
-			rangeVar := from.GetRangeVar()
-			if rangeVar != nil {
-				tables = append(tables, rangeVar.Relname)
-			}
-		}
-	}
-
 	for _, stmt := range stmt.Stmts {
-		if isMulti(stmt.Stmt.GetSelectStmt()) {
-			getMultiTable(stmt.Stmt.GetSelectStmt())
+		// Get the tables from the left and right side of the complex query.
+		selectStatement := stmt.Stmt.GetSelectStmt()
+		if isMulti(selectStatement) {
+			tables = append(tables, getMultiTable(selectStatement)...)
 		}
 
 		// Get the table from the WITH clause.
 		if withClause := stmt.Stmt.GetSelectStmt().GetWithClause(); withClause != nil {
 			for _, cte := range withClause.Ctes {
 				selectStmt := cte.GetCommonTableExpr().Ctequery.GetSelectStmt()
 				if isMulti(selectStmt) {
-					getMultiTable(selectStmt)
+					tables = append(tables, getMultiTable(selectStmt)...)
 				} else {
-					getSingleTable(selectStmt.FromClause)
+					tables = append(tables, getSingleTable(selectStmt)...)
 				}
 			}
 		} else {
 			// Get the table from the FROM clause.
-			if selectQuery := stmt.Stmt.GetSelectStmt(); selectQuery != nil {
-				getSingleTable(selectQuery.FromClause)
+			if selectStatement := stmt.Stmt.GetSelectStmt(); selectStatement != nil {
+				tables = append(tables, getSingleTable(selectStatement)...)
 			}
 		}
 
@@ -114,6 +115,24 @@ func GetTablesFromQuery(query string) ([]string, error) {
 		if deleteQuery := stmt.Stmt.GetDeleteStmt(); deleteQuery != nil {
 			tables = append(tables, deleteQuery.Relation.Relname)
 		}
+
+		if truncateQuery := stmt.Stmt.GetTruncateStmt(); truncateQuery != nil {
+			for _, relation := range truncateQuery.Relations {
+				tables = append(tables, relation.GetRangeVar().Relname)
+			}
+		}
+
+		if dropTable := stmt.Stmt.GetDropStmt(); dropTable != nil {
+			for _, object := range dropTable.GetObjects() {
+				for _, table := range object.GetList().GetItems() {
+					tables = append(tables, table.GetString_().Str)
+				}
+			}
+		}
+
+		if alterTable := stmt.Stmt.GetAlterTableStmt(); alterTable != nil {
+			tables = append(tables, alterTable.Relation.Relname)
+		}
 	}
 
 	return tables, nil

plugin/utils_test.go

@@ -240,6 +240,48 @@ func Test_GetTablesFromQuery_DeleteMultiple(t *testing.T) {
 	assert.Equal(t, []string{"users", "comments"}, tables)
 }
 
+func Test_GetTablesFromQuery_Truncate(t *testing.T) {
+	query := "TRUNCATE users"
+	tables, err := GetTablesFromQuery(query)
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"users"}, tables)
+}
+
+func Test_GetTablesFromQuery_TruncateMultiple(t *testing.T) {
+	query := "TRUNCATE users, comments"
+	tables, err := GetTablesFromQuery(query)
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"users", "comments"}, tables)
+}
+
+func Test_GetTablesFromQuery_Drop(t *testing.T) {
+	query := "DROP TABLE users"
+	tables, err := GetTablesFromQuery(query)
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"users"}, tables)
+}
+
+func Test_GetTablesFromQuery_DropMultiple(t *testing.T) {
+	query := "DROP TABLE users, comments"
+	tables, err := GetTablesFromQuery(query)
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"users", "comments"}, tables)
+}
+
+func Test_GetTablesFromQuery_Alter(t *testing.T) {
+	query := "ALTER TABLE users ADD COLUMN name VARCHAR(255)"
+	tables, err := GetTablesFromQuery(query)
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"users"}, tables)
+}
+
+func Test_GetTablesFromQuery_AlterMultiple(t *testing.T) {
+	query := "ALTER TABLE users ADD COLUMN name VARCHAR(255); ALTER TABLE comments ADD COLUMN name VARCHAR(255)"
+	tables, err := GetTablesFromQuery(query)
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"users", "comments"}, tables)
+}
+
 func Test_validateAddressPort_Hostname(t *testing.T) {
 	assert.True(t, validateAddressPort("localhost:5432"))
 	assert.True(t, validateAddressPort("	localhost:5432"))