gitlabds 2.3.1__tar.gz → 2.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. {gitlabds-2.3.1 → gitlabds-2.4.0}/PKG-INFO +10 -6
  2. {gitlabds-2.3.1 → gitlabds-2.4.0}/README.md +9 -5
  3. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/snowflake_connectors.py +90 -24
  4. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/PKG-INFO +10 -6
  5. {gitlabds-2.3.1 → gitlabds-2.4.0}/tests/test_snowflake_connectors.py +232 -59
  6. {gitlabds-2.3.1 → gitlabds-2.4.0}/LICENSE +0 -0
  7. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/__init__.py +0 -0
  8. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/baselines.py +0 -0
  9. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/config_generator.py +0 -0
  10. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/dummy.py +0 -0
  11. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/feature_descriptions.py +0 -0
  12. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/feature_reduction.py +0 -0
  13. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/insights.py +0 -0
  14. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/memory_optimization.py +0 -0
  15. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/missing.py +0 -0
  16. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/model_evaluator.py +0 -0
  17. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/monitoring_metrics.py +0 -0
  18. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/outliers.py +0 -0
  19. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/record_benchmarking.py +0 -0
  20. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/serving_features.py +0 -0
  21. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/split_data.py +0 -0
  22. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/trends.py +0 -0
  23. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/SOURCES.txt +0 -0
  24. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/dependency_links.txt +0 -0
  25. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/requires.txt +0 -0
  26. {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/top_level.txt +0 -0
  27. {gitlabds-2.3.1 → gitlabds-2.4.0}/pyproject.toml +0 -0
  28. {gitlabds-2.3.1 → gitlabds-2.4.0}/setup.cfg +0 -0
  29. {gitlabds-2.3.1 → gitlabds-2.4.0}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gitlabds
3
- Version: 2.3.1
3
+ Version: 2.4.0
4
4
  Summary: GitLab Data Science Tools
5
5
  Author-email: Kevin Dietz <kdietz@gitlab.com>
6
6
  License-Expression: MIT
@@ -1135,7 +1135,7 @@ feature_drift = results['feature_drift']
1135
1135
 
1136
1136
  #### Snowflake Connectors
1137
1137
 
1138
- Thin wrappers over the native Snowflake SDKs that resolve GitLab's two credential paths — local (dbt profile + SSO) and GitLab CI (env vars + password or key pair) — and then hand off to the native objects. Migrated from `gitlabdata.orchestration_utils`; see the migration notes at the end of this section.
1138
+ Thin wrappers over the native Snowflake SDKs that resolve GitLab's two credential paths — local (dbt profile + SSO) and GitLab CI (env vars + key-pair auth) — and then hand off to the native objects. Migrated from `gitlabdata.orchestration_utils`; see the migration notes at the end of this section.
1139
1139
 
1140
1140
  All of these require `pip install gitlabds[snowflake]`.
1141
1141
 
@@ -1144,7 +1144,7 @@ All of these require `pip install gitlabds[snowflake]`.
1144
1144
  #### Description
1145
1145
  Open a Snowflake connection, resolving credentials from your dbt profile locally or from GitLab CI environment variables in CI.
1146
1146
 
1147
- `gitlabds.snowflake_engine_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="dev", schema="", auth_method="password", profiles_path=None, **connect_kwargs):`
1147
+ `gitlabds.snowflake_engine_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="dev", schema="", auth_method="keypair", profiles_path=None, **connect_kwargs):`
1148
1148
 
1149
1149
  #### Parameters:
1150
1150
  - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
@@ -1152,10 +1152,12 @@ Open a Snowflake connection, resolving credentials from your dbt profile locally
1152
1152
  - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1153
1153
  - **_profile_target_** : dbt profile output name. Only used on the local path.
1154
1154
  - **_schema_** : Schema to connect to. Falls back to the dbt profile's schema locally; when it resolves empty the connection is opened with no default schema.
1155
- - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path. Key pair secrets are not yet provisioned in CI.
1155
+ - **_auth_method_** : CI authentication method. `"keypair"` is the only supported value and the default, so passing it explicitly is the same as omitting it. Retained so callers that pass `auth_method="keypair"` keep working. `"password"` is accepted but ignored, with a logged warning naming the fix; any other value raises `ValueError`.
1156
1156
  - **_profiles_path_** : Override the dbt `profiles.yml` location (defaults to `~/.dbt/profiles.yml`). Only used on the local path.
1157
1157
  - **_\*\*connect_kwargs_** : Passed through to `snowflake.connector.connect()`, e.g. `autocommit=False`, `client_session_keep_alive=True`.
1158
1158
 
1159
+ > **CI authentication** uses key-pair auth exclusively (password-based CI auth was removed ahead of Snowflake's 2026-09-22 service-user password deprecation). In CI, reads `SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY` (raw PEM or base64-encoded PEM — base64 is required for GitLab CI/CD variables since multi-line values cannot be masked) and optionally `SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE` (or the `_PASSWORD` variant) for encrypted keys.
1160
+
1159
1161
  #### Returns
1160
1162
  - An open `snowflake.connector.SnowflakeConnection`. The caller owns its lifecycle.
1161
1163
 
@@ -1300,17 +1302,19 @@ gitlabds.snowflake_stage_load_copy_remove(
1300
1302
  #### Description
1301
1303
  Create a Snowflake Feature Store connection (Snowpark Session + native `FeatureStore`), resolving credentials the same way as `snowflake_engine_factory`.
1302
1304
 
1303
- `gitlabds.feature_store_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="feature_store", auth_method="password", creation_mode="CreationMode.FAIL_IF_NOT_EXIST", profiles_path=None):`
1305
+ `gitlabds.feature_store_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="feature_store", auth_method="keypair", creation_mode="CreationMode.FAIL_IF_NOT_EXIST", profiles_path=None):`
1304
1306
 
1305
1307
  #### Parameters:
1306
1308
  - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1307
1309
  - **_role_** : Snowflake role for CI credentials. The feature store's database and schema come from `SNOWFLAKE_FEATURE_STORE_DATABASE` / `SNOWFLAKE_FEATURE_STORE_SCHEMA` on the CI path, not from the role's usual database/schema.
1308
1310
  - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1309
1311
  - **_profile_target_** : dbt profile output name. Only used on the local path.
1310
- - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path.
1312
+ - **_auth_method_** : CI authentication method. `"keypair"` is the only supported value and the default, so passing it explicitly is the same as omitting it. Retained so callers that pass `auth_method="keypair"` keep working. `"password"` is accepted but ignored, with a logged warning naming the fix; any other value raises `ValueError`.
1311
1313
  - **_creation_mode_** : Either a `"CreationMode.X"` string or a `snowflake.ml.feature_store.CreationMode` member.
1312
1314
  - **_profiles_path_** : Override the dbt `profiles.yml` location. Only used on the local path.
1313
1315
 
1316
+ > **CI authentication** uses key-pair auth exclusively — see `snowflake_engine_factory` for the env var names and base64 encoding details.
1317
+
1314
1318
  #### Returns
1315
1319
  - A `snowflake.ml.feature_store.FeatureStore`. The underlying Snowpark session is reachable as `fs._session` and is not closed by this function.
1316
1320
 
@@ -1099,7 +1099,7 @@ feature_drift = results['feature_drift']
1099
1099
 
1100
1100
  #### Snowflake Connectors
1101
1101
 
1102
- Thin wrappers over the native Snowflake SDKs that resolve GitLab's two credential paths — local (dbt profile + SSO) and GitLab CI (env vars + password or key pair) — and then hand off to the native objects. Migrated from `gitlabdata.orchestration_utils`; see the migration notes at the end of this section.
1102
+ Thin wrappers over the native Snowflake SDKs that resolve GitLab's two credential paths — local (dbt profile + SSO) and GitLab CI (env vars + key-pair auth) — and then hand off to the native objects. Migrated from `gitlabdata.orchestration_utils`; see the migration notes at the end of this section.
1103
1103
 
1104
1104
  All of these require `pip install gitlabds[snowflake]`.
1105
1105
 
@@ -1108,7 +1108,7 @@ All of these require `pip install gitlabds[snowflake]`.
1108
1108
  #### Description
1109
1109
  Open a Snowflake connection, resolving credentials from your dbt profile locally or from GitLab CI environment variables in CI.
1110
1110
 
1111
- `gitlabds.snowflake_engine_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="dev", schema="", auth_method="password", profiles_path=None, **connect_kwargs):`
1111
+ `gitlabds.snowflake_engine_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="dev", schema="", auth_method="keypair", profiles_path=None, **connect_kwargs):`
1112
1112
 
1113
1113
  #### Parameters:
1114
1114
  - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
@@ -1116,10 +1116,12 @@ Open a Snowflake connection, resolving credentials from your dbt profile locally
1116
1116
  - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1117
1117
  - **_profile_target_** : dbt profile output name. Only used on the local path.
1118
1118
  - **_schema_** : Schema to connect to. Falls back to the dbt profile's schema locally; when it resolves empty the connection is opened with no default schema.
1119
- - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path. Key pair secrets are not yet provisioned in CI.
1119
+ - **_auth_method_** : CI authentication method. `"keypair"` is the only supported value and the default, so passing it explicitly is the same as omitting it. Retained so callers that pass `auth_method="keypair"` keep working. `"password"` is accepted but ignored, with a logged warning naming the fix; any other value raises `ValueError`.
1120
1120
  - **_profiles_path_** : Override the dbt `profiles.yml` location (defaults to `~/.dbt/profiles.yml`). Only used on the local path.
1121
1121
  - **_\*\*connect_kwargs_** : Passed through to `snowflake.connector.connect()`, e.g. `autocommit=False`, `client_session_keep_alive=True`.
1122
1122
 
1123
+ > **CI authentication** uses key-pair auth exclusively (password-based CI auth was removed ahead of Snowflake's 2026-09-22 service-user password deprecation). In CI, reads `SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY` (raw PEM or base64-encoded PEM — base64 is required for GitLab CI/CD variables since multi-line values cannot be masked) and optionally `SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE` (or the `_PASSWORD` variant) for encrypted keys.
1124
+
1123
1125
  #### Returns
1124
1126
  - An open `snowflake.connector.SnowflakeConnection`. The caller owns its lifecycle.
1125
1127
 
@@ -1264,17 +1266,19 @@ gitlabds.snowflake_stage_load_copy_remove(
1264
1266
  #### Description
1265
1267
  Create a Snowflake Feature Store connection (Snowpark Session + native `FeatureStore`), resolving credentials the same way as `snowflake_engine_factory`.
1266
1268
 
1267
- `gitlabds.feature_store_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="feature_store", auth_method="password", creation_mode="CreationMode.FAIL_IF_NOT_EXIST", profiles_path=None):`
1269
+ `gitlabds.feature_store_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="feature_store", auth_method="keypair", creation_mode="CreationMode.FAIL_IF_NOT_EXIST", profiles_path=None):`
1268
1270
 
1269
1271
  #### Parameters:
1270
1272
  - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1271
1273
  - **_role_** : Snowflake role for CI credentials. The feature store's database and schema come from `SNOWFLAKE_FEATURE_STORE_DATABASE` / `SNOWFLAKE_FEATURE_STORE_SCHEMA` on the CI path, not from the role's usual database/schema.
1272
1274
  - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1273
1275
  - **_profile_target_** : dbt profile output name. Only used on the local path.
1274
- - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path.
1276
+ - **_auth_method_** : CI authentication method. `"keypair"` is the only supported value and the default, so passing it explicitly is the same as omitting it. Retained so callers that pass `auth_method="keypair"` keep working. `"password"` is accepted but ignored, with a logged warning naming the fix; any other value raises `ValueError`.
1275
1277
  - **_creation_mode_** : Either a `"CreationMode.X"` string or a `snowflake.ml.feature_store.CreationMode` member.
1276
1278
  - **_profiles_path_** : Override the dbt `profiles.yml` location. Only used on the local path.
1277
1279
 
1280
+ > **CI authentication** uses key-pair auth exclusively — see `snowflake_engine_factory` for the env var names and base64 encoding details.
1281
+
1278
1282
  #### Returns
1279
1283
  - A `snowflake.ml.feature_store.FeatureStore`. The underlying Snowpark session is reachable as `fs._session` and is not closed by this function.
1280
1284
 
@@ -4,7 +4,7 @@ Snowflake connectors for Data Science pipelines.
4
4
  Thin wrappers around the native Snowflake SDKs (snowflake-connector-python,
5
5
  snowflake-snowpark-python, snowflake-ml-python) that resolve GitLab's two
6
6
  credential paths -- local (dbt profile + SSO) and GitLab CI (env vars +
7
- password or key pair) -- so callers don't need to reimplement that branch.
7
+ key-pair auth) -- so callers don't need to reimplement that branch.
8
8
  The native SDK objects (SnowflakeConnection, Session, FeatureStore) do the
9
9
  actual connecting; this module only decides which credentials to hand them.
10
10
 
@@ -12,6 +12,7 @@ Requires the optional Snowflake dependencies: pip install gitlabds[snowflake]
12
12
 
13
13
  """
14
14
 
15
+ import base64
15
16
  import logging
16
17
  import os
17
18
  from pathlib import Path
@@ -39,18 +40,15 @@ DOCKER_PROFILES_PATH = "/usr/local/snowflake_profile/profiles.yml"
39
40
  _CI_ROLE_VARS: Dict[str, Dict[str, str]] = {
40
41
  "DATA_SCIENCE_LOADER": {
41
42
  "user": "SNOWFLAKE_DATA_SCIENCE_LOAD_USER",
42
- "password": "SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD",
43
43
  "private_key": "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY",
44
44
  "private_key_password": "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSWORD",
45
+ "private_key_passphrase": "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE",
45
46
  "account": "SNOWFLAKE_ACCOUNT",
46
47
  "database": "SNOWFLAKE_PROD_DATABASE",
47
48
  "warehouse": "SNOWFLAKE_LOAD_WAREHOUSE",
48
49
  },
49
50
  }
50
51
 
51
- # The private_key env vars are not provisioned in CI yet, so auth_method="keypair"
52
- # will report them missing. Password auth remains the default.
53
-
54
52
 
55
53
  def _dbt_profile_config(
56
54
  profile_target: str, profiles_path: Optional[Union[str, Path]] = None
@@ -78,9 +76,53 @@ def _dbt_profile_config(
78
76
 
79
77
 
80
78
  def _private_key_der(private_key_pem: str, password: Optional[str] = None) -> bytes:
81
- """Convert a PEM-encoded private key to the DER/PKCS8 bytes Snowflake expects."""
79
+ """
80
+ Convert a PEM private key to the DER/PKCS8 bytes Snowflake expects.
81
+
82
+ Accepts either a raw PEM string (starts with ``-----BEGIN``) or a
83
+ base64-encoded PEM string. Base64 encoding is required when the key is
84
+ stored in a GitLab CI/CD variable, because GitLab cannot mask multi-line
85
+ values; callers should base64-encode the PEM before storing it and this
86
+ function will decode it transparently.
87
+
88
+ Parameters
89
+ ----------
90
+ private_key_pem : str
91
+ Raw PEM string or base64-encoded PEM string.
92
+ password : str, optional
93
+ Passphrase protecting the private key, if any.
94
+
95
+ Returns
96
+ -------
97
+ bytes
98
+ Unencrypted DER-encoded private key bytes suitable for
99
+ ``snowflake.connector.connect(private_key=...)``.
100
+
101
+ Raises
102
+ ------
103
+ ValueError
104
+ When the input is neither a valid PEM string nor valid base64-encoded
105
+ PEM.
106
+ """
107
+ pem_str = private_key_pem.strip()
108
+ if pem_str.startswith("-----BEGIN"):
109
+ pem_bytes = pem_str.encode("utf-8")
110
+ else:
111
+ try:
112
+ pem_bytes = base64.b64decode(pem_str)
113
+ except Exception:
114
+ raise ValueError(
115
+ "private_key is neither a valid PEM string (expected to start with "
116
+ "'-----BEGIN') nor valid base64-encoded PEM. "
117
+ "Encode the PEM key with base64 before storing it in a CI/CD variable."
118
+ )
119
+ if not pem_bytes.lstrip().startswith(b"-----BEGIN"):
120
+ raise ValueError(
121
+ "private_key decoded from base64 is not a valid PEM string. "
122
+ "Encode the PEM key with base64 before storing it in a CI/CD variable."
123
+ )
82
124
  key = load_pem_private_key(
83
- private_key_pem.encode("utf-8"),
125
+ pem_bytes,
84
126
  password=password.encode("utf-8") if password else None,
85
127
  backend=default_backend(),
86
128
  )
@@ -103,7 +145,8 @@ def _resolve_connect_params(
103
145
  """
104
146
  Build connect() / Session.builder.configs() kwargs for either credential path.
105
147
 
106
- "local" reads the dbt profile and uses SSO; "ci" reads env vars for `role`.
148
+ "local" reads the dbt profile and uses SSO; "ci" reads env vars for `role`
149
+ and uses key-pair authentication -- the only CI method still supported.
107
150
  The schema key is omitted when it resolves empty, so the connection gets no
108
151
  default schema rather than one named "".
109
152
  """
@@ -131,14 +174,25 @@ def _resolve_connect_params(
131
174
  raise ValueError(f"run_target must be 'local' or 'ci', got {run_target!r}")
132
175
  if role not in _CI_ROLE_VARS:
133
176
  raise ValueError(f"Unknown role {role!r} for CI auth. Known roles: {sorted(_CI_ROLE_VARS)}")
134
- if auth_method not in ("password", "keypair"):
135
- raise ValueError(f"auth_method must be 'password' or 'keypair', got {auth_method!r}")
177
+ if auth_method == "password":
178
+ # Don't fail the run: password auth stops working at Snowflake's
179
+ # 2026-09-22 deadline anyway, so key-pair is the only thing that could
180
+ # succeed. Warn loudly and use it.
181
+ logger.warning(
182
+ "auth_method='password' is no longer supported and was ignored; connecting "
183
+ "with key-pair auth instead. Password-based CI authentication was removed "
184
+ "ahead of Snowflake's 2026-09-22 service-user password deprecation. To fix: "
185
+ "change auth_method='password' to auth_method='keypair' (or drop the "
186
+ "parameter -- 'keypair' is the default) and set %s in your CI/CD variables.",
187
+ _CI_ROLE_VARS[role]["private_key"],
188
+ )
189
+ elif auth_method != "keypair":
190
+ raise ValueError(f"auth_method must be 'keypair', got {auth_method!r}")
136
191
 
137
192
  var_names = _CI_ROLE_VARS[role]
138
- credential = "private_key" if auth_method == "keypair" else "password"
139
193
  missing = [
140
194
  var_names[key]
141
- for key in ("user", "account", "database", "warehouse", credential)
195
+ for key in ("user", "account", "database", "warehouse", "private_key")
142
196
  if not env.get(var_names[key])
143
197
  ]
144
198
  if missing:
@@ -155,12 +209,14 @@ def _resolve_connect_params(
155
209
  )
156
210
  if schema:
157
211
  params["schema"] = schema
158
- if auth_method == "keypair":
159
- params["private_key"] = _private_key_der(
160
- env[var_names["private_key"]], env.get(var_names["private_key_password"])
161
- )
162
- else:
163
- params["password"] = env[var_names["password"]]
212
+ # Accept either _PASSWORD or _PASSPHRASE; the group CI/CD variable uses
213
+ # the _PASSPHRASE suffix, which is also the more accurate term.
214
+ key_passphrase = env.get(var_names["private_key_password"]) or env.get(
215
+ var_names["private_key_passphrase"]
216
+ )
217
+ params["private_key"] = _private_key_der(
218
+ env[var_names["private_key"]], key_passphrase
219
+ )
164
220
  return params
165
221
 
166
222
 
@@ -170,7 +226,7 @@ def snowflake_engine_factory(
170
226
  run_target: Optional[str] = None,
171
227
  profile_target: str = "dev",
172
228
  schema: str = "",
173
- auth_method: str = "password",
229
+ auth_method: str = "keypair",
174
230
  profiles_path: Optional[Union[str, Path]] = None,
175
231
  **connect_kwargs: Any,
176
232
  ) -> SnowflakeConnection:
@@ -192,8 +248,13 @@ def snowflake_engine_factory(
192
248
  dbt profile output name. Only used on the local path.
193
249
  schema : str, default=""
194
250
  Schema to connect to. Falls back to the dbt profile's schema locally.
195
- auth_method : str, default="password"
196
- "password" or "keypair". Only used on the CI path.
251
+ auth_method : str, default="keypair"
252
+ Authentication method for the CI path. Only "keypair" is supported, so
253
+ passing it explicitly is equivalent to omitting it. Retained for
254
+ backward compatibility with callers that pass auth_method="keypair".
255
+ "password" is accepted but ignored with a warning -- password-based CI
256
+ auth was removed ahead of Snowflake's 2026-09-22 service-user password
257
+ deprecation. Any other value raises ValueError.
197
258
  profiles_path : str or Path, optional
198
259
  Override the dbt profiles.yml location. Only used on the local path.
199
260
  **connect_kwargs
@@ -366,7 +427,7 @@ def feature_store_factory(
366
427
  role: str = "DATA_SCIENCE_LOADER",
367
428
  run_target: Optional[str] = None,
368
429
  profile_target: str = "feature_store",
369
- auth_method: str = "password",
430
+ auth_method: str = "keypair",
370
431
  creation_mode: str = "CreationMode.FAIL_IF_NOT_EXIST",
371
432
  profiles_path: Optional[Union[str, Path]] = None,
372
433
  ) -> "FeatureStore": # noqa: F821
@@ -387,8 +448,13 @@ def feature_store_factory(
387
448
  not given.
388
449
  profile_target : str, default="feature_store"
389
450
  dbt profile output name. Only used on the local path.
390
- auth_method : str, default="password"
391
- "password" or "keypair". Only used on the CI path.
451
+ auth_method : str, default="keypair"
452
+ Authentication method for the CI path. Only "keypair" is supported, so
453
+ passing it explicitly is equivalent to omitting it. Retained for
454
+ backward compatibility with callers that pass auth_method="keypair".
455
+ "password" is accepted but ignored with a warning -- password-based CI
456
+ auth was removed ahead of Snowflake's 2026-09-22 service-user password
457
+ deprecation. Any other value raises ValueError.
392
458
  creation_mode : str or CreationMode, default="CreationMode.FAIL_IF_NOT_EXIST"
393
459
  Either a "CreationMode.X" string or a CreationMode member.
394
460
  profiles_path : str or Path, optional
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gitlabds
3
- Version: 2.3.1
3
+ Version: 2.4.0
4
4
  Summary: GitLab Data Science Tools
5
5
  Author-email: Kevin Dietz <kdietz@gitlab.com>
6
6
  License-Expression: MIT
@@ -1135,7 +1135,7 @@ feature_drift = results['feature_drift']
1135
1135
 
1136
1136
  #### Snowflake Connectors
1137
1137
 
1138
- Thin wrappers over the native Snowflake SDKs that resolve GitLab's two credential paths — local (dbt profile + SSO) and GitLab CI (env vars + password or key pair) — and then hand off to the native objects. Migrated from `gitlabdata.orchestration_utils`; see the migration notes at the end of this section.
1138
+ Thin wrappers over the native Snowflake SDKs that resolve GitLab's two credential paths — local (dbt profile + SSO) and GitLab CI (env vars + key-pair auth) — and then hand off to the native objects. Migrated from `gitlabdata.orchestration_utils`; see the migration notes at the end of this section.
1139
1139
 
1140
1140
  All of these require `pip install gitlabds[snowflake]`.
1141
1141
 
@@ -1144,7 +1144,7 @@ All of these require `pip install gitlabds[snowflake]`.
1144
1144
  #### Description
1145
1145
  Open a Snowflake connection, resolving credentials from your dbt profile locally or from GitLab CI environment variables in CI.
1146
1146
 
1147
- `gitlabds.snowflake_engine_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="dev", schema="", auth_method="password", profiles_path=None, **connect_kwargs):`
1147
+ `gitlabds.snowflake_engine_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="dev", schema="", auth_method="keypair", profiles_path=None, **connect_kwargs):`
1148
1148
 
1149
1149
  #### Parameters:
1150
1150
  - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
@@ -1152,10 +1152,12 @@ Open a Snowflake connection, resolving credentials from your dbt profile locally
1152
1152
  - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1153
1153
  - **_profile_target_** : dbt profile output name. Only used on the local path.
1154
1154
  - **_schema_** : Schema to connect to. Falls back to the dbt profile's schema locally; when it resolves empty the connection is opened with no default schema.
1155
- - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path. Key pair secrets are not yet provisioned in CI.
1155
+ - **_auth_method_** : CI authentication method. `"keypair"` is the only supported value and the default, so passing it explicitly is the same as omitting it. Retained so callers that pass `auth_method="keypair"` keep working. `"password"` is accepted but ignored, with a logged warning naming the fix; any other value raises `ValueError`.
1156
1156
  - **_profiles_path_** : Override the dbt `profiles.yml` location (defaults to `~/.dbt/profiles.yml`). Only used on the local path.
1157
1157
  - **_\*\*connect_kwargs_** : Passed through to `snowflake.connector.connect()`, e.g. `autocommit=False`, `client_session_keep_alive=True`.
1158
1158
 
1159
+ > **CI authentication** uses key-pair auth exclusively (password-based CI auth was removed ahead of Snowflake's 2026-09-22 service-user password deprecation). In CI, reads `SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY` (raw PEM or base64-encoded PEM — base64 is required for GitLab CI/CD variables since multi-line values cannot be masked) and optionally `SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE` (or the `_PASSWORD` variant) for encrypted keys.
1160
+
1159
1161
  #### Returns
1160
1162
  - An open `snowflake.connector.SnowflakeConnection`. The caller owns its lifecycle.
1161
1163
 
@@ -1300,17 +1302,19 @@ gitlabds.snowflake_stage_load_copy_remove(
1300
1302
  #### Description
1301
1303
  Create a Snowflake Feature Store connection (Snowpark Session + native `FeatureStore`), resolving credentials the same way as `snowflake_engine_factory`.
1302
1304
 
1303
- `gitlabds.feature_store_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="feature_store", auth_method="password", creation_mode="CreationMode.FAIL_IF_NOT_EXIST", profiles_path=None):`
1305
+ `gitlabds.feature_store_factory(args=None, role="DATA_SCIENCE_LOADER", run_target=None, profile_target="feature_store", auth_method="keypair", creation_mode="CreationMode.FAIL_IF_NOT_EXIST", profiles_path=None):`
1304
1306
 
1305
1307
  #### Parameters:
1306
1308
  - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1307
1309
  - **_role_** : Snowflake role for CI credentials. The feature store's database and schema come from `SNOWFLAKE_FEATURE_STORE_DATABASE` / `SNOWFLAKE_FEATURE_STORE_SCHEMA` on the CI path, not from the role's usual database/schema.
1308
1310
  - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1309
1311
  - **_profile_target_** : dbt profile output name. Only used on the local path.
1310
- - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path.
1312
+ - **_auth_method_** : CI authentication method. `"keypair"` is the only supported value and the default, so passing it explicitly is the same as omitting it. Retained so callers that pass `auth_method="keypair"` keep working. `"password"` is accepted but ignored, with a logged warning naming the fix; any other value raises `ValueError`.
1311
1313
  - **_creation_mode_** : Either a `"CreationMode.X"` string or a `snowflake.ml.feature_store.CreationMode` member.
1312
1314
  - **_profiles_path_** : Override the dbt `profiles.yml` location. Only used on the local path.
1313
1315
 
1316
+ > **CI authentication** uses key-pair auth exclusively — see `snowflake_engine_factory` for the env var names and base64 encoding details.
1317
+
1314
1318
  #### Returns
1315
1319
  - A `snowflake.ml.feature_store.FeatureStore`. The underlying Snowpark session is reachable as `fs._session` and is not closed by this function.
1316
1320
 
@@ -65,14 +65,21 @@ class FakeConnection:
65
65
  return self._cursor
66
66
 
67
67
 
68
- CI_ENV = {
68
+ # CI_ENV_BASE intentionally omits SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY so that
69
+ # individual tests can inject a freshly-generated key via rsa_pem. Tests that
70
+ # need a fully-valid CI env build one with make_ci_env(rsa_pem).
71
+ CI_ENV_BASE = {
69
72
  "SNOWFLAKE_DATA_SCIENCE_LOAD_USER": "ds_user",
70
- "SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD": "ds_password",
71
73
  "SNOWFLAKE_ACCOUNT": "gitlab_account",
72
74
  "SNOWFLAKE_PROD_DATABASE": "PROD",
73
75
  "SNOWFLAKE_LOAD_WAREHOUSE": "LOADING",
74
76
  }
75
77
 
78
+
79
+ def make_ci_env(private_key_pem: str) -> dict:
80
+ """Return a complete CI env dict with the given PEM key."""
81
+ return dict(CI_ENV_BASE, SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=private_key_pem)
82
+
76
83
  PROFILE = {
77
84
  "user": "kdietz@gitlab.com",
78
85
  "account": "gitlab_account",
@@ -131,9 +138,9 @@ class TestCiRoleVars:
131
138
  def test_data_science_loader_env_var_names(self):
132
139
  assert sc._CI_ROLE_VARS["DATA_SCIENCE_LOADER"] == {
133
140
  "user": "SNOWFLAKE_DATA_SCIENCE_LOAD_USER",
134
- "password": "SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD",
135
141
  "private_key": "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY",
136
142
  "private_key_password": "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSWORD",
143
+ "private_key_passphrase": "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE",
137
144
  "account": "SNOWFLAKE_ACCOUNT",
138
145
  "database": "SNOWFLAKE_PROD_DATABASE",
139
146
  "warehouse": "SNOWFLAKE_LOAD_WAREHOUSE",
@@ -184,6 +191,7 @@ class TestDbtProfileConfig:
184
191
 
185
192
  class TestPrivateKeyDer:
186
193
  def test_unencrypted_pem_round_trips_to_der(self, rsa_pem):
194
+ """Raw PEM input (existing behavior) must keep working."""
187
195
  der = sc._private_key_der(rsa_pem)
188
196
  assert isinstance(der, bytes)
189
197
  assert serialization.load_der_private_key(der, password=None).key_size == 2048
@@ -197,6 +205,43 @@ class TestPrivateKeyDer:
197
205
  with pytest.raises(ValueError):
198
206
  sc._private_key_der(encrypted_pem(b"hunter2"), password="wrong")
199
207
 
208
+ def test_base64_encoded_pem_is_accepted(self, rsa_pem):
209
+ """GitLab CI/CD variables cannot mask multi-line PEM; base64 is the workaround."""
210
+ import base64
211
+
212
+ b64_pem = base64.b64encode(rsa_pem.encode("utf-8")).decode("utf-8")
213
+ assert not b64_pem.startswith("-----BEGIN")
214
+ der = sc._private_key_der(b64_pem)
215
+ assert isinstance(der, bytes)
216
+ assert serialization.load_der_private_key(der, password=None).key_size == 2048
217
+
218
+ def test_base64_encoded_encrypted_pem_is_accepted(self):
219
+ """Base64-encoded encrypted PEM with a passphrase must also work."""
220
+ import base64
221
+
222
+ pem = encrypted_pem(b"hunter2")
223
+ b64_pem = base64.b64encode(pem.encode("utf-8")).decode("utf-8")
224
+ der = sc._private_key_der(b64_pem, password="hunter2")
225
+ assert serialization.load_der_private_key(der, password=None) is not None
226
+
227
+ def test_invalid_input_raises_clear_value_error(self):
228
+ """Garbage input must raise ValueError with a helpful message, not an opaque crypto error."""
229
+ with pytest.raises(ValueError, match="neither a valid PEM string.*nor valid base64-encoded PEM"):
230
+ sc._private_key_der("this-is-not-a-key")
231
+
232
+ def test_base64_of_non_pem_raises_clear_value_error(self):
233
+ """base64('hello world') decodes fine but is not a PEM -- must raise clearly."""
234
+ import base64
235
+
236
+ b64_garbage = base64.b64encode(b"hello world").decode("utf-8")
237
+ with pytest.raises(ValueError, match="not a valid PEM string"):
238
+ sc._private_key_der(b64_garbage)
239
+
240
+ def test_raw_pem_with_leading_whitespace_is_accepted(self, rsa_pem):
241
+ """strip() must handle leading/trailing whitespace that CI env vars sometimes add."""
242
+ der = sc._private_key_der("\n" + rsa_pem + "\n")
243
+ assert serialization.load_der_private_key(der, password=None).key_size == 2048
244
+
200
245
 
201
246
  # ---------------------------------------------------------------------------
202
247
  # _resolve_connect_params -- local path
@@ -206,7 +251,7 @@ class TestPrivateKeyDer:
206
251
  class TestResolveConnectParamsLocal:
207
252
  def _resolve(self, profiles_file, schema="", profile_target="dev"):
208
253
  return sc._resolve_connect_params(
209
- "local", "DATA_SCIENCE_LOADER", profile_target, schema, "password", {}, profiles_file
254
+ "local", "DATA_SCIENCE_LOADER", profile_target, schema, "keypair", {}, profiles_file
210
255
  )
211
256
 
212
257
  def test_uses_sso(self, profiles_file):
@@ -252,94 +297,140 @@ class TestResolveConnectParamsLocal:
252
297
 
253
298
 
254
299
  class TestResolveConnectParamsCi:
255
- def _resolve(self, env=None, role="DATA_SCIENCE_LOADER", schema="", auth_method="password"):
300
+ def _resolve(self, rsa_pem, env=None, role="DATA_SCIENCE_LOADER", schema="", auth_method="keypair"):
256
301
  return sc._resolve_connect_params(
257
- "ci", role, "dev", schema, auth_method, env if env is not None else dict(CI_ENV)
302
+ "ci", role, "dev", schema, auth_method, env if env is not None else make_ci_env(rsa_pem)
258
303
  )
259
304
 
260
- def test_reads_the_roles_env_vars(self):
261
- params = self._resolve()
305
+ def test_reads_the_roles_env_vars(self, rsa_pem):
306
+ params = self._resolve(rsa_pem)
262
307
  assert params["user"] == "ds_user"
263
- assert params["password"] == "ds_password"
308
+ assert isinstance(params["private_key"], bytes)
309
+ assert "password" not in params
264
310
  assert params["account"] == "gitlab_account"
265
311
  assert params["database"] == "PROD"
266
312
  assert params["warehouse"] == "LOADING"
267
313
  assert params["role"] == "DATA_SCIENCE_LOADER"
268
314
 
269
- def test_never_uses_sso(self):
270
- assert "authenticator" not in self._resolve()
315
+ def test_never_uses_sso(self, rsa_pem):
316
+ assert "authenticator" not in self._resolve(rsa_pem)
271
317
 
272
- def test_schema_is_included_when_given(self):
273
- assert self._resolve(schema="data_science")["schema"] == "data_science"
318
+ def test_schema_is_included_when_given(self, rsa_pem):
319
+ assert self._resolve(rsa_pem, schema="data_science")["schema"] == "data_science"
274
320
 
275
- def test_schema_key_is_omitted_when_empty(self):
276
- assert "schema" not in self._resolve()
321
+ def test_schema_key_is_omitted_when_empty(self, rsa_pem):
322
+ assert "schema" not in self._resolve(rsa_pem)
277
323
 
278
324
  def test_invalid_run_target_raises(self):
279
325
  with pytest.raises(ValueError, match="run_target must be 'local' or 'ci'"):
280
- sc._resolve_connect_params("prod", "DATA_SCIENCE_LOADER", "dev", "", "password", {})
326
+ sc._resolve_connect_params("prod", "DATA_SCIENCE_LOADER", "dev", "", "keypair", {})
281
327
 
282
- def test_unknown_role_lists_the_known_ones(self):
328
+ def test_password_auth_method_warns_and_falls_back_to_keypair(self, rsa_pem, caplog):
329
+ """
330
+ Password auth is gone, but failing the run helps nobody -- key-pair is the
331
+ only method that can succeed. Warn with the fix and connect anyway.
332
+ """
333
+ with caplog.at_level(logging.WARNING, logger="gitlabds.snowflake_connectors"):
334
+ params = self._resolve(rsa_pem, auth_method="password")
335
+ assert isinstance(params["private_key"], bytes)
336
+ assert "password" not in params
337
+ assert "auth_method='password' is no longer supported" in caplog.text
338
+
339
+ def test_password_warning_names_the_fix_and_the_env_var(self, rsa_pem, caplog):
340
+ with caplog.at_level(logging.WARNING, logger="gitlabds.snowflake_connectors"):
341
+ self._resolve(rsa_pem, auth_method="password")
342
+ assert "auth_method='keypair'" in caplog.text
343
+ assert "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY" in caplog.text
344
+
345
+ def test_keypair_is_silent(self, rsa_pem, caplog):
346
+ """The supported call style must not nag."""
347
+ with caplog.at_level(logging.WARNING, logger="gitlabds.snowflake_connectors"):
348
+ self._resolve(rsa_pem, auth_method="keypair")
349
+ assert caplog.text == ""
350
+
351
+ def test_unknown_auth_method_raises(self, rsa_pem):
352
+ """Only 'keypair' is accepted; anything else must fail."""
353
+ with pytest.raises(ValueError, match="auth_method must be 'keypair'"):
354
+ self._resolve(rsa_pem, auth_method="oauth")
355
+
356
+ def test_unknown_role_lists_the_known_ones(self, rsa_pem):
283
357
  with pytest.raises(ValueError, match="Unknown role 'NOPE' for CI auth"):
284
- self._resolve(role="NOPE")
285
-
286
- def test_unknown_auth_method_raises(self):
287
- with pytest.raises(ValueError, match="auth_method must be 'password' or 'keypair'"):
288
- self._resolve(auth_method="oauth")
358
+ self._resolve(rsa_pem, role="NOPE")
289
359
 
290
360
  @pytest.mark.parametrize(
291
361
  "missing_var",
292
362
  [
293
363
  "SNOWFLAKE_DATA_SCIENCE_LOAD_USER",
294
- "SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD",
364
+ "SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY",
295
365
  "SNOWFLAKE_ACCOUNT",
296
366
  "SNOWFLAKE_PROD_DATABASE",
297
367
  "SNOWFLAKE_LOAD_WAREHOUSE",
298
368
  ],
299
369
  )
300
- def test_each_missing_credential_names_itself(self, missing_var):
301
- env = {k: v for k, v in CI_ENV.items() if k != missing_var}
370
+ def test_each_missing_credential_names_itself(self, rsa_pem, missing_var):
371
+ env = {k: v for k, v in make_ci_env(rsa_pem).items() if k != missing_var}
302
372
  with pytest.raises(ValueError, match=missing_var):
303
- self._resolve(env=env)
373
+ self._resolve(rsa_pem, env=env)
304
374
 
305
- def test_empty_credential_counts_as_missing(self):
375
+ def test_empty_private_key_counts_as_missing(self, rsa_pem):
306
376
  """An env var set to "" in CI is a misconfiguration, not a valid credential."""
307
- env = dict(CI_ENV, SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD="")
308
- with pytest.raises(ValueError, match="SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD"):
309
- self._resolve(env=env)
377
+ env = dict(make_ci_env(rsa_pem), SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY="")
378
+ with pytest.raises(ValueError, match="SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY"):
379
+ self._resolve(rsa_pem, env=env)
310
380
 
311
- def test_all_missing_variables_are_reported_at_once(self):
381
+ def test_all_missing_variables_are_reported_at_once(self, rsa_pem):
312
382
  with pytest.raises(ValueError) as exc:
313
- self._resolve(env={})
383
+ self._resolve(rsa_pem, env={})
314
384
  message = str(exc.value)
315
385
  assert "SNOWFLAKE_DATA_SCIENCE_LOAD_USER" in message
316
386
  assert "SNOWFLAKE_ACCOUNT" in message
317
387
  assert "SNOWFLAKE_LOAD_WAREHOUSE" in message
318
388
 
319
- def test_keypair_sends_der_bytes_and_no_password(self, rsa_pem):
320
- env = dict(CI_ENV, SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=rsa_pem)
321
- params = self._resolve(env=env, auth_method="keypair")
389
+ def test_sends_der_bytes_and_no_password(self, rsa_pem):
390
+ params = self._resolve(rsa_pem)
322
391
  assert isinstance(params["private_key"], bytes)
323
392
  assert "password" not in params
324
393
 
325
- def test_keypair_uses_the_key_password_when_present(self):
394
+ def test_uses_the_key_password_when_present(self):
326
395
  env = dict(
327
- CI_ENV,
396
+ CI_ENV_BASE,
328
397
  SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=encrypted_pem(b"s3cret"),
329
398
  SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSWORD="s3cret",
330
399
  )
331
- assert isinstance(self._resolve(env=env, auth_method="keypair")["private_key"], bytes)
400
+ assert isinstance(self._resolve(None, env=env)["private_key"], bytes)
332
401
 
333
- def test_keypair_without_the_secret_names_the_variable(self):
334
- """Key pair vars are not provisioned in CI yet -- fail with a named var."""
335
- with pytest.raises(ValueError, match="SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY"):
336
- self._resolve(auth_method="keypair")
402
+ def test_passphrase_env_var_is_accepted_as_fallback(self):
403
+ """SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE is the actual group CI/CD
404
+ variable name; gitlabds must accept it when the _PASSWORD variant is absent."""
405
+ env = dict(
406
+ CI_ENV_BASE,
407
+ SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=encrypted_pem(b"s3cret"),
408
+ SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE="s3cret",
409
+ )
410
+ params = self._resolve(None, env=env)
411
+ assert isinstance(params["private_key"], bytes)
412
+
413
+ def test_password_takes_precedence_over_passphrase(self):
414
+ """When both _PASSWORD and _PASSPHRASE are set, _PASSWORD wins."""
415
+ env = dict(
416
+ CI_ENV_BASE,
417
+ SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=encrypted_pem(b"s3cret"),
418
+ SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSWORD="s3cret",
419
+ SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY_PASSPHRASE="wrong",
420
+ )
421
+ # If _PASSWORD is used (correct), decryption succeeds; if _PASSPHRASE is
422
+ # used instead, it would fail with a ValueError.
423
+ params = self._resolve(None, env=env)
424
+ assert isinstance(params["private_key"], bytes)
337
425
 
338
- def test_keypair_does_not_require_a_password_variable(self, rsa_pem):
339
- env = dict(CI_ENV, SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=rsa_pem)
340
- del env["SNOWFLAKE_DATA_SCIENCE_LOAD_PASSWORD"]
341
- assert "private_key" in self._resolve(env=env, auth_method="keypair")
426
+ def test_base64_pem_via_env_var(self, rsa_pem):
427
+ """End-to-end: a base64-encoded PEM stored in the CI env var resolves correctly."""
428
+ import base64
342
429
 
430
+ b64_pem = base64.b64encode(rsa_pem.encode("utf-8")).decode("utf-8")
431
+ env = dict(CI_ENV_BASE, SNOWFLAKE_DATA_SCIENCE_LOAD_PRIVATE_KEY=b64_pem)
432
+ params = self._resolve(rsa_pem, env=env)
433
+ assert isinstance(params["private_key"], bytes)
343
434
 
344
435
  # ---------------------------------------------------------------------------
345
436
  # snowflake_engine_factory
@@ -363,26 +454,32 @@ class TestSnowflakeEngineFactory:
363
454
  sc.snowflake_engine_factory(profiles_path=profiles_file)
364
455
  assert captured_connect["authenticator"] == "externalbrowser"
365
456
 
366
- def test_ci_path_when_gitlab_ci_is_set(self, captured_connect, monkeypatch):
457
+ def test_ci_path_when_gitlab_ci_is_set(self, captured_connect, monkeypatch, rsa_pem):
367
458
  monkeypatch.setenv("GITLAB_CI", "true")
368
- sc.snowflake_engine_factory(args=CI_ENV)
459
+ sc.snowflake_engine_factory(args=make_ci_env(rsa_pem))
369
460
  assert captured_connect["user"] == "ds_user"
370
461
  assert "authenticator" not in captured_connect
371
462
 
372
- def test_ci_is_detected_even_when_args_omits_gitlab_ci(self, captured_connect, monkeypatch):
463
+ def test_ci_is_detected_even_when_args_omits_gitlab_ci(
464
+ self, captured_connect, monkeypatch, rsa_pem
465
+ ):
373
466
  """
374
467
  Regression: inference must read os.environ, not `args`. A caller passing a
375
468
  curated credential dict in CI used to be routed to the local SSO path and
376
469
  hang on browser auth.
377
470
  """
378
471
  monkeypatch.setenv("GITLAB_CI", "true")
379
- sc.snowflake_engine_factory(args=dict(CI_ENV)) # no GITLAB_CI key inside
472
+ sc.snowflake_engine_factory(args=make_ci_env(rsa_pem)) # no GITLAB_CI key inside
380
473
  assert "authenticator" not in captured_connect
381
- assert captured_connect["password"] == "ds_password"
474
+ assert isinstance(captured_connect["private_key"], bytes)
382
475
 
383
- def test_gitlab_ci_inside_args_does_not_force_the_ci_path(self, captured_connect, profiles_file):
476
+ def test_gitlab_ci_inside_args_does_not_force_the_ci_path(
477
+ self, captured_connect, profiles_file, rsa_pem
478
+ ):
384
479
  """The converse: a stray GITLAB_CI in the passed dict is not a CI signal."""
385
- sc.snowflake_engine_factory(args=dict(CI_ENV, GITLAB_CI="true"), profiles_path=profiles_file)
480
+ sc.snowflake_engine_factory(
481
+ args=dict(make_ci_env(rsa_pem), GITLAB_CI="true"), profiles_path=profiles_file
482
+ )
386
483
  assert captured_connect["authenticator"] == "externalbrowser"
387
484
 
388
485
  def test_explicit_run_target_wins_over_inference(
@@ -392,9 +489,9 @@ class TestSnowflakeEngineFactory:
392
489
  sc.snowflake_engine_factory(run_target="local", profiles_path=profiles_file)
393
490
  assert captured_connect["authenticator"] == "externalbrowser"
394
491
 
395
- def test_defaults_args_to_os_environ(self, captured_connect, monkeypatch):
492
+ def test_defaults_args_to_os_environ(self, captured_connect, monkeypatch, rsa_pem):
396
493
  monkeypatch.setenv("GITLAB_CI", "true")
397
- for name, value in CI_ENV.items():
494
+ for name, value in make_ci_env(rsa_pem).items():
398
495
  monkeypatch.setenv(name, value)
399
496
  sc.snowflake_engine_factory() # no args passed at all
400
497
  assert captured_connect["user"] == "ds_user"
@@ -423,6 +520,50 @@ class TestSnowflakeEngineFactory:
423
520
  with pytest.raises(ValueError, match="run_target must be"):
424
521
  sc.snowflake_engine_factory(run_target="prod")
425
522
 
523
+ def test_explicit_keypair_auth_method_is_accepted(
524
+ self, captured_connect, monkeypatch, rsa_pem
525
+ ):
526
+ """
527
+ Downstream callers migrated to key-pair auth pass auth_method="keypair"
528
+ explicitly. Dropping the parameter would TypeError at their call site, so
529
+ it stays in the signature until 3.0.
530
+ """
531
+ monkeypatch.setenv("GITLAB_CI", "true")
532
+ sc.snowflake_engine_factory(args=make_ci_env(rsa_pem), auth_method="keypair")
533
+ assert isinstance(captured_connect["private_key"], bytes)
534
+
535
+ def test_auth_method_is_not_forwarded_to_connect(self, captured_connect, monkeypatch, rsa_pem):
536
+ """It is consumed by the signature, not swallowed into **connect_kwargs."""
537
+ monkeypatch.setenv("GITLAB_CI", "true")
538
+ sc.snowflake_engine_factory(args=make_ci_env(rsa_pem), auth_method="keypair")
539
+ assert "auth_method" not in captured_connect
540
+
541
+ def test_explicit_keypair_is_identical_to_omitting_it(self, monkeypatch, rsa_pem):
542
+ """
543
+ The contract downstream callers rely on: auth_method="keypair" and no
544
+ auth_method at all must produce the same connection.
545
+ """
546
+ monkeypatch.setenv("GITLAB_CI", "true")
547
+ env = make_ci_env(rsa_pem)
548
+ captured = []
549
+
550
+ def fake_connect(**kwargs):
551
+ captured.append(kwargs)
552
+ return FakeConnection()
553
+
554
+ monkeypatch.setattr(snowflake.connector, "connect", fake_connect)
555
+ sc.snowflake_engine_factory(args=env)
556
+ sc.snowflake_engine_factory(args=env, auth_method="keypair")
557
+ assert captured[0] == captured[1]
558
+
559
+ def test_password_auth_method_warns_and_still_connects(self, captured_connect, monkeypatch, rsa_pem, caplog):
560
+ monkeypatch.setenv("GITLAB_CI", "true")
561
+ with caplog.at_level(logging.WARNING, logger="gitlabds.snowflake_connectors"):
562
+ sc.snowflake_engine_factory(args=make_ci_env(rsa_pem), auth_method="password")
563
+ assert isinstance(captured_connect["private_key"], bytes)
564
+ assert "password" not in captured_connect
565
+ assert "auth_method='password' is no longer supported" in caplog.text
566
+
426
567
 
427
568
  # ---------------------------------------------------------------------------
428
569
  # ds_query_dataframe
@@ -739,11 +880,11 @@ class TestFeatureStoreFactory:
739
880
  assert captured_feature_store["default_warehouse"] == PROFILE["warehouse"]
740
881
 
741
882
  def test_ci_path_uses_the_dedicated_feature_store_env_vars(
742
- self, captured_feature_store, monkeypatch
883
+ self, captured_feature_store, monkeypatch, rsa_pem
743
884
  ):
744
885
  monkeypatch.setenv("GITLAB_CI", "true")
745
886
  env = dict(
746
- CI_ENV,
887
+ make_ci_env(rsa_pem),
747
888
  SNOWFLAKE_FEATURE_STORE_DATABASE="FS_DB",
748
889
  SNOWFLAKE_FEATURE_STORE_SCHEMA="FS_SCHEMA",
749
890
  )
@@ -753,10 +894,42 @@ class TestFeatureStoreFactory:
753
894
  # The role's own PROD database must not leak through.
754
895
  assert captured_feature_store["session_params"]["database"] == "FS_DB"
755
896
 
756
- def test_ci_path_without_the_feature_store_vars_names_them(self, monkeypatch):
897
+ def test_explicit_keypair_auth_method_is_accepted(
898
+ self, captured_feature_store, monkeypatch, rsa_pem
899
+ ):
900
+ """
901
+ The feature-store repo calls feature_store_factory(..., auth_method="keypair")
902
+ in update_features.py and both notebooks; removing the parameter would break
903
+ every scheduled run on the next version bump.
904
+ """
905
+ monkeypatch.setenv("GITLAB_CI", "true")
906
+ env = dict(
907
+ make_ci_env(rsa_pem),
908
+ SNOWFLAKE_FEATURE_STORE_DATABASE="FS_DB",
909
+ SNOWFLAKE_FEATURE_STORE_SCHEMA="FS_SCHEMA",
910
+ )
911
+ sc.feature_store_factory(args=env, auth_method="keypair")
912
+ assert captured_feature_store["database"] == "FS_DB"
913
+ assert isinstance(captured_feature_store["session_params"]["private_key"], bytes)
914
+
915
+ def test_password_auth_method_warns_and_still_connects(
916
+ self, captured_feature_store, monkeypatch, rsa_pem, caplog
917
+ ):
918
+ monkeypatch.setenv("GITLAB_CI", "true")
919
+ env = dict(
920
+ make_ci_env(rsa_pem),
921
+ SNOWFLAKE_FEATURE_STORE_DATABASE="FS_DB",
922
+ SNOWFLAKE_FEATURE_STORE_SCHEMA="FS_SCHEMA",
923
+ )
924
+ with caplog.at_level(logging.WARNING, logger="gitlabds.snowflake_connectors"):
925
+ sc.feature_store_factory(args=env, auth_method="password")
926
+ assert isinstance(captured_feature_store["session_params"]["private_key"], bytes)
927
+ assert "auth_method='password' is no longer supported" in caplog.text
928
+
929
+ def test_ci_path_without_the_feature_store_vars_names_them(self, monkeypatch, rsa_pem):
757
930
  monkeypatch.setenv("GITLAB_CI", "true")
758
931
  with pytest.raises(ValueError, match="SNOWFLAKE_FEATURE_STORE_DATABASE"):
759
- sc.feature_store_factory(args=dict(CI_ENV))
932
+ sc.feature_store_factory(args=make_ci_env(rsa_pem))
760
933
 
761
934
  def test_missing_schema_is_reported_clearly(self, captured_feature_store, tmp_path):
762
935
  """FeatureStore(name="") fails deep in the SDK; fail fast with context instead."""
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes