Public/Connect-OnePAMDatabase.ps1
|
function Connect-OnePAMDatabase { <# .SYNOPSIS Opens an interactive database session to a OnePAM resource. .DESCRIPTION Connects to a database resource through the OnePAM gateway and provides an interactive query REPL with formatted table output. Supports PostgreSQL, MySQL, MSSQL, MongoDB, and Elasticsearch. Metacommands: \q quit, \h help, \dt tables, \d describe, \databases list DBs. .PARAMETER Resource Resource name or UUID. .EXAMPLE Connect-OnePAMDatabase -Resource "prod-postgres" .EXAMPLE Connect-OnePAMDatabase -Resource "my-mongodb" #> [CmdletBinding()] param( [Parameter(Mandatory, Position = 0)] [string]$Resource ) $res = Get-OnePAMResource -Name $Resource if (-not $res) { throw "Resource '$Resource' not found." } $resId = if ($res.id) { $res.id } elseif ($res.ID) { $res.ID } else { throw "Cannot determine resource ID." } Assert-OpSafePathSegment -Value $resId -Label 'resource ID' if ($res.type -and $res.type -ne 'database') { throw "Resource '$Resource' is type '$($res.type)', not database." } $session = Invoke-OpApi -Method POST -Path "/api/v1/resources/$resId/connect" if (-not $session.session_id) { throw 'Failed to create session.' } Assert-OpSafePathSegment -Value $session.session_id -Label 'session ID' Write-Host "Connecting to database: $Resource" -ForegroundColor Cyan Write-Host "Session: $($session.session_id)" if ($session.proxy_host) { Write-Host "Gateway: $($session.proxy_host)" } Write-Host '' $cfg = Get-OpConfig if ($session.direct) { $baseUri = [Uri]$cfg.api_base $wsScheme = if ($baseUri.Scheme -eq 'https') { 'wss' } else { 'ws' } $hostPort = $baseUri.Authority } else { if (-not $session.proxy_host) { throw 'Session response missing proxy_host.' } $wsScheme = 'wss' $hostPort = $session.proxy_host if ($session.proxy_port -and $session.proxy_port -ne 443) { $hostPort = "$($session.proxy_host):$($session.proxy_port)" } } $queryParts = [System.Collections.Generic.List[string]]::new() if (-not $session.direct -and $session.token) { $queryParts.Add("token=$([System.Uri]::EscapeDataString($session.token))") } $wsUrl = "${wsScheme}://${hostPort}/gateway/db/$($session.session_id)" if ($queryParts.Count -gt 0) { $wsUrl += "?$($queryParts -join '&')" } $wsHeaders = @{} if ($session.direct) { $authToken = Get-OpToken if ($authToken) { $wsHeaders['Authorization'] = "Bearer $($authToken.access_token)" } } $ws = New-OpWebSocket -Uri $wsUrl -Headers $wsHeaders try { $connResp = Receive-OpWsMessage -WebSocket $ws -TimeoutSeconds 30 if ($connResp.MessageType -eq 'Close') { throw 'Connection closed by server.' } $connMsg = $connResp.Text | ConvertFrom-Json if ($connMsg.type -eq 'error') { throw "Database connection failed: $($connMsg.message)" } if ($connMsg.type -ne 'connected') { throw "Unexpected server response: $($connMsg.type)" } $driver = $connMsg.driver $dbVersion = if ($connMsg.version) { $connMsg.version } else { $driver } Write-Host "Connected to: $dbVersion" -ForegroundColor Green Write-Host 'Type \q to quit, \h for help.' Write-Host '' $isDocDB = $driver -in @('mongodb', 'mongo', 'elasticsearch', 'elastic') $queryCounter = 0 $prompt = "$Resource=> " while ($true) { Write-Host -NoNewline $prompt $line = Read-Host if ($null -eq $line) { break } $line = $line.Trim() if ($line -eq '') { continue } if ($line.StartsWith('\')) { $metaResult = Invoke-OpDbMetaCommand -Line $line -Driver $driver -WebSocket $ws -Resource $Resource if ($metaResult -eq 'quit') { break } if ($metaResult -eq 'handled') { continue } if ($metaResult -and $metaResult -ne 'unknown') { $line = $metaResult } elseif ($metaResult -eq 'unknown') { continue } } $queryText = $line if (-not $isDocDB) { $upperTrimmed = $queryText.TrimStart().ToUpper() $needsSemicolon = -not $queryText.TrimEnd().EndsWith(';') -and -not $upperTrimmed.StartsWith('SELECT') -and -not $upperTrimmed.StartsWith('SHOW') -and -not $upperTrimmed.StartsWith('DESCRIBE') while ($needsSemicolon -and -not $queryText.TrimEnd().EndsWith(';')) { Write-Host -NoNewline ' -> ' $nextLine = Read-Host if ($null -eq $nextLine) { break } $queryText += "`n$nextLine" } } $queryCounter++ $queryID = "q$queryCounter" $queryMsg = @{ type = 'query' id = $queryID sql = $queryText } | ConvertTo-Json -Compress Send-OpWsMessage -WebSocket $ws -Message $queryMsg Read-OpQueryResults -WebSocket $ws -QueryId $queryID } } finally { Close-OpWebSocket -WebSocket $ws Write-Host 'Disconnected.' -ForegroundColor Gray } } function Invoke-OpDbMetaCommand { param( [string]$Line, [string]$Driver, [System.Net.WebSockets.ClientWebSocket]$WebSocket, [string]$Resource ) switch -Regex ($Line) { '^\\q(uit)?$' { return 'quit' } '^\\h(elp)?$' { Write-Host @" Common commands: \q, \quit Quit \h, \help Show this help \databases List databases "@ switch -Wildcard ($Driver) { 'mongo*' { Write-Host @" MongoDB commands: \collections List collections \show dbs List all databases \show collections List collections Enter queries as JSON. "@ } 'elastic*' { Write-Host @" Elasticsearch commands: \indices List indices \health Cluster health \info Cluster info \mapping <index> Show index mapping Enter SQL queries or use the query protocol. "@ } default { Write-Host @" SQL commands: \dt List tables \d List tables with types \d <table> Describe table columns Enter SQL queries directly. Multi-line queries end with a semicolon. "@ } } return 'handled' } '^\\dt$' { switch ($Driver) { 'mysql' { return 'SHOW TABLES' } { $_ -in 'mssql', 'sqlserver' } { return "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE = 'BASE TABLE' ORDER BY TABLE_NAME" } { $_ -in 'mongodb', 'mongo', 'elasticsearch', 'elastic' } { return 'handled' } default { return "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename" } } } '^\\d$' { switch ($Driver) { 'mysql' { return 'SHOW TABLES' } { $_ -in 'mssql', 'sqlserver' } { return "SELECT TABLE_NAME, TABLE_TYPE FROM INFORMATION_SCHEMA.TABLES ORDER BY TABLE_NAME" } default { return "SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = 'public' ORDER BY table_name" } } } '^\\d\s+(.+)$' { $table = $Matches[1].Trim() if ($table -notmatch '^[a-zA-Z0-9_.\-]+$') { Write-Host "Invalid table name: $table" -ForegroundColor Red return 'handled' } switch ($Driver) { 'mysql' { return "DESCRIBE ``$table``" } { $_ -in 'mssql', 'sqlserver' } { return "SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '$table' ORDER BY ORDINAL_POSITION" } default { return "SELECT column_name, data_type, is_nullable, column_default FROM information_schema.columns WHERE table_name = '$table' AND table_schema = 'public' ORDER BY ordinal_position" } } } '^\\(databases|show dbs)$' { switch ($Driver) { { $_ -in 'mongodb', 'mongo' } { $msg = @{ type = 'query'; id = 'meta'; operation = 'listDatabases' } | ConvertTo-Json -Compress Send-OpWsMessage -WebSocket $WebSocket -Message $msg Read-OpQueryResults -WebSocket $WebSocket -QueryId 'meta' return 'handled' } 'mysql' { return 'SHOW DATABASES' } { $_ -in 'mssql', 'sqlserver' } { return 'SELECT name FROM sys.databases ORDER BY name' } default { return 'SELECT datname FROM pg_database ORDER BY datname' } } } '^\\(collections|show collections)$' { if ($Driver -in @('mongodb', 'mongo')) { $msg = @{ type = 'query'; id = 'meta'; operation = 'listCollections' } | ConvertTo-Json -Compress Send-OpWsMessage -WebSocket $WebSocket -Message $msg Read-OpQueryResults -WebSocket $WebSocket -QueryId 'meta' return 'handled' } Write-Host '\collections is only available for MongoDB' -ForegroundColor Red return 'handled' } '^\\indices$' { if ($Driver -in @('elasticsearch', 'elastic')) { $msg = @{ type = 'query'; id = 'meta'; operation = 'cat_indices' } | ConvertTo-Json -Compress Send-OpWsMessage -WebSocket $WebSocket -Message $msg Read-OpQueryResults -WebSocket $WebSocket -QueryId 'meta' return 'handled' } Write-Host '\indices is only available for Elasticsearch' -ForegroundColor Red return 'handled' } '^\\health$' { if ($Driver -in @('elasticsearch', 'elastic')) { $msg = @{ type = 'query'; id = 'meta'; operation = 'cat_health' } | ConvertTo-Json -Compress Send-OpWsMessage -WebSocket $WebSocket -Message $msg Read-OpQueryResults -WebSocket $WebSocket -QueryId 'meta' return 'handled' } Write-Host '\health is only available for Elasticsearch' -ForegroundColor Red return 'handled' } '^\\info$' { if ($Driver -in @('elasticsearch', 'elastic')) { $msg = @{ type = 'query'; id = 'meta'; operation = 'info' } | ConvertTo-Json -Compress Send-OpWsMessage -WebSocket $WebSocket -Message $msg Read-OpQueryResults -WebSocket $WebSocket -QueryId 'meta' return 'handled' } Write-Host '\info is only available for Elasticsearch' -ForegroundColor Red return 'handled' } '^\\mapping\s+(.+)$' { if ($Driver -in @('elasticsearch', 'elastic')) { $index = $Matches[1].Trim() if (-not $index) { Write-Host 'Usage: \mapping <index>' -ForegroundColor Red return 'handled' } if ($index -notmatch '^[a-zA-Z0-9_.\-\*]+$') { Write-Host "Invalid index name: $index" -ForegroundColor Red return 'handled' } return "DESCRIBE `"$index`"" } Write-Host '\mapping is only available for Elasticsearch' -ForegroundColor Red return 'handled' } default { Write-Host "Unknown command: $Line (type \h for help)" -ForegroundColor Yellow return 'unknown' } } } function Read-OpQueryResults { param( [System.Net.WebSockets.ClientWebSocket]$WebSocket, [string]$QueryId ) $columns = @() $columnWidths = @() $allRows = [System.Collections.Generic.List[string[]]]::new() while ($true) { $resp = Receive-OpWsMessage -WebSocket $WebSocket -TimeoutSeconds 120 if ($resp.MessageType -eq 'Close') { Write-Host 'Connection closed.' -ForegroundColor Red return } $msg = $resp.Text | ConvertFrom-Json switch ($msg.type) { 'columns' { $columns = @($msg.columns) $columnWidths = @($columns | ForEach-Object { $_.Length }) } 'row' { $row = @() $values = @($msg.values) for ($i = 0; $i -lt $values.Count; $i++) { $val = if ($null -eq $values[$i]) { 'NULL' } else { "$($values[$i])" } $row += $val if ($i -lt $columnWidths.Count) { $w = [Math]::Min($val.Length, 60) if ($w -gt $columnWidths[$i]) { $columnWidths[$i] = $w } } } $allRows.Add($row) } 'complete' { if ($columns.Count -gt 0) { Write-OpResultTable -Columns $columns -Widths $columnWidths -Rows $allRows } $rowCount = if ($msg.rows_affected) { $msg.rows_affected } else { $allRows.Count } $plural = if ($rowCount -eq 1) { '' } else { 's' } $duration = if ($msg.duration_ms) { $msg.duration_ms } else { 0 } Write-Host "($rowCount row$plural, $([math]::Round($duration, 1)) ms)" Write-Host '' return } 'error' { Write-Host "ERROR: $($msg.message)" -ForegroundColor Red if ($msg.duration_ms) { Write-Host "($([math]::Round($msg.duration_ms, 1)) ms)" -ForegroundColor Red } Write-Host '' return } } } } function Write-OpResultTable { param( [string[]]$Columns, [int[]]$Widths, [System.Collections.Generic.List[string[]]]$Rows ) $headerParts = for ($i = 0; $i -lt $Columns.Count; $i++) { $w = if ($i -lt $Widths.Count) { $Widths[$i] } else { $Columns[$i].Length } " {0,-$w} " -f $Columns[$i] } Write-Host ($headerParts -join '|') $sepParts = for ($i = 0; $i -lt $Columns.Count; $i++) { $w = if ($i -lt $Widths.Count) { $Widths[$i] } else { $Columns[$i].Length } '-' * ($w + 2) } Write-Host ($sepParts -join '+') foreach ($row in $Rows) { $rowParts = for ($i = 0; $i -lt $row.Count; $i++) { $w = if ($i -lt $Widths.Count) { $Widths[$i] } else { $row[$i].Length } $display = $row[$i] if ($display.Length -gt 60) { $display = $display.Substring(0, 57) + '...' } " {0,-$w} " -f $display } Write-Host ($rowParts -join '|') } } |