SNOW-1335075: SNOW-1209099 Cannot connect using private key auth over private link with regionless account notation

Please answer these questions before submitting your issue. Thanks!

What version of Python are you using?

Python 3.11.8 (main, Feb 6 2024, 21:21:21) [Clang 15.0.0 (clang-1500.1.0.2.5)]
What operating system and processor architecture are you using?

macOS-14.4.1-x86_64-i386-64bit
What are the component versions in the environment (pip freeze)?

annotated-types==0.6.0 anyio==4.3.0 asn1crypto==1.5.1 asttokens==2.4.1 attrs==23.2.0 bandit==1.7.7 -e git+https://github.com/biocommons/anyvar.git@ab96430665cbaa1697e3d69ea3a0b1ffb407ff95#egg=biocommons.anyvar biocommons.seqrepo==0.6.9 bioutils==0.5.8.post1 black==24.2.0 canonicaljson==2.0.0 certifi==2024.2.2 cffi==1.16.0 cfgv==3.4.0 charset-normalizer==3.3.2 click==8.1.7 coloredlogs==15.0.1 configparser==6.0.1 coverage==7.4.3 cryptography==41.0.7 decorator==5.1.1 dill==0.3.8 distlib==0.3.8 executing==2.0.1 fastapi==0.110.0 filelock==3.13.1 ga4gh.vrs @ file:///workspaces/gdh/vrs-python greenlet==3.0.3 gunicorn==21.2.0 h11==0.14.0 hgvs==1.5.4 httpcore==1.0.4 httpx==0.27.0 humanfriendly==10.0 identify==2.5.35 idna==3.6 importlib-metadata==7.0.1 iniconfig==2.0.0 ipython==8.22.1 jedi==0.19.1 markdown-it-py==3.0.0 matplotlib-inline==0.1.6 mdurl==0.1.2 mypy-extensions==1.0.0 nodeenv==1.8.0 numpy==1.26.4 packaging==23.2 Parsley==1.3 parso==0.8.3 pathspec==0.12.1 pbr==6.0.0 pexpect==4.9.0 platformdirs==3.11.0 pluggy==1.4.0 pre-commit==3.6.2 prompt-toolkit==3.0.43 psycopg==3.1.18 psycopg-binary==3.1.18 psycopg-pool==3.2.1 psycopg2==2.9.9 psycopg2-binary==2.9.9 ptyprocess==0.7.0 pure-eval==0.2.2 pycparser==2.21 pydantic==2.6.2 pydantic_core==2.16.3 Pygments==2.17.2 PyJWT==2.8.0 pyOpenSSL==23.3.0 pysam==0.22.0 pytest==8.0.2 pytest-cov==4.1.0 pytest-mock==3.12.0 python-multipart==0.0.9 pytz==2024.1 PyYAML==6.0.1 requests==2.31.0 rich==13.7.0 ruff==0.2.2 six==1.16.0 sniffio==1.3.1 snowflake-connector-python==3.4.1 snowflake-sqlalchemy==1.5.1 sortedcontainers==2.4.0 SQLAlchemy==1.4.51 sqlparse==0.4.4 stack-data==0.6.3 starlette==0.36.3 stevedore==5.2.0 tabulate==0.9.0 tomlkit==0.12.3 tqdm==4.66.2 traitlets==5.14.1 typing_extensions==4.10.0 urllib3==1.26.18 uvicorn==0.27.1 virtualenv==20.25.1 wcwidth==0.2.13 yoyo-migrations==8.2.0 zipp==3.17.0
What did you do?

from snowflake.sqlalchemy import URL
from sqlalchemy import create_engine

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization

with open("rsa_key.p8", "rb") as key:
    p_key = serialization.load_pem_private_key(
        key.read(),
        password=None,
        backend=default_backend(),
    )

pkb = p_key.private_bytes(
    encoding=serialization.Encoding.DER,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption(),
)

engine = create_engine(
    URL(
        account="mgb-mgbdev.privatelink",
        user="***REDACTED****",
    ),
    connect_args={
        "private_key": pkb,
    },
)

engine.connect()

Results in error snowflake.connector.errors.DatabaseError: 250001 (08001): None: Failed to connect to DB: mgb-mgbdev.privatelink.snowflakecomputing.com:443. JWT token is invalid. [45834078-d20d-4e94-b034-7dcbb157e096]

The root cause is that the SnowflakeDialect.create_connect_args() returning account=mgb, which in turns makes the subject of the JWT incorrect. Instead, the account should be mgb-mgbdev.

I have implemented a monkey patch to workaround the issue:

# Create a new pointer to the existing create_connect_args method
SnowflakeDialect._orig_create_connect_args = SnowflakeDialect.create_connect_args

# Define a new create_connect_args method that calls the original method
#   and then fixes the result so that the account name is not mangled
#   when using privatelink
def sf_create_connect_args_override(self, url: URL):

    # retval is tuple of empty array and dict ([], {})
    retval = self._orig_create_connect_args(url)

    # the dict has the options including the mangled account name
    opts = retval[1]
    if "host" in opts and "account" in opts and ".privatelink.snowflakecomputing.com" in opts["host"]:
        opts["account"] = opts["host"].split(".")[0]

    return retval

# Replace the create_connect_args method with the override
SnowflakeDialect.create_connect_args = sf_create_connect_args_override

snowflakedb / snowflake-sqlalchemy

SNOW-1335075: SNOW-1209099 Cannot connect using private key auth over private link with regionless account notation #489