Write-DbaDbTableData: Fix for connections to Azure SQL Database with …

…AccessToken (#9401)
dataplat · Jul 1, 2024 · fe48ad8 · fe48ad8
1 parent 00e3f97
commit fe48ad8
Show file tree

Hide file tree

Showing 7 changed files with 25 additions and 39 deletions.
diff --git a/public/Connect-DbaInstance.ps1 b/public/Connect-DbaInstance.ps1
@@ -142,6 +142,8 @@ function Connect-DbaInstance {
 
         Note that the token is valid for only one hour and cannot be renewed automatically.
 
+        Note that the returned SMO is not a fully functional SMO. It can only be used in a limited list of commands like Invoke-DbaQuery, Import-DbaCsv or Write-DbaDbTableData.
+
     .PARAMETER DedicatedAdminConnection
         Connects using "ADMIN:" to create a dedicated admin connection (DAC) as a non-pooled connection.
         If the instance is on a remote server, the remote access has to be enabled via "Set-DbaSpConfigure -Name RemoteDacConnectionsEnabled -Value $true" or "sp_configure 'remote admin connections', 1".

diff --git a/public/Get-DbaTopResourceUsage.ps1 b/public/Get-DbaTopResourceUsage.ps1
@@ -254,7 +254,7 @@ function Get-DbaTopResourceUsage {
     process {
         foreach ($instance in $SqlInstance) {
             try {
-                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -MinimumVersion 10 -StatementTimeout 0
+                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -MinimumVersion 10
             } catch {
                 Stop-Function -Message "Failure" -Category ConnectionError -ErrorRecord $_ -Target $instance -Continue
             }

diff --git a/public/Import-DbaCsv.ps1 b/public/Import-DbaCsv.ps1
@@ -456,7 +456,7 @@ function Import-DbaCsv {
                 $elapsed = [System.Diagnostics.Stopwatch]::StartNew()
                 # Open Connection to SQL Server
                 try {
-                    $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -Database $Database -StatementTimeout 0 -MinimumVersion 9
+                    $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -Database $Database -MinimumVersion 9
                     $sqlconn = $server.ConnectionContext.SqlConnectionObject
                     if ($sqlconn.State -ne 'Open') {
                         $sqlconn.Open()

diff --git a/public/Invoke-DbaDbUpgrade.ps1 b/public/Invoke-DbaDbUpgrade.ps1
@@ -132,7 +132,7 @@ function Invoke-DbaDbUpgrade {
 
         foreach ($instance in $SqlInstance) {
             try {
-                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -StatementTimeout 0
+                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential
             } catch {
                 Stop-Function -Message "Failure" -Category ConnectionError -ErrorRecord $_ -Target $instance -Continue
             }

diff --git a/public/Set-DbaDbCompression.ps1 b/public/Set-DbaDbCompression.ps1
@@ -128,7 +128,7 @@ function Set-DbaDbCompression {
         $starttime = Get-Date
         foreach ($instance in $SqlInstance) {
             try {
-                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -MinimumVersion 10 -StatementTimeout 0
+                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -MinimumVersion 10
             } catch {
                 Stop-Function -Message "Failure" -Category ConnectionError -ErrorRecord $_ -Target $instance -Continue
             }

diff --git a/public/Write-DbaDbTableData.ps1 b/public/Write-DbaDbTableData.ps1
@@ -468,6 +468,7 @@ function Write-DbaDbTableData {
         }
         #endregion Connect to server
 
+        #region Resolve Full Qualified Table Name
         if ($server.ServerType -ne 'SqlAzureDatabase') {
             <#
                 Skip adding database name to Fully Qualified Tablename for Azure SQL DB
@@ -502,47 +503,30 @@ function Write-DbaDbTableData {
         Write-Message -Level SomewhatVerbose -Message "FQTN processed: $fqtn"
         #endregion Resolve Full Qualified Table Name
 
-
-        #region Get database
-        # we used to do a try catch on $server.Databases if $server.ServerType -eq 'SqlAzureDatabase' here
-        # but it seems this was fixed in the newest SMO
-        try {
-            # This works for both onprem and azure -- using a hash only works for onprem
-            $databaseObject = $server.Databases | Where-Object Name -eq $databaseName
-            #endregion Get database
-
-            #region Prepare database and bulk operations
-            if ($null -eq $databaseObject) {
-                Stop-Function -Message "Database $databaseName does not exist." -Target $SqlInstance
-                return
-            }
-
-            $databaseObject.Tables.Refresh()
-            if ($schemaName -notin $databaseObject.Schemas.Name) {
-                Stop-Function -Message "Schema $schemaName does not exist."
-                return
+        #region Test if table exists
+        if ($tableName.StartsWith('#')) {
+            try {
+                Write-Message -Level Verbose -Message "The table $tableName should be in tempdb and we try to find it."
+                $null = $server.ConnectionContext.ExecuteScalar("SELECT TOP(1) 1 FROM [$tableName]")
+                $tableExists = $true
+            } catch {
+                $tableExists = $false
             }
-
-            if ($tableName.StartsWith('#')) {
-                try {
-                    Write-Message -Level Verbose -Message "The table $tableName should be in tempdb and we try to find it."
-                    $null = $databaseObject.Query("SELECT TOP(1) 1 FROM [$tableName]")
-                    $tableExists = $true
-                } catch {
-                    $tableExists = $false
-                }
-            } else {
-                $targetTable = $databaseObject.Tables | Where-Object { $_.Name -eq $tableName -and $_.Schema -eq $schemaName }
-                $tableExists = $targetTable.Count -eq 1
+        } else {
+            # We don't use SMO here because it does not work for Azure SQL Database connected with AccessToken.
+            try {
+                $null = $server.ConnectionContext.ExecuteScalar("SELECT TOP(1) 1 FROM $fqtn")
+                $tableExists = $true
+            } catch {
+                $tableExists = $false
             }
-        } catch {
-            Stop-Function -Message "Failure" -ErrorRecord $_ -Continue
         }
 
         if ((-not $tableExists) -and (-not $AutoCreateTable)) {
             Stop-Function -Message "Table does not exist and automatic creation of the table has not been selected. Specify the '-AutoCreateTable'-parameter to generate a suitable table."
             return
         }
+        #endregion Test if table exists
 
         $bulkCopyOptions = 0
         $options = "TableLock", "CheckConstraints", "FireTriggers", "KeepIdentity", "KeepNulls", "Default"

diff --git a/tests/Connect-DbaInstance.Tests.ps1 b/tests/Connect-DbaInstance.Tests.ps1
@@ -195,9 +195,9 @@ Describe "$commandname Integration Tests" -Tags "IntegrationTests" {
         }
 
         It "clones when using parameter StatementTimeout" {
-            $serverClone = Connect-DbaInstance -SqlInstance $server -StatementTimeout 0
+            $serverClone = Connect-DbaInstance -SqlInstance $server -StatementTimeout 123
             $server.ConnectionContext.StatementTimeout | Should -Be (Get-DbatoolsConfigValue -FullName 'sql.execution.timeout')
-            $serverClone.ConnectionContext.StatementTimeout | Should -Be 0
+            $serverClone.ConnectionContext.StatementTimeout | Should -Be 123
         }
 
         It "clones when using parameter DedicatedAdminConnection" {