@@ -2,89 +2,60 @@ package main
22
33import (
44 "context"
5- "crypto/sha256"
6- "encoding/hex"
7- "errors"
85 "fmt"
9- "os"
10- "slices"
11-
126 "obot-platform/database/pkg/cmd"
13-
14- "github.com/gptscript-ai/go-gptscript"
15- _ "github.com/ncruces/go-sqlite3/driver"
16- _ "github.com/ncruces/go-sqlite3/embed"
7+ "os"
178)
189
19- var workspaceID = os .Getenv ("DATABASE_WORKSPACE_ID" )
20-
2110func main () {
2211 if len (os .Args ) != 2 {
2312 fmt .Println ("Usage: gptscript-go-tool <command>" )
2413 os .Exit (1 )
2514 }
2615 command := os .Args [1 ]
16+ ctx := context .Background ()
2717
28- g , err := gptscript . NewGPTScript ( )
29- if err != nil {
30- fmt . Printf ( "Error creating GPTScript: %v \n " , err )
31- os .Exit ( 1 )
18+ workspaceID := os . Getenv ( "DATABASE_WORKSPACE_ID" )
19+ if workspaceID == "" {
20+ // TODO(njhale): Figure out why DATABASE_WORKSPACE_ID is not set here for the UI tools.
21+ workspaceID = os .Getenv ( "GPTSCRIPT_WORKSPACE_ID" )
3222 }
33- defer g .Close ()
34-
35- var (
36- ctx = context .Background ()
37- dbFileName = "obot.db"
38- dbWorkspacePath = "/databases/" + dbFileName
39- revisionID string = "-1"
40- initialDBData []byte
41- )
42-
43- workspaceDB , err := g .ReadFileWithRevisionInWorkspace (ctx , dbWorkspacePath , gptscript.ReadFileInWorkspaceOptions {
44- WorkspaceID : workspaceID ,
45- })
4623
47- var notFoundErr * gptscript.NotFoundInWorkspaceError
48- if err != nil && ! errors .As (err , & notFoundErr ) {
49- fmt .Printf ("Error reading DB file: %v\n " , err )
50- os .Exit (1 )
51- }
24+ // Get admin DSN from environment variable
25+ adminDSN := os .Getenv ("POSTGRES_DSN" )
5226
53- // Create a temporary file for the SQLite database
54- dbFile , err := os . CreateTemp ( "" , dbFileName )
27+ // Setup database and user with admin credentials
28+ dsn , err := cmd . EnsureTenantSchema ( ctx , adminDSN , workspaceID )
5529 if err != nil {
56- fmt .Printf ("Error creating temp file : %v\n " , err )
30+ fmt .Printf ("Error setting up database : %v\n " , err )
5731 os .Exit (1 )
5832 }
59- defer dbFile .Close ()
60- defer os .Remove (dbFile .Name ())
6133
62- // Write the data to the temporary file
63- if workspaceDB != nil && workspaceDB .Content != nil {
64- initialDBData = workspaceDB .Content
65- if err := os .WriteFile (dbFile .Name (), initialDBData , 0644 ); err != nil {
66- fmt .Printf ("Error writing to temp file: %v\n " , err )
67- os .Exit (1 )
68- }
69- if workspaceDB .RevisionID != "" {
70- revisionID = workspaceDB .RevisionID
71- }
72- }
73-
74- // Run the requested command
34+ // Run the requested command using the user credentials
7535 var result string
7636 switch command {
7737 case "listDatabaseTables" :
78- result , err = cmd .ListDatabaseTables (ctx , dbFile )
38+ result , err = cmd .ListDatabaseTables (ctx , dsn )
39+
7940 case "listDatabaseTableRows" :
80- result , err = cmd .ListDatabaseTableRows (ctx , dbFile , os .Getenv ("TABLE" ))
41+ table := os .Getenv ("TABLE" )
42+ if table == "" {
43+ err = fmt .Errorf ("TABLE environment variable is required" )
44+ break
45+ }
46+ result , err = cmd .ListDatabaseTableRows (ctx , dsn , table )
47+
8148 case "runDatabaseSQL" :
82- result , err = cmd .RunDatabaseCommand (ctx , dbFile , os .Getenv ("SQL" ), "-header" )
83- if err == nil {
84- err = saveWorkspaceDB (ctx , g , dbWorkspacePath , revisionID , dbFile , initialDBData )
49+ sql := os .Getenv ("SQL" )
50+ if sql == "" {
51+ err = fmt .Errorf ("SQL environment variable is required" )
52+ break
8553 }
54+ result , err = cmd .RunDatabaseCommand (ctx , dsn , sql )
55+
8656 case "databaseContext" :
87- result , err = cmd .DatabaseContext (ctx , dbFile )
57+ result , err = cmd .DatabaseContext (ctx , dsn )
58+
8859 default :
8960 err = fmt .Errorf ("unknown command: %s" , command )
9061 }
@@ -96,66 +67,3 @@ func main() {
9667
9768 fmt .Print (result )
9869}
99-
100- // saveWorkspaceDB saves the updated database file to the workspace if the content of the database has changed.
101- func saveWorkspaceDB (
102- ctx context.Context ,
103- g * gptscript.GPTScript ,
104- dbWorkspacePath string ,
105- revisionID string ,
106- dbFile * os.File ,
107- initialDBData []byte ,
108- ) error {
109- updatedDBData , err := os .ReadFile (dbFile .Name ())
110- if err != nil {
111- return fmt .Errorf ("Error reading updated DB file: %v" , err )
112- }
113-
114- if hash (initialDBData ) == hash (updatedDBData ) {
115- return nil
116- }
117-
118- if err := g .WriteFileInWorkspace (ctx , dbWorkspacePath , updatedDBData , gptscript.WriteFileInWorkspaceOptions {
119- WorkspaceID : workspaceID ,
120- CreateRevision : & ([]bool {true }[0 ]),
121- LatestRevisionID : revisionID ,
122- }); err != nil {
123- return fmt .Errorf ("Error writing updated DB file to workspace: %v" , err )
124- }
125-
126- // Delete old revisions after successfully writing the new revision
127- revisions , err := g .ListRevisionsForFileInWorkspace (ctx , dbWorkspacePath , gptscript.ListRevisionsForFileInWorkspaceOptions {
128- WorkspaceID : workspaceID ,
129- })
130- if err != nil {
131- fmt .Fprintf (os .Stderr , "Error listing revisions: %v\n " , err )
132- return nil
133- }
134-
135- lastRevisionIndex := slices .IndexFunc (revisions , func (rev gptscript.FileInfo ) bool {
136- return rev .RevisionID == revisionID
137- })
138-
139- if lastRevisionIndex < 0 {
140- return nil
141- }
142-
143- for _ , rev := range revisions [:lastRevisionIndex + 1 ] {
144- if err := g .DeleteRevisionForFileInWorkspace (ctx , dbWorkspacePath , rev .RevisionID , gptscript.DeleteRevisionForFileInWorkspaceOptions {
145- WorkspaceID : workspaceID ,
146- }); err != nil {
147- fmt .Fprintf (os .Stderr , "Error deleting revision %s: %v\n " , rev .RevisionID , err )
148- }
149- }
150-
151- return nil
152- }
153-
154- // hash computes the SHA-256 hash of the given data and returns it as a hexadecimal string
155- func hash (data []byte ) string {
156- if data == nil {
157- return ""
158- }
159- hash := sha256 .Sum256 (data )
160- return hex .EncodeToString (hash [:])
161- }
0 commit comments