pgm/cmd/parse.go

// Patial Tech.
// Author, Ankit Patial

package main

import (
	"bytes"
	"fmt"
	"regexp"
	"strings"
)

type (
	TableInfo struct {
		Schema     string
		Table      string
		Columns    []*Column
		PrimaryKey []string
	}

	Column struct {
		// Name of column
		Name string
		// Type of column
		Type string
	}
)

func parse(scheam []byte) ([]*TableInfo, error) {
	var (
		t      = false
		n      = bytes.Count(scheam, []byte("CREATE TABLE"))
		tables = make([]*TableInfo, n)
		i      = 0
		sb     strings.Builder
	)

	for l := range bytes.SplitSeq(scheam, []byte("\n")) {
		if strings.HasPrefix(string(l), "CREATE TABLE") {
			t = true
			sb.Write(l)
		} else if t {
			sb.Write(l)
			if strings.Contains(string(l), ";") {
				t = false
				t, err := parseTableStmt(sb.String())
				if err != nil {
					return nil, err
				}
				tables[i] = t
				i++
				sb.Reset()
			}

		}
	}

	return tables, nil
}

func parseTableStmt(sqlStatement string) (*TableInfo, error) {
	// Regex to match CREATE TABLE statement and extract table name and column definitions
	tableRegex := regexp.MustCompile(`(?i)CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(?:` + "`" + `?(\w+)` + "`" + `?\.)?` + "`" + `?(\w+)` + "`" + `?\s*\(([\s\S]*?)\)(?:\s*ENGINE.*?)?;`)
	matches := tableRegex.FindStringSubmatch(sqlStatement)
	if matches == nil {
		return nil, fmt.Errorf("no CREATE TABLE statement found")
	}

	schema := matches[1]
	tbl := matches[2]
	cols := matches[3]

	// Parse column definitions by splitting first on commas
	// This is a simplistic approach - a real SQL parser would be more robust
	columns := parseColumns(cols)

	return &TableInfo{
		Schema:  schema,
		Table:   tbl,
		Columns: columns,
	}, nil
}

func parseColumns(colsStr string) []*Column {
	var columns []*Column
	var currentDef strings.Builder
	parenthesesCount := 0
	inQuote := false

	// First, intelligently split the column definitions
	var columnDefs []string
	for _, char := range colsStr {
		if char == '(' {
			parenthesesCount++
			currentDef.WriteRune(char)
		} else if char == ')' {
			parenthesesCount--
			currentDef.WriteRune(char)
		} else if char == '\'' || char == '"' {
			inQuote = !inQuote
			currentDef.WriteRune(char)
		} else if char == ',' && parenthesesCount == 0 && !inQuote {
			// End of a column definition
			columnDefs = append(columnDefs, strings.TrimSpace(currentDef.String()))
			currentDef.Reset()
		} else {
			currentDef.WriteRune(char)
		}
	}

	// Add the last column definition if there's anything left
	if currentDef.Len() > 0 {
		columnDefs = append(columnDefs, strings.TrimSpace(currentDef.String()))
	}

	// Now parse each column definition
	for _, colDef := range columnDefs {
		// Skip constraints and keys that don't define columns
		if strings.HasPrefix(strings.ToUpper(colDef), "PRIMARY KEY") ||
			strings.HasPrefix(strings.ToUpper(colDef), "UNIQUE KEY") ||
			strings.HasPrefix(strings.ToUpper(colDef), "FOREIGN KEY") ||
			strings.HasPrefix(strings.ToUpper(colDef), "KEY") ||
			strings.HasPrefix(strings.ToUpper(colDef), "INDEX") ||
			strings.HasPrefix(strings.ToUpper(colDef), "CONSTRAINT") {
			continue
		}

		colNameRegex := regexp.MustCompile(`^` + "`" + `?(\w+)` + "`" + `?\s+(.+)$`)
		matches := colNameRegex.FindStringSubmatch(colDef)
		if matches != nil {
			columns = append(columns, &Column{
				Name: matches[1],
				Type: strings.TrimSpace(matches[2]),
			})
		}
	}

	return columns
}