[RDBMS] Add flexible-server connect command

Dockerfile

@@ -36,7 +36,7 @@ LABEL maintainer="Microsoft" \
 # libintl and icu-libs - required by azure devops artifact (az extension add --name azure-devops)
 RUN apk add --no-cache bash openssh ca-certificates jq curl openssl perl git zip \
  && apk add --no-cache --virtual .build-deps gcc make openssl-dev libffi-dev musl-dev linux-headers \
- && apk add --no-cache libintl icu-libs libc6-compat \
+ && apk add --no-cache libintl icu-libs libc6-compat postgresql-dev \
  && apk add --no-cache bash-completion \
  && update-ca-certificates
 

azure-pipelines.yml

@@ -107,7 +107,7 @@ jobs:
   - task: UsePythonVersion@0
     displayName: 'Use Python 3'
     inputs:
-      versionSpec: 3.x
+      versionSpec: 3.8
 
   - bash: ./scripts/ci/dependency_check.sh
     displayName: 'Verify src/azure-cli/requirements.py3.Darwin.txt'

src/azure-cli/azure/cli/command_modules/rdbms/_helptext_mysql.py

@@ -291,3 +291,13 @@
     text: az mysql flexible-server wait --exists --resource-group testGroup --name testServer
     crafted: true
 """
+
+helps['mysql flexible-server connect'] = """
+type: command
+short-summary: Connect to a flexible server.
+example:
+  - name: Connect to a flexible server using administrator username and password. Default database is used if not specified.
+    text: az mysql flexible-server connect -u username -p password -d flexibleserverdb
+  - name: Connect to a flexible server and execute a query with results outputted in a table.
+    text: az mysql flexible-server connect -u username -p password --mysql-query "select host, user from mysql.user;" --output table
+"""

src/azure-cli/azure/cli/command_modules/rdbms/_helptext_pg.py

@@ -227,3 +227,13 @@
     text: az postgres server wait --exists --resource-group testGroup --name testServer
     crafted: true
 """
+
+helps['postgres flexible-server connect'] = """
+type: command
+short-summary: Connect to a flexible server.
+example:
+  - name: Connect to a flexible server using administrator username and password. Default database is used if not specified.
+    text: az postgres flexible-server connect -u username -p password -d postgres
+  - name: Connect to a flexible server and execute a query with results outputted in a table.
+    text: az postgres flexible-server connect -u username -p password --postgres-query "select * from pg_user;" --output table
+"""

src/azure-cli/azure/cli/command_modules/rdbms/_params.py

@@ -337,6 +337,18 @@ def _flexible_server_params(command_group):
         with self.argument_context('{} flexible-server list-skus'.format(command_group)) as c:
             c.argument('location', arg_type=get_location_type(self.cli_ctx))
 
+        with self.argument_context('{} flexible-server connect'.format(command_group)) as c:
+            c.argument('server_name', id_part=None, options_list=['--name', '-n'], arg_type=server_name_arg_type)
+            c.argument('administrator_login', arg_group='Authentication', arg_type=administrator_login_arg_type, options_list=['--admin-user', '-u'],
+                       help='The login username of the administrator.')
+            c.argument('administrator_login_password', arg_group='Authentication', options_list=['--admin-password', '-p'],
+                       help='The login password of the administrator.')
+            c.argument('database_name', arg_type=database_name_arg_type, options_list=['--database-name', '-d'], help='The name of a database.')
+            if command_group == "mysql":
+                c.argument('mysql_query', options_list=['--mysql-query', '-c'], help='A query to run against the flexible server.')
+            elif command_group == "postgres":
+                c.argument('postgres_query', options_list=['--postgres-query', '-c'], help='A query to run against the flexible server.')
+
         # flexible-server parameter
         for scope in ['list', 'set', 'show']:
             argument_context_string = '{} flexible-server parameter {}'.format(command_group, scope)

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_commands.py

@@ -136,6 +136,7 @@ def load_flexibleserver_command_table(self, _):
                             is_preview=True) as g:
         g.custom_command('list-skus', 'flexible_list_skus', table_transformer=table_transform_output_list_skus)
         g.custom_command('show-connection-string', 'flexible_server_connection_string')
+        g.custom_command('connect', 'connect_to_flexible_server_postgresql')
 
     # MySQL commands
     with self.command_group('mysql flexible-server', mysql_flexible_servers_sdk,
@@ -193,6 +194,7 @@ def load_flexibleserver_command_table(self, _):
                             is_preview=True) as g:
         g.custom_command('list-skus', 'flexible_list_skus', table_transformer=table_transform_output_list_skus)
         g.custom_command('show-connection-string', 'flexible_server_connection_string')
+        g.custom_command('connect', 'connect_to_flexible_server_mysql')
 
     with self.command_group('mysql flexible-server replica', mysql_flexible_replica_sdk,
                             is_preview=True) as g:

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_custom_mysql.py

@@ -8,6 +8,7 @@
 from msrestazure.azure_exceptions import CloudError
 from msrestazure.tools import resource_id, is_valid_resource_id, parse_resource_id  # pylint: disable=import-error
 from knack.log import get_logger
+import mysql.connector as mysql_connector
 from azure.cli.core.commands.client_factory import get_subscription_id
 from azure.cli.core.util import CLIError, sdk_no_wait
 from azure.cli.core.local_context import ALL
@@ -397,9 +398,10 @@ def _create_server(db_context, cmd, resource_group_name, server_name, location,
 
 
 def flexible_server_connection_string(
-        server_name='{server}', database_name='{database}', administrator_login='{login}',
+        cmd, server_name='{server}', database_name='{database}', administrator_login='{login}',
         administrator_login_password='{password}'):
-    host = '{}.mysql.database.azure.com'.format(server_name)
+    mysql_server_endpoint = cmd.cli_ctx.cloud.suffixes.mysql_server_endpoint
+    host = '{}{}'.format(server_name, mysql_server_endpoint)
     if database_name is None:
         database_name = 'mysql'
     return {
@@ -408,6 +410,53 @@ def flexible_server_connection_string(
     }
 
 
+def connect_to_flexible_server_mysql(cmd, server_name, administrator_login, administrator_login_password, database_name=None, mysql_query=None):
+    mysql_server_endpoint = cmd.cli_ctx.cloud.suffixes.mysql_server_endpoint
+    host = '{}{}'.format(server_name, mysql_server_endpoint)
+    cursor = None
+    json_data = None
+    if database_name is None:
+        database_name = DEFAULT_DB_NAME
+        logger.warning("Connecting to %s database by default.", DEFAULT_DB_NAME)
+    # Connect to mysql and get cursor to run sql commands
+    try:
+        connection_kwargs = {
+            'host': host,
+            'database': database_name,
+            'user': administrator_login,
+            'password': administrator_login_password
+        }
+        connection = mysql_connector.connect(**connection_kwargs)
+        logger.warning('Successfully connected to %s.', server_name)
+    except Exception as e:
+        logger.warning('Failed connection to %s. Check error and validate firewall and public access settings.', server_name)
+        raise CLIError("Unable to connect to flexible server: {}".format(e))
+
+    # execute query if passed in
+    if mysql_query is not None:
+        try:
+            cursor = connection.cursor()
+            cursor.execute(mysql_query)
+            logger.warning("Ran Database Query: '%s'", mysql_query)
+            logger.warning("Retrieving first 30 rows of query output.")
+            result = cursor.fetchmany(30)  # limit to 30 rows of output for now
+            row_headers = [x[0] for x in cursor.description]  # this will extract row headers
+            # format the result for a clean display
+            json_data = []
+            for rv in result:
+                json_data.append(dict(zip(row_headers, rv)))
+        except mysql_connector.errors.Error as e:
+            raise CLIError("Unable to execute query '{0}': {1}".format(mysql_query, e))
+        finally:
+            if cursor is not None:
+                try:
+                    cursor.close()
+                    logger.warning("Closed the connection to %s", server_name)
+                except Exception:  # pylint: disable=broad-except
+                    logger.warning('Unable to close connection cursor.')
+    return json_data
+
+
 def _create_mysql_connection_strings(host, user, password, database):
     result = {
         'mysql_cmd': "mysql {database} --host {host} --user {user} --password={password}",

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_custom_postgres.py

@@ -8,6 +8,7 @@
 from msrestazure.azure_exceptions import CloudError
 from msrestazure.tools import resource_id, is_valid_resource_id, parse_resource_id  # pylint: disable=import-error
 from knack.log import get_logger
+import psycopg2
 from azure.cli.core.commands.client_factory import get_subscription_id
 from azure.cli.core.local_context import ALL
 from azure.cli.core.util import CLIError, sdk_no_wait
@@ -333,15 +334,64 @@ def _create_server(db_context, cmd, resource_group_name, server_name, location,
 
 
 def flexible_server_connection_string(
-        server_name='{server}', database_name='{database}', administrator_login='{login}',
+        cmd, server_name='{server}', database_name='{database}', administrator_login='{login}',
         administrator_login_password='{password}'):
-    host = '{}.postgres.database.azure.com'.format(server_name)
+    postgresql_server_endpoint = cmd.cli_ctx.cloud.suffixes.postgresql_server_endpoint
+    host = '{}{}'.format(server_name, postgresql_server_endpoint)
     return {
         'connectionStrings': _create_postgresql_connection_strings(host, administrator_login,
                                                                    administrator_login_password, database_name)
     }
 
 
+def connect_to_flexible_server_postgresql(cmd, server_name, administrator_login, administrator_login_password, database_name=None, postgres_query=None):
+    postgresql_server_endpoint = cmd.cli_ctx.cloud.suffixes.postgresql_server_endpoint
+    host = '{}{}'.format(server_name, postgresql_server_endpoint)
+    cursor = None
+    json_data = None
+    if database_name is None:
+        logger.warning("Connecting to postgres database by default.")
+        database_name = "postgres"
+    # Connect to mysql and get cursor to run sql commands
+    try:
+        connection_kwargs = {
+            'host': host,
+            'database': database_name,
+            'user': administrator_login,
+            'password': administrator_login_password
+        }
+        connection = psycopg2.connect(**connection_kwargs)
+        connection.set_session(autocommit=True)
+        logger.warning('Successfully connected to %s.', server_name)
+    except Exception as e:
+        logger.warning('Failed connection to %s. Check error and validate firewall and public access settings.', server_name)
+        raise CLIError("Unable to connect to flexible server: {}".format(e))
+
+    # execute query if passed in
+    if postgres_query is not None:
+        try:
+            cursor = connection.cursor()
+            cursor.execute(postgres_query)
+            logger.warning("Ran Database Query: '%s'", postgres_query)
+            logger.warning("Retrieving first 30 rows of query output.")
+            result = cursor.fetchmany(30)  # limit to 30 rows of output for now
+            row_headers = [x[0] for x in cursor.description]  # this will extract row headers
+            # format the result for a clean display
+            json_data = []
+            for rv in result:
+                json_data.append(dict(zip(row_headers, rv)))
+        except Exception as e:
+            raise CLIError("Unable to execute query '{0}': {1}".format(postgres_query, e))
+        finally:
+            if cursor is not None:
+                try:
+                    cursor.close()
+                    logger.warning("Closed the connection to %s", server_name)
+                except Exception:  # pylint: disable=broad-except
+                    logger.warning('Unable to close connection cursor.')
+    return json_data
+
+
 def _create_postgresql_connection_strings(host, user, password, database):
     result = {
         'psql_cmd': "postgresql://{user}:{password}@{host}/postgres?sslmode=require",