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.
- {gitlabds-2.3.1 → gitlabds-2.4.0}/PKG-INFO +10 -6
- {gitlabds-2.3.1 → gitlabds-2.4.0}/README.md +9 -5
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/snowflake_connectors.py +90 -24
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/PKG-INFO +10 -6
- {gitlabds-2.3.1 → gitlabds-2.4.0}/tests/test_snowflake_connectors.py +232 -59
- {gitlabds-2.3.1 → gitlabds-2.4.0}/LICENSE +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/__init__.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/baselines.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/config_generator.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/dummy.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/feature_descriptions.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/feature_reduction.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/insights.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/memory_optimization.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/missing.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/model_evaluator.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/monitoring_metrics.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/outliers.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/record_benchmarking.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/serving_features.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/split_data.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds/trends.py +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/SOURCES.txt +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/dependency_links.txt +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/requires.txt +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/gitlabds.egg-info/top_level.txt +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/pyproject.toml +0 -0
- {gitlabds-2.3.1 → gitlabds-2.4.0}/setup.cfg +0 -0
- {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
|
+
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 +
|
|
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="
|
|
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_** : `"
|
|
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="
|
|
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_** : `"
|
|
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 +
|
|
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="
|
|
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_** : `"
|
|
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="
|
|
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_** : `"
|
|
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
|
-
|
|
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
|
-
"""
|
|
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
|
-
|
|
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
|
|
135
|
-
|
|
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",
|
|
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
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
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 = "
|
|
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="
|
|
196
|
-
|
|
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 = "
|
|
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="
|
|
391
|
-
|
|
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
|
+
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 +
|
|
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="
|
|
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_** : `"
|
|
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="
|
|
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_** : `"
|
|
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
|
-
|
|
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, "
|
|
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="
|
|
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
|
|
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["
|
|
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", "", "
|
|
326
|
+
sc._resolve_connect_params("prod", "DATA_SCIENCE_LOADER", "dev", "", "keypair", {})
|
|
281
327
|
|
|
282
|
-
def
|
|
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
|
-
"
|
|
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
|
|
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
|
|
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(
|
|
308
|
-
with pytest.raises(ValueError, match="
|
|
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
|
|
320
|
-
|
|
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
|
|
394
|
+
def test_uses_the_key_password_when_present(self):
|
|
326
395
|
env = dict(
|
|
327
|
-
|
|
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
|
|
400
|
+
assert isinstance(self._resolve(None, env=env)["private_key"], bytes)
|
|
332
401
|
|
|
333
|
-
def
|
|
334
|
-
"""
|
|
335
|
-
|
|
336
|
-
|
|
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
|
|
339
|
-
env
|
|
340
|
-
|
|
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=
|
|
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(
|
|
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=
|
|
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["
|
|
474
|
+
assert isinstance(captured_connect["private_key"], bytes)
|
|
382
475
|
|
|
383
|
-
def test_gitlab_ci_inside_args_does_not_force_the_ci_path(
|
|
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(
|
|
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
|
|
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
|
-
|
|
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
|
|
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=
|
|
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
|
|
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
|
|
File without changes
|
|
File without changes
|