From d57168db5eebc0ccdfd0af650cb5815434d6193a Mon Sep 17 00:00:00 2001
From: Xiaofan Zhou <81600993+xfz11@users.noreply.github.com>
Date: Wed, 6 Mar 2024 15:40:14 +0800
Subject: [PATCH] {ServiceConnector}: Fix user-agent missing and getting
 postgres db exception (#7342)

* fix error

* update

* update

* update

* update

* update user agent

* lint

* updated
---
 src/serviceconnector-passwordless/HISTORY.rst      |  4 ++++
 .../_client_factory.py                             | 14 +++++---------
 .../_credential_free.py                            | 12 ++++++++++++
 .../azext_serviceconnector_passwordless/config.py  |  2 +-
 src/serviceconnector-passwordless/setup.py         |  2 +-
 5 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/serviceconnector-passwordless/HISTORY.rst b/src/serviceconnector-passwordless/HISTORY.rst
index 0d24af75c00..041a80a2d4d 100644
--- a/src/serviceconnector-passwordless/HISTORY.rst
+++ b/src/serviceconnector-passwordless/HISTORY.rst
@@ -2,6 +2,10 @@
 
 Release History
 ===============
+1.0.2
+++++++
+* Bypass error of `az postgres flexible-server db show`
+
 1.0.1
 ++++++
 * Make some improvements to fix non-json output issue
diff --git a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_client_factory.py b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_client_factory.py
index cd371074f2e..fc7b86bf185 100644
--- a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_client_factory.py
+++ b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_client_factory.py
@@ -2,22 +2,18 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+import os
 
 
 # pylint: disable=consider-using-f-string
 def cf_connection_cl(cli_ctx, *_):
     from azure.mgmt.servicelinker import ServiceLinkerManagementClient
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client, get_az_user_agent
-    from azure.core.pipeline import policies
+    from azure.cli.core.commands.client_factory import get_mgmt_service_client
     from .config import NAME, VERSION
-
-    user_agent_policy = policies.UserAgentPolicy(
-        user_agent=get_az_user_agent())
-    user_agent_policy.add_user_agent(
-        "CliExtension/{}({})".format(NAME, VERSION))
+    os.environ['AZURE_HTTP_USER_AGENT'] = (os.environ.get('AZURE_HTTP_USER_AGENT')
+                                           or '') + " CliExtension/{}({})".format(NAME, VERSION)
     return get_mgmt_service_client(cli_ctx, ServiceLinkerManagementClient,
-                                   subscription_bound=False, api_version="2022-11-01-preview",
-                                   user_agent_policy=user_agent_policy)
+                                   subscription_bound=False, api_version="2022-11-01-preview")
 
 
 def cf_linker(cli_ctx, *_):
diff --git a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py
index c5c7bdfc889..0eea328446c 100644
--- a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py
+++ b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py
@@ -641,6 +641,18 @@ def __init__(self, cmd, target_id, target_type, auth_info, connection_name, skip
 
     def check_db_existence(self):
         try:
+            # `az postgres flexible-server db show -d postgres` will throw exception
+            if self.dbname == "postgres":
+                server_info = run_cli_cmd(
+                    'az postgres flexible-server show -n {} -g {} --subscription {}'.format(
+                        self.db_server, self.resource_group, self.subscription))
+                if server_info is None:
+                    e = ResourceNotFoundError(
+                        "No server found for '{}'".format(self.db_server))
+                    telemetry.set_exception(e, "No-Server")
+                    raise e
+                else:
+                    return
             db_info = run_cli_cmd(
                 'az postgres flexible-server db show --server-name {} --database-name {} -g {} --subscription {}'.format(
                     self.db_server, self.dbname, self.resource_group, self.subscription))
diff --git a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py
index b0b962702a3..78cfa671e19 100644
--- a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py
+++ b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py
@@ -4,5 +4,5 @@
 # --------------------------------------------------------------------------------------------
 
 
-VERSION = '1.0.1'
+VERSION = '1.0.2'
 NAME = 'serviceconnector-passwordless'
diff --git a/src/serviceconnector-passwordless/setup.py b/src/serviceconnector-passwordless/setup.py
index 091d10647ed..89442b4d374 100644
--- a/src/serviceconnector-passwordless/setup.py
+++ b/src/serviceconnector-passwordless/setup.py
@@ -15,7 +15,7 @@
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
 
 
-VERSION = '1.0.1'
+VERSION = '1.0.2'
 try:
     from azext_serviceconnector_passwordless.config import VERSION
 except ImportError: