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

[ehclark]

Please answer these questions before submitting your issue. Thanks!

1. 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)]

2. What operating system and processor architecture are you using?

   macOS-14.4.1-x86_64-i386-64bit

3. 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

4. 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

[sfc-gh-dszmolka]

hi - thanks for raising this. Indeed this is a known issue for a while and we're working on resolving it. Thank you for sharing your approach too !

As an alternative mitigation, the account locator version (xy12345.eu-central-1.privatelink) can be too used in the account field; it should work as the issue only seems to affect the regionless privatelink notation (myorg-myaccount.privatelink) account format.