Skip to content

Commit 9b3143c

Browse files
pan93412pan93412
committed
feat(sqlrunner): implement GetDatabaseStructure method
1 parent b47f1e6 commit 9b3143c

File tree

3 files changed

+204
-0
lines changed

3 files changed

+204
-0
lines changed

‎internal/sqlrunner/models.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,14 @@ type ErrorResponse struct {
4545
func (e ErrorResponse) Error() string {
4646
return fmt.Sprintf("%s: %s", e.Code, e.Message)
4747
}
48+
49+
// DatabaseStructure is the database structure of a schema.
50+
type DatabaseStructure struct {
51+
Tables []DatabaseTable `json:"tables"`
52+
}
53+
54+
// DatabaseTable is the table structure of a schema.
55+
type DatabaseTable struct {
56+
Name string `json:"name"`
57+
Columns []string `json:"columns"`
58+
}

‎internal/sqlrunner/sqlrunner.go‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,50 @@ func (s *SqlRunner) Query(ctx context.Context, schema, query string) (DataRespon
6666
return respBody.Data, nil
6767
}
6868

69+
func (s *SqlRunner) GetDatabaseStructure(ctx context.Context, schema string) (DatabaseStructure, error) {
70+
// Query SQLite's master table to get all table names
71+
tablesQuery := "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
72+
tablesResp, err := s.Query(ctx, schema, tablesQuery)
73+
if err != nil {
74+
return DatabaseStructure{}, fmt.Errorf("failed to query tables: %w", err)
75+
}
76+
77+
var tables []DatabaseTable
78+
79+
// For each table, get its column information
80+
for _, row := range tablesResp.Rows {
81+
if len(row) == 0 {
82+
continue
83+
}
84+
tableName := row[0]
85+
86+
// Use PRAGMA table_info to get column information
87+
columnsQuery := fmt.Sprintf("PRAGMA table_info(%s)", tableName)
88+
columnsResp, err := s.Query(ctx, schema, columnsQuery)
89+
if err != nil {
90+
return DatabaseStructure{}, fmt.Errorf("query columns for table %s: %w", tableName, err)
91+
}
92+
93+
var columns []string
94+
// PRAGMA table_info returns: cid, name, type, notnull, dflt_value, pk
95+
// We only need the name (index 1)
96+
for _, columnRow := range columnsResp.Rows {
97+
if len(columnRow) > 1 {
98+
columns = append(columns, columnRow[1])
99+
}
100+
}
101+
102+
tables = append(tables, DatabaseTable{
103+
Name: tableName,
104+
Columns: columns,
105+
})
106+
}
107+
108+
return DatabaseStructure{
109+
Tables: tables,
110+
}, nil
111+
}
112+
69113
func (s *SqlRunner) IsHealthy(ctx context.Context) bool {
70114
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/healthz", s.cfg.URI), nil)
71115
if err != nil {

‎internal/sqlrunner/sqlrunner_test.go‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,3 +61,152 @@ func TestQuery_SchemaError(t *testing.T) {
6161
t.Errorf("Expected SCHEMA_ERROR, got %v", errResp.Code)
6262
}
6363
}
64+
65+
func TestGetDatabaseStructure_Success(t *testing.T) {
66+
s := testhelper.NewSQLRunnerClient(t)
67+
68+
// Create a schema with multiple tables and columns
69+
schema := `
70+
CREATE TABLE users (
71+
id INTEGER PRIMARY KEY,
72+
name TEXT NOT NULL,
73+
email TEXT UNIQUE
74+
);
75+
CREATE TABLE posts (
76+
id INTEGER PRIMARY KEY,
77+
title TEXT NOT NULL,
78+
content TEXT,
79+
user_id INTEGER,
80+
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
81+
);
82+
CREATE TABLE categories (
83+
id INTEGER PRIMARY KEY,
84+
name TEXT NOT NULL
85+
);
86+
`
87+
88+
structure, err := s.GetDatabaseStructure(context.Background(), schema)
89+
if err != nil {
90+
t.Fatalf("Expected success, got error: %v", err)
91+
}
92+
93+
// Verify we have the expected number of tables
94+
if len(structure.Tables) != 3 {
95+
t.Errorf("Expected 3 tables, got %d", len(structure.Tables))
96+
}
97+
98+
// Helper function to find a table by name
99+
findTable := func(name string) *sqlrunner.DatabaseTable {
100+
for _, table := range structure.Tables {
101+
if table.Name == name {
102+
return &table
103+
}
104+
}
105+
return nil
106+
}
107+
108+
// Verify users table
109+
usersTable := findTable("users")
110+
if usersTable == nil {
111+
t.Error("Expected to find 'users' table")
112+
} else {
113+
expectedColumns := []string{"id", "name", "email"}
114+
if len(usersTable.Columns) != len(expectedColumns) {
115+
t.Errorf("Expected %d columns in users table, got %d", len(expectedColumns), len(usersTable.Columns))
116+
}
117+
for i, expected := range expectedColumns {
118+
if i >= len(usersTable.Columns) || usersTable.Columns[i] != expected {
119+
t.Errorf("Expected column %d to be '%s', got '%s'", i, expected, usersTable.Columns[i])
120+
}
121+
}
122+
}
123+
124+
// Verify posts table
125+
postsTable := findTable("posts")
126+
if postsTable == nil {
127+
t.Error("Expected to find 'posts' table")
128+
} else {
129+
expectedColumns := []string{"id", "title", "content", "user_id", "created_at"}
130+
if len(postsTable.Columns) != len(expectedColumns) {
131+
t.Errorf("Expected %d columns in posts table, got %d", len(expectedColumns), len(postsTable.Columns))
132+
}
133+
}
134+
135+
// Verify categories table
136+
categoriesTable := findTable("categories")
137+
if categoriesTable == nil {
138+
t.Error("Expected to find 'categories' table")
139+
} else {
140+
expectedColumns := []string{"id", "name"}
141+
if len(categoriesTable.Columns) != len(expectedColumns) {
142+
t.Errorf("Expected %d columns in categories table, got %d", len(expectedColumns), len(categoriesTable.Columns))
143+
}
144+
}
145+
}
146+
147+
func TestGetDatabaseStructure_EmptyDatabase(t *testing.T) {
148+
s := testhelper.NewSQLRunnerClient(t)
149+
150+
// Schema that doesn't create any tables - just a comment
151+
schema := "-- Empty database with no tables"
152+
153+
structure, err := s.GetDatabaseStructure(context.Background(), schema)
154+
if err != nil {
155+
t.Fatalf("Expected success, got error: %v", err)
156+
}
157+
158+
// Should have no tables
159+
if len(structure.Tables) != 0 {
160+
t.Errorf("Expected 0 tables in empty database, got %d", len(structure.Tables))
161+
}
162+
}
163+
164+
func TestGetDatabaseStructure_ErrorHandling(t *testing.T) {
165+
s := testhelper.NewSQLRunnerClient(t)
166+
167+
// Create a schema with syntax error that should fail
168+
schema := "CREATE TABLE invalid_syntax (id int"
169+
170+
_, err := s.GetDatabaseStructure(context.Background(), schema)
171+
if err == nil {
172+
t.Error("Expected error for invalid schema, got nil")
173+
}
174+
175+
// Should be a schema error since the schema creation fails
176+
var errResp *sqlrunner.ErrorResponse
177+
if !errors.As(err, &errResp) {
178+
t.Errorf("Expected ErrorResponse, got %v", err)
179+
}
180+
if errResp.Code != sqlrunner.ErrorCodeSchemaError {
181+
t.Errorf("Expected SCHEMA_ERROR, got %v", errResp.Code)
182+
}
183+
}
184+
185+
func TestGetDatabaseStructure_WithViews(t *testing.T) {
186+
s := testhelper.NewSQLRunnerClient(t)
187+
188+
// Create schema with both tables and views
189+
schema := `
190+
CREATE TABLE products (
191+
id INTEGER PRIMARY KEY,
192+
name TEXT NOT NULL,
193+
price DECIMAL(10,2)
194+
);
195+
CREATE VIEW expensive_products AS
196+
SELECT * FROM products WHERE price > 100;
197+
`
198+
199+
structure, err := s.GetDatabaseStructure(context.Background(), schema)
200+
if err != nil {
201+
t.Fatalf("Expected success, got error: %v", err)
202+
}
203+
204+
// Should only include tables, not views (since we filter by type='table')
205+
if len(structure.Tables) != 1 {
206+
t.Errorf("Expected 1 table (views should be excluded), got %d", len(structure.Tables))
207+
}
208+
209+
if structure.Tables[0].Name != "products" {
210+
t.Errorf("Expected table name 'products', got '%s'", structure.Tables[0].Name)
211+
}
212+
}

0 commit comments

Comments
 (0)