Bracket IPv6 literals in prompt strings

changelog.md

@@ -1,6 +1,11 @@
 Upcoming (TBD)
 ==============
 
+Features
+--------
+* Bracket IPv6 literals in prompt strings.
+
+
 Bug Fixes
 --------
 * Read `/source` files and the `\e` editor file as UTF-8 instead of the locale encoding, so non-ASCII SQL no longer fails or is mis-decoded under non-UTF-8 locales.

mycli/main_modes/repl.py

@@ -98,7 +98,7 @@
     need_completion_refresh,
     need_completion_reset,
 )
-from mycli.packages.utils.string_utils import sanitize_terminal_title
+from mycli.packages.utils.string_utils import bracket_ipv6_literal, sanitize_terminal_title
 from mycli.types import Query
 
 if TYPE_CHECKING:
@@ -321,6 +321,8 @@ def render_prompt_string(
     short_prompt_host, _, _ = prompt_host.partition('.')
     if re.match(r'^[\d\.]+$', short_prompt_host):
         short_prompt_host = prompt_host
+    prompt_host = bracket_ipv6_literal(prompt_host)
+    short_prompt_host = bracket_ipv6_literal(short_prompt_host)
     now = datetime.now()
     species_name = sql_execute.server_info.species.name if sql_execute.server_info and sql_execute.server_info.species else 'MySQL'
     strings = string.split('\\\\')

mycli/packages/special_commands/special_command_utils.py

@@ -9,6 +9,8 @@
 import pymysql
 from pymysql.cursors import Cursor
 
+from mycli.packages.utils.string_utils import bracket_ipv6_literal
+
 logger = logging.getLogger(__name__)
 
 CACHED_SSL_VERSION: dict[tuple, str | None] = {}
@@ -166,10 +168,7 @@ def format_connection_dsn(
 ) -> str:
     if user_part := urlquote(user or ''):
         user_part = f'{user_part}@'
-    host_part = host or 'localhost'
-    if ':' in host_part and not host_part.startswith('['):
-        # IPv6 literals must be bracketed so the port can be told apart
-        host_part = f'[{host_part}]'
+    host_part = bracket_ipv6_literal(host or 'localhost')
     port_part = f':{port}' if port is not None else ''
     if db_part := urlquote(database or ''):
         db_part = f'/{db_part}'

mycli/packages/utils/string_utils.py

@@ -1,3 +1,4 @@
+import ipaddress
 import re
 
 from cli_helpers.utils import strip_ansi
@@ -13,3 +14,13 @@ def sanitize_terminal_title(title: FormattedText) -> str:
     sanitized = sanitized.replace('\n', ' ')
     sanitized = re.sub('[\x00-\x1f\x7f]', '', sanitized)
     return sanitized
+
+
+def bracket_ipv6_literal(host: str) -> str:
+    try:
+        ip_obj = ipaddress.ip_address(host)
+    except ValueError:
+        return host
+    if isinstance(ip_obj, ipaddress.IPv6Address) and not host.startswith('['):
+        return f'[{host}]'
+    return host

mycli_test/pytests/test_main_modes_repl.py

@@ -1178,6 +1178,28 @@ def test_transaction_prompt_reuses_cached_render_without_ping() -> None:
     ping.assert_called_once_with()
 
 
+@pytest.mark.parametrize('escape', [r'\h', r'\H'])
+@pytest.mark.parametrize(
+    ('host', 'expected'),
+    [('::1', '[::1]:3306'), ('2001:db8::1', '[2001:db8::1]:3306'), ('fe80::1%en0', '[fe80::1%en0]:3306'), ('[::1]', '[::1]:3306')],
+)
+def test_render_prompt_string_brackets_ipv6_host(escape: str, host: str, expected: str) -> None:
+    cli = make_repl_cli(SimpleNamespace(user='alice', host=host, dbname='db', port=3306, socket=None, server_info=None, conn=None))
+
+    prompt = repl_mode.render_prompt_string(cli, escape + r':\p', 0)
+
+    assert to_plain_text(prompt) == expected
+
+
+@pytest.mark.parametrize('escape', [r'\h', r'\H'])
+def test_render_prompt_string_preserves_colon_in_login_path(escape: str) -> None:
+    cli = make_repl_cli(SimpleNamespace(user='alice', host='::1', dbname='db', port=3306, socket=None, server_info=None, conn=None))
+    cli.login_path = 'prod:readonly'
+    cli.login_path_as_host = True
+
+    assert to_plain_text(repl_mode.render_prompt_string(cli, escape, 0)) == 'prod:readonly'
+
+
 def test_render_prompt_string_includes_current_edit_mode() -> None:
     cli = make_repl_cli(
         SimpleNamespace(

mycli_test/pytests/test_special_utils.py

@@ -328,18 +328,20 @@ def test_format_connection_dsn_includes_ssh_jump() -> None:
     )
 
 
-def test_format_connection_dsn_brackets_ipv6_host() -> None:
+@pytest.mark.parametrize('host', ['::1', '2001:db8::1', 'fe80::1%en0', '::ffff:192.0.2.1'])
+@pytest.mark.parametrize('already_bracketed', [False, True])
+def test_format_connection_dsn_brackets_ipv6_host(host: str, already_bracketed: bool) -> None:
     dsn = format_connection_dsn(
         user='alice',
-        host='::1',
+        host=f'[{host}]' if already_bracketed else host,
         port=3307,
         database='prod',
         socket=None,
         character_set=None,
     )
-    assert dsn == 'mysql://alice@[::1]:3307/prod'
+    assert dsn == f'mysql://alice@[{host}]:3307/prod'
     uri = urlparse(dsn)
-    assert uri.hostname == '::1'
+    assert uri.hostname == host
     assert uri.port == 3307
 
 

mycli_test/pytests/test_string_utils.py

@@ -1,6 +1,23 @@
 # type: ignore
 
-from mycli.packages.utils.string_utils import sanitize_terminal_title
+import pytest
+
+from mycli.packages.utils.string_utils import bracket_ipv6_literal, sanitize_terminal_title
+
+
+@pytest.mark.parametrize('host', ['::1', '2001:db8::1', 'fe80::1%en0', '::ffff:192.0.2.1'])
+def test_bracket_ipv6_literal_wraps_addresses(host: str) -> None:
+    assert bracket_ipv6_literal(host) == f'[{host}]'
+
+
+@pytest.mark.parametrize('host', ['', 'localhost', 'db.example.com', '192.0.2.1', 'prod:readonly', '2001:db8::invalid'])
+def test_bracket_ipv6_literal_preserves_non_ipv6_values(host: str) -> None:
+    assert bracket_ipv6_literal(host) == host
+
+
+@pytest.mark.parametrize('host', ['[::1]', '[2001:db8::1]', '[fe80::1%en0]'])
+def test_bracket_ipv6_literal_preserves_existing_brackets(host: str) -> None:
+    assert bracket_ipv6_literal(host) == host
 
 
 def test_sanitize_terminal_title_strips_ansi_sequences() -> None: