11import { JobManager } from "../config" ;
22import * as vscode from "vscode" ;
33import Statement from "../database/statement" ;
4+ import { ContextItem } from "@continuedev/core" ;
45
56export function canTalkToDb ( ) {
67 return JobManager . getSelection ( ) !== undefined ;
@@ -14,73 +15,31 @@ export function getCurrentSchema(): string {
1415} ;
1516
1617export type TableRefs = { [ key : string ] : TableColumn [ ] } ;
17-
18- export async function getTableMetaData ( schema : string , tableName : string ) : Promise < TableColumn [ ] > {
19- const objectFindStatement = [
20- `SELECT ` ,
21- ` column.TABLE_SCHEMA,` ,
22- ` column.TABLE_NAME,` ,
23- ` column.COLUMN_NAME,` ,
24- ` key.CONSTRAINT_NAME,` ,
25- ` column.DATA_TYPE, ` ,
26- ` column.CHARACTER_MAXIMUM_LENGTH,` ,
27- ` column.NUMERIC_SCALE, ` ,
28- ` column.NUMERIC_PRECISION,` ,
29- ` column.IS_NULLABLE, ` ,
30- // ` column.HAS_DEFAULT, `,
31- // ` column.COLUMN_DEFAULT, `,
32- ` column.COLUMN_TEXT, ` ,
33- ` column.IS_IDENTITY` ,
34- `FROM QSYS2.SYSCOLUMNS2 as column` ,
35- `LEFT JOIN QSYS2.syskeycst as key` ,
36- ` on ` ,
37- ` column.table_schema = key.table_schema and` ,
38- ` column.table_name = key.table_name and` ,
39- ` column.column_name = key.column_name` ,
40- `WHERE column.TABLE_SCHEMA = '${ Statement . delimName ( schema , true ) } ,
41- `AND column.TABLE_NAME = '${ Statement . delimName ( tableName , true ) } ,
42- `ORDER BY column.ORDINAL_POSITION` ,
43- ] . join ( ` ` ) ;
44-
45- return await JobManager . runSQL ( objectFindStatement ) ;
46- }
47-
48- export async function parsePromptForRefs ( stream : vscode . ChatResponseStream , prompt : string [ ] ) : Promise < TableRefs > {
49- const tables : TableRefs = { } ;
50- for ( const word of prompt ) {
51- const [ schema , table ] = word . split ( `.` ) ;
52- const cleanedTable = table . replace ( / [ , \/ # ! ? $ % \^ & \* ; : { } = \- _ ` ~ ( ) ] / g, "" ) ;
53- if ( schema && cleanedTable ) {
54- if ( stream !== null ) {
55- stream . progress ( `looking up information for ${ schema } ${ cleanedTable } )
56- }
57- const data = await getTableMetaData ( schema , cleanedTable ) ;
58- tables [ cleanedTable ] = tables [ cleanedTable ] || [ ] ;
59- tables [ cleanedTable ] . push ( ...data ) ;
60- }
61- }
62- return tables ;
18+ interface MarkdownRef {
19+ TABLE_NAME : string ,
20+ COLUMN_INFO ?: string ,
21+ SCHMEA ?: string ,
6322}
6423
6524export async function findPossibleTables ( stream : vscode . ChatResponseStream , schema : string , words : string [ ] ) {
6625
6726 let tables : TableRefs = { }
6827
69- // parse all SCHEMA.TABLE references first
70- tables = await parsePromptForRefs ( stream , words . filter ( word => word . includes ( '.' ) ) ) ;
71-
28+ // Parse all SCHEMA.TABLE references first
29+ const schemaTableRefs = words . filter ( word => word . includes ( '.' ) ) ;
7230 const justWords = words . map ( word => word . replace ( / [ , \/ # ! ? $ % \^ & \* ; : { } = \- _ ` ~ ( ) ] / g, "" ) ) ;
7331
7432 // Remove plurals from words
7533 justWords . push ( ...justWords . filter ( word => word . endsWith ( 's' ) ) . map ( word => word . slice ( 0 , - 1 ) ) ) ;
7634
77- // filter prompt for possible refs to tables
35+ // Filter prompt for possible refs to tables
7836 const validWords = justWords
7937 . filter ( word => word . length > 2 && ! word . endsWith ( 's' ) && ! word . includes ( `'` ) )
8038 . map ( word => `'${ Statement . delimName ( word , true ) } ) ;
8139
8240 const objectFindStatement = [
8341 `SELECT ` ,
42+ ` column.TABLE_SCHEMA,` ,
8443 ` column.TABLE_NAME,` ,
8544 ` column.COLUMN_NAME,` ,
8645 ` key.CONSTRAINT_NAME,` ,
@@ -99,15 +58,18 @@ export async function findPossibleTables(stream: vscode.ChatResponseStream, sche
9958 ` column.table_schema = key.table_schema and` ,
10059 ` column.table_name = key.table_name and` ,
10160 ` column.column_name = key.column_name` ,
102- `WHERE column.TABLE_SCHEMA = '${ schema } ,
61+ `WHERE column.TABLE_SCHEMA = '${ Statement . delimName ( schema , true ) } ,
10362 ...[
104- words . length > 0
105- ? `AND column.TABLE_NAME in (${ validWords . join ( `, ` ) }
106- : `` ,
63+ schemaTableRefs . length > 0
64+ ? `AND (column.TABLE_NAME in (${ validWords . join ( `, ` ) } ${ schemaTableRefs . map ( ref => {
65+ const [ schema , table ] = ref . split ( '.' ) ;
66+ const cleanedTable = table . replace ( / [ , \/ # ! ? $ % \^ & \* ; : { } = \- _ ` ~ ( ) ] / g, "" ) ;
67+ return `(column.TABLE_SCHEMA = '${ Statement . delimName ( schema , true ) } ${ Statement . delimName ( cleanedTable , true ) } ;
68+ } ) . join ( ' OR ' ) } ))`
69+ : `AND column.TABLE_NAME in (${ validWords . join ( `, ` ) } ,
10770 ] ,
10871 `ORDER BY column.ORDINAL_POSITION` ,
10972 ] . join ( ` ` ) ;
110-
11173 // TODO
11274 const result : TableColumn [ ] = await JobManager . runSQL ( objectFindStatement ) ;
11375
@@ -138,41 +100,44 @@ export async function findPossibleTables(stream: vscode.ChatResponseStream, sche
138100 *
139101 * Tables with names starting with 'SYS' are skipped.
140102 */
141- export function refsToMarkdown ( refs : TableRefs ) {
103+ export function refsToMarkdown ( refs : TableRefs ) : MarkdownRef [ ] {
142104 const condensedResult = Object . keys ( refs ) . length > 5 ;
143105
144- let markdown : string [ ] = [ ] ;
145-
106+ let markdownRefs : MarkdownRef [ ] = [ ] ;
146107 for ( const tableName in refs ) {
147108 if ( tableName . startsWith ( `SYS` ) ) continue ;
148109
149- markdown . push ( `# ${ tableName } , `` ) ;
150-
151- if ( condensedResult ) {
152- markdown . push ( `| Column | Type | Text |` ) ;
153- markdown . push ( `| - | - | - |` ) ;
154- } else {
155- markdown . push (
156- `| Column | Type | Nullable | Identity | Text | Constraint |`
157- ) ;
158- markdown . push ( `| - | - | - | - | - | - |` ) ;
159- }
160- for ( const column of refs [ tableName ] ) {
161- if ( condensedResult ) {
162- markdown . push (
163- `| ${ column . COLUMN_NAME } ${ column . DATA_TYPE } ${ column . COLUMN_TEXT }
164- ) ;
165- } else {
166- markdown . push (
167- `| ${ column . COLUMN_NAME } ${ column . DATA_TYPE } ${ column . IS_NULLABLE } ${ column . IS_IDENTITY } ${ column . COLUMN_TEXT } ${ column . CONSTRAINT_NAME }
168- ) ;
169- }
170- }
171-
172- markdown . push ( `` ) ;
110+ const curRef : MarkdownRef = {
111+ TABLE_NAME : tableName ,
112+ SCHMEA : refs [ tableName ] [ 0 ] . TABLE_SCHEMA ,
113+ COLUMN_INFO : refs [ tableName ] . map ( column => {
114+ const lengthPrecision = column . CHARACTER_MAXIMUM_LENGTH
115+ ? `(${ column . CHARACTER_MAXIMUM_LENGTH } ${ column . NUMERIC_PRECISION ? `:${ column . NUMERIC_PRECISION } : `` }
116+ : `` ;
117+ return `${ column . COLUMN_NAME } ${ column . COLUMN_TEXT ? ` - ${ column . COLUMN_TEXT } : `` } ${ column . DATA_TYPE } ${ lengthPrecision } ${ column . IS_IDENTITY } ${ column . IS_NULLABLE } ;
118+ } ) . join ( `\n` )
119+ } ;
120+ markdownRefs . push ( curRef ) ;
121+ }
122+
123+ return markdownRefs ;
124+ }
125+
126+ export function createContinueContextItems ( refs : MarkdownRef [ ] ) {
127+ const contextItems : ContextItem [ ] = [ ] ;
128+ const job = JobManager . getSelection ( ) ;
129+ for ( const tableRef of refs ) {
130+ let prompt = `Table: ${ tableRef . TABLE_NAME } ${ tableRef . SCHMEA } ;
131+ prompt += `Format: column_name (column_text) type(length:precision) is_identity is_nullable\n`
132+ prompt += `${ tableRef . COLUMN_INFO } ;
133+ contextItems . push ( {
134+ name : `${ job . name } ${ tableRef . SCHMEA } ${ tableRef . TABLE_NAME } ,
135+ description : `Column information for ${ tableRef . TABLE_NAME } ,
136+ content : prompt ,
137+ } ) ;
173138 }
174139
175- return markdown . join ( `\n` ) ;
140+ return contextItems ;
176141}
177142
178143export async function getSystemStatus ( ) : Promise < string > {
0 commit comments