relspecgo/pkg/readers/mssql/queries.go

package mssql

import (
	"strings"

	"git.warky.dev/wdevs/relspecgo/pkg/models"
)

// querySchemas retrieves all user-defined schemas from the database
func (r *Reader) querySchemas() ([]*models.Schema, error) {
	query := `
		SELECT s.name, ISNULL(ep.value, '') as description
		FROM sys.schemas s
		LEFT JOIN sys.extended_properties ep
			ON ep.major_id = s.schema_id
			AND ep.minor_id = 0
			AND ep.class = 3
			AND ep.name = 'MS_Description'
		WHERE s.name NOT IN ('dbo', 'guest', 'INFORMATION_SCHEMA', 'sys')
		ORDER BY s.name
	`

	rows, err := r.db.QueryContext(r.ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	schemas := make([]*models.Schema, 0)
	for rows.Next() {
		var name, description string

		if err := rows.Scan(&name, &description); err != nil {
			return nil, err
		}

		schema := models.InitSchema(name)
		if description != "" {
			schema.Description = description
		}

		schemas = append(schemas, schema)
	}

	// Always include dbo schema if it has tables
	dboSchema := models.InitSchema("dbo")
	schemas = append(schemas, dboSchema)

	return schemas, rows.Err()
}

// queryTables retrieves all tables for a given schema
func (r *Reader) queryTables(schemaName string) ([]*models.Table, error) {
	query := `
		SELECT t.table_schema, t.table_name, ISNULL(ep.value, '') as description
		FROM information_schema.tables t
		LEFT JOIN sys.extended_properties ep
			ON ep.major_id = OBJECT_ID(QUOTENAME(t.table_schema) + '.' + QUOTENAME(t.table_name))
			AND ep.minor_id = 0
			AND ep.class = 1
			AND ep.name = 'MS_Description'
		WHERE t.table_schema = ? AND t.table_type = 'BASE TABLE'
		ORDER BY t.table_name
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	tables := make([]*models.Table, 0)
	for rows.Next() {
		var schema, tableName, description string

		if err := rows.Scan(&schema, &tableName, &description); err != nil {
			return nil, err
		}

		table := models.InitTable(tableName, schema)
		if description != "" {
			table.Description = description
		}

		tables = append(tables, table)
	}

	return tables, rows.Err()
}

// queryColumns retrieves all columns for tables in a schema
// Returns map[schema.table]map[columnName]*Column
func (r *Reader) queryColumns(schemaName string) (map[string]map[string]*models.Column, error) {
	query := `
		SELECT
			c.table_schema,
			c.table_name,
			c.column_name,
			c.ordinal_position,
			c.column_default,
			c.is_nullable,
			c.data_type,
			c.character_maximum_length,
			c.numeric_precision,
			c.numeric_scale,
			ISNULL(ep.value, '') as description,
			COLUMNPROPERTY(OBJECT_ID(QUOTENAME(c.table_schema) + '.' + QUOTENAME(c.table_name)), c.column_name, 'IsIdentity') as is_identity
		FROM information_schema.columns c
		LEFT JOIN sys.extended_properties ep
			ON ep.major_id = OBJECT_ID(QUOTENAME(c.table_schema) + '.' + QUOTENAME(c.table_name))
			AND ep.minor_id = COLUMNPROPERTY(OBJECT_ID(QUOTENAME(c.table_schema) + '.' + QUOTENAME(c.table_name)), c.column_name, 'ColumnId')
			AND ep.class = 1
			AND ep.name = 'MS_Description'
		WHERE c.table_schema = ?
		ORDER BY c.table_schema, c.table_name, c.ordinal_position
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	columnsMap := make(map[string]map[string]*models.Column)

	for rows.Next() {
		var schema, tableName, columnName, isNullable, dataType, description string
		var ordinalPosition int
		var columnDefault, charMaxLength, numPrecision, numScale, isIdentity *int

		if err := rows.Scan(&schema, &tableName, &columnName, &ordinalPosition, &columnDefault, &isNullable, &dataType, &charMaxLength, &numPrecision, &numScale, &description, &isIdentity); err != nil {
			return nil, err
		}

		column := models.InitColumn(columnName, tableName, schema)
		column.Type = r.mapDataType(dataType)
		column.NotNull = (isNullable == "NO")
		column.Sequence = uint(ordinalPosition)

		if description != "" {
			column.Description = description
		}

		// Check if this is an identity column (auto-increment)
		if isIdentity != nil && *isIdentity == 1 {
			column.AutoIncrement = true
		}

		if charMaxLength != nil && *charMaxLength > 0 {
			column.Length = *charMaxLength
		}

		if numPrecision != nil && *numPrecision > 0 {
			column.Precision = *numPrecision
		}

		if numScale != nil && *numScale > 0 {
			column.Scale = *numScale
		}

		// Create table key
		tableKey := schema + "." + tableName
		if columnsMap[tableKey] == nil {
			columnsMap[tableKey] = make(map[string]*models.Column)
		}
		columnsMap[tableKey][columnName] = column
	}

	return columnsMap, rows.Err()
}

// queryPrimaryKeys retrieves all primary key constraints for a schema
// Returns map[schema.table]*Constraint
func (r *Reader) queryPrimaryKeys(schemaName string) (map[string]*models.Constraint, error) {
	query := `
		SELECT
			s.name as schema_name,
			t.name as table_name,
			i.name as constraint_name,
			STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) as columns
		FROM sys.tables t
		INNER JOIN sys.indexes i ON t.object_id = i.object_id AND i.is_primary_key = 1
		INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
		INNER JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
		INNER JOIN sys.columns c ON t.object_id = c.object_id AND ic.column_id = c.column_id
		WHERE s.name = ?
		GROUP BY s.name, t.name, i.name
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	primaryKeys := make(map[string]*models.Constraint)

	for rows.Next() {
		var schema, tableName, constraintName, columnsStr string

		if err := rows.Scan(&schema, &tableName, &constraintName, &columnsStr); err != nil {
			return nil, err
		}

		columns := strings.Split(columnsStr, ",")

		constraint := models.InitConstraint(constraintName, models.PrimaryKeyConstraint)
		constraint.Schema = schema
		constraint.Table = tableName
		constraint.Columns = columns

		tableKey := schema + "." + tableName
		primaryKeys[tableKey] = constraint
	}

	return primaryKeys, rows.Err()
}

// queryForeignKeys retrieves all foreign key constraints for a schema
// Returns map[schema.table][]*Constraint
func (r *Reader) queryForeignKeys(schemaName string) (map[string][]*models.Constraint, error) {
	query := `
		SELECT
			s.name as schema_name,
			t.name as table_name,
			fk.name as constraint_name,
			rs.name as referenced_schema,
			rt.name as referenced_table,
			STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY fkc.constraint_column_id) as columns,
			STRING_AGG(rc.name, ',') WITHIN GROUP (ORDER BY fkc.constraint_column_id) as referenced_columns,
			fk.delete_referential_action_desc,
			fk.update_referential_action_desc
		FROM sys.foreign_keys fk
		INNER JOIN sys.tables t ON fk.parent_object_id = t.object_id
		INNER JOIN sys.tables rt ON fk.referenced_object_id = rt.object_id
		INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
		INNER JOIN sys.schemas rs ON rt.schema_id = rs.schema_id
		INNER JOIN sys.foreign_key_columns fkc ON fk.object_id = fkc.constraint_object_id
		INNER JOIN sys.columns c ON fkc.parent_object_id = c.object_id AND fkc.parent_column_id = c.column_id
		INNER JOIN sys.columns rc ON fkc.referenced_object_id = rc.object_id AND fkc.referenced_column_id = rc.column_id
		WHERE s.name = ?
		GROUP BY s.name, t.name, fk.name, rs.name, rt.name, fk.delete_referential_action_desc, fk.update_referential_action_desc
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	foreignKeys := make(map[string][]*models.Constraint)

	for rows.Next() {
		var schema, tableName, constraintName, refSchema, refTable, columnsStr, refColumnsStr, deleteAction, updateAction string

		if err := rows.Scan(&schema, &tableName, &constraintName, &refSchema, &refTable, &columnsStr, &refColumnsStr, &deleteAction, &updateAction); err != nil {
			return nil, err
		}

		columns := strings.Split(columnsStr, ",")
		refColumns := strings.Split(refColumnsStr, ",")

		constraint := models.InitConstraint(constraintName, models.ForeignKeyConstraint)
		constraint.Schema = schema
		constraint.Table = tableName
		constraint.Columns = columns
		constraint.ReferencedSchema = refSchema
		constraint.ReferencedTable = refTable
		constraint.ReferencedColumns = refColumns
		constraint.OnDelete = strings.ToUpper(deleteAction)
		constraint.OnUpdate = strings.ToUpper(updateAction)

		tableKey := schema + "." + tableName
		foreignKeys[tableKey] = append(foreignKeys[tableKey], constraint)
	}

	return foreignKeys, rows.Err()
}

// queryUniqueConstraints retrieves all unique constraints for a schema
// Returns map[schema.table][]*Constraint
func (r *Reader) queryUniqueConstraints(schemaName string) (map[string][]*models.Constraint, error) {
	query := `
		SELECT
			s.name as schema_name,
			t.name as table_name,
			i.name as constraint_name,
			STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) as columns
		FROM sys.tables t
		INNER JOIN sys.indexes i ON t.object_id = i.object_id AND i.is_unique = 1 AND i.is_primary_key = 0
		INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
		INNER JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
		INNER JOIN sys.columns c ON t.object_id = c.object_id AND ic.column_id = c.column_id
		WHERE s.name = ?
		GROUP BY s.name, t.name, i.name
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	uniqueConstraints := make(map[string][]*models.Constraint)

	for rows.Next() {
		var schema, tableName, constraintName, columnsStr string

		if err := rows.Scan(&schema, &tableName, &constraintName, &columnsStr); err != nil {
			return nil, err
		}

		columns := strings.Split(columnsStr, ",")

		constraint := models.InitConstraint(constraintName, models.UniqueConstraint)
		constraint.Schema = schema
		constraint.Table = tableName
		constraint.Columns = columns

		tableKey := schema + "." + tableName
		uniqueConstraints[tableKey] = append(uniqueConstraints[tableKey], constraint)
	}

	return uniqueConstraints, rows.Err()
}

// queryCheckConstraints retrieves all check constraints for a schema
// Returns map[schema.table][]*Constraint
func (r *Reader) queryCheckConstraints(schemaName string) (map[string][]*models.Constraint, error) {
	query := `
		SELECT
			s.name as schema_name,
			t.name as table_name,
			cc.name as constraint_name,
			cc.definition
		FROM sys.tables t
		INNER JOIN sys.check_constraints cc ON t.object_id = cc.parent_object_id
		INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
		WHERE s.name = ?
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	checkConstraints := make(map[string][]*models.Constraint)

	for rows.Next() {
		var schema, tableName, constraintName, definition string

		if err := rows.Scan(&schema, &tableName, &constraintName, &definition); err != nil {
			return nil, err
		}

		constraint := models.InitConstraint(constraintName, models.CheckConstraint)
		constraint.Schema = schema
		constraint.Table = tableName
		constraint.Expression = definition

		tableKey := schema + "." + tableName
		checkConstraints[tableKey] = append(checkConstraints[tableKey], constraint)
	}

	return checkConstraints, rows.Err()
}

// queryIndexes retrieves all indexes for a schema
// Returns map[schema.table][]*Index
func (r *Reader) queryIndexes(schemaName string) (map[string][]*models.Index, error) {
	query := `
		SELECT
			s.name as schema_name,
			t.name as table_name,
			i.name as index_name,
			i.is_unique,
			STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) as columns
		FROM sys.tables t
		INNER JOIN sys.indexes i ON t.object_id = i.object_id AND i.is_primary_key = 0 AND i.name IS NOT NULL
		INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
		INNER JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
		INNER JOIN sys.columns c ON t.object_id = c.object_id AND ic.column_id = c.column_id
		WHERE s.name = ?
		GROUP BY s.name, t.name, i.name, i.is_unique
	`

	rows, err := r.db.QueryContext(r.ctx, query, schemaName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	indexes := make(map[string][]*models.Index)

	for rows.Next() {
		var schema, tableName, indexName, columnsStr string
		var isUnique int

		if err := rows.Scan(&schema, &tableName, &indexName, &isUnique, &columnsStr); err != nil {
			return nil, err
		}

		columns := strings.Split(columnsStr, ",")

		index := models.InitIndex(indexName, tableName, schema)
		index.Columns = columns
		index.Unique = (isUnique == 1)
		index.Type = "btree" // MSSQL uses btree by default

		tableKey := schema + "." + tableName
		indexes[tableKey] = append(indexes[tableKey], index)
	}

	return indexes, rows.Err()
}