gitlabds 2.2.3__tar.gz → 2.3.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 (30) hide show
  1. {gitlabds-2.2.3/gitlabds.egg-info → gitlabds-2.3.0}/PKG-INFO +241 -9
  2. gitlabds-2.2.3/PKG-INFO → gitlabds-2.3.0/README.md +234 -37
  3. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/__init__.py +17 -7
  4. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/feature_descriptions.py +3 -22
  5. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/serving_features.py +4 -26
  6. gitlabds-2.3.0/gitlabds/snowflake_connectors.py +450 -0
  7. gitlabds-2.2.3/README.md → gitlabds-2.3.0/gitlabds.egg-info/PKG-INFO +269 -5
  8. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds.egg-info/SOURCES.txt +3 -1
  9. gitlabds-2.3.0/gitlabds.egg-info/requires.txt +16 -0
  10. {gitlabds-2.2.3 → gitlabds-2.3.0}/pyproject.toml +6 -3
  11. gitlabds-2.3.0/tests/test_snowflake_connectors.py +769 -0
  12. gitlabds-2.2.3/gitlabds.egg-info/requires.txt +0 -13
  13. {gitlabds-2.2.3 → gitlabds-2.3.0}/LICENSE +0 -0
  14. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/baselines.py +0 -0
  15. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/config_generator.py +0 -0
  16. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/dummy.py +0 -0
  17. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/feature_reduction.py +0 -0
  18. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/insights.py +0 -0
  19. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/memory_optimization.py +0 -0
  20. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/missing.py +0 -0
  21. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/model_evaluator.py +0 -0
  22. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/monitoring_metrics.py +0 -0
  23. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/outliers.py +0 -0
  24. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/record_benchmarking.py +0 -0
  25. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/split_data.py +0 -0
  26. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds/trends.py +0 -0
  27. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds.egg-info/dependency_links.txt +0 -0
  28. {gitlabds-2.2.3 → gitlabds-2.3.0}/gitlabds.egg-info/top_level.txt +0 -0
  29. {gitlabds-2.2.3 → gitlabds-2.3.0}/setup.cfg +0 -0
  30. {gitlabds-2.2.3 → gitlabds-2.3.0}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gitlabds
3
- Version: 2.2.3
3
+ Version: 2.3.0
4
4
  Summary: GitLab Data Science Tools
5
5
  Author-email: Kevin Dietz <kdietz@gitlab.com>
6
6
  License-Expression: MIT
@@ -25,9 +25,12 @@ Requires-Dist: seaborn>=0.13.2
25
25
  Requires-Dist: shap>=0.50.0
26
26
  Requires-Dist: tqdm>=4.66.2
27
27
  Requires-Dist: xgboost>=1.6.1
28
- Provides-Extra: feature-store
29
- Requires-Dist: snowflake-snowpark-python; extra == "feature-store"
30
- Requires-Dist: snowflake-ml-python; extra == "feature-store"
28
+ Requires-Dist: pyyaml>=6.0.0
29
+ Provides-Extra: snowflake
30
+ Requires-Dist: cryptography>=41.0.0; extra == "snowflake"
31
+ Requires-Dist: snowflake-connector-python[pandas]>=3.0.0; extra == "snowflake"
32
+ Requires-Dist: snowflake-snowpark-python>=1.37.0; extra == "snowflake"
33
+ Requires-Dist: snowflake-ml-python>=1.11.0; extra == "snowflake"
31
34
  Dynamic: license-file
32
35
 
33
36
  [![PyPI version](https://badge.fury.io/py/gitlabds.svg)](https://badge.fury.io/py/gitlabds)
@@ -61,19 +64,23 @@ pip install gitlabds
61
64
 
62
65
  ### Optional Dependencies
63
66
 
64
- #### Snowflake Feature Store
67
+ #### Snowflake
65
68
 
66
- The `serve_features`, `clear_feature_serving_locks`, and `get_feature_descriptions` functions require additional Snowflake dependencies.
69
+ The Snowflake connectors (`snowflake_engine_factory`, `ds_query_dataframe`, `query_executor`, `snowflake_stage_load_copy_remove`, `feature_store_factory`) and the feature serving functions (`serve_features`, `clear_feature_serving_locks`, `get_feature_descriptions`) require additional dependencies.
67
70
 
68
71
  To use these functions, install with:
69
72
 
70
73
  ```bash
71
- pip install gitlabds[feature-store]
74
+ pip install gitlabds[snowflake]
72
75
  ```
73
76
 
74
77
  This will install:
78
+ - `snowflake-connector-python[pandas]`
75
79
  - `snowflake-snowpark-python`
76
80
  - `snowflake-ml-python`
81
+ - `cryptography`
82
+
83
+ These are optional rather than core because `snowflake-ml-python` constrains the numpy/scikit-learn/xgboost stack that the core package also depends on.
77
84
 
78
85
  If you try to use these functions without installing the optional dependencies, you'll get an error message with installation instructions.
79
86
 
@@ -1125,6 +1132,231 @@ feature_drift = results['feature_drift']
1125
1132
 
1126
1133
  ### Deployment
1127
1134
 
1135
+ #### Snowflake Connectors
1136
+
1137
+ 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
+
1139
+ All of these require `pip install gitlabds[snowflake]`.
1140
+
1141
+ <details><summary> Snowflake Engine Factory </summary>
1142
+
1143
+ #### Description
1144
+ Open a Snowflake connection, resolving credentials from your dbt profile locally or from GitLab CI environment variables in CI.
1145
+
1146
+ `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
+
1148
+ #### Parameters:
1149
+ - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1150
+ - **_role_** : CI credential set to use. Currently only supports `DATA_SCIENCE_LOADER`
1151
+ - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1152
+ - **_profile_target_** : dbt profile output name. Only used on the local path.
1153
+ - **_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.
1154
+ - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path. Key pair secrets are not yet provisioned in CI.
1155
+ - **_profiles_path_** : Override the dbt `profiles.yml` location (defaults to `~/.dbt/profiles.yml`). Only used on the local path.
1156
+ - **_\*\*connect_kwargs_** : Passed through to `snowflake.connector.connect()`, e.g. `autocommit=False`, `client_session_keep_alive=True`.
1157
+
1158
+ #### Returns
1159
+ - An open `snowflake.connector.SnowflakeConnection`. The caller owns its lifecycle.
1160
+
1161
+ #### Requirements
1162
+ ```bash
1163
+ pip install gitlabds[snowflake]
1164
+ ```
1165
+
1166
+ #### Examples:
1167
+
1168
+ ```python
1169
+ import gitlabds
1170
+
1171
+ # Local: reads the "dev" target from ~/.dbt/profiles.yml and authenticates via SSO
1172
+ conn = gitlabds.snowflake_engine_factory()
1173
+ try:
1174
+ df = gitlabds.ds_query_dataframe(conn, "select * from prod.data_science.my_table")
1175
+ finally:
1176
+ conn.close()
1177
+ ```
1178
+ ```python
1179
+ # CI: credentials come from os.environ; run_target is inferred from GITLAB_CI
1180
+ conn = gitlabds.snowflake_engine_factory(role="DATA_SCIENCE_LOADER", schema="data_science")
1181
+
1182
+ # Opt out of autocommit to group statements into one transaction
1183
+ conn = gitlabds.snowflake_engine_factory(autocommit=False)
1184
+ ```
1185
+ </details>
1186
+
1187
+ <details><summary> DS Query Dataframe </summary>
1188
+
1189
+ #### Description
1190
+ Run a SELECT and return a pandas DataFrame, using Snowflake's Arrow-based fetch. Prefer this over `query_executor` for large reads such as pulling training data.
1191
+
1192
+ `gitlabds.ds_query_dataframe(conn, query, optimize_memory=False, verbose=False):`
1193
+
1194
+ #### Parameters:
1195
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1196
+ - **_query_** : SELECT query to run
1197
+ - **_optimize_memory_** : Run the result through `reduce_memory_usage()` (dtype downcasting) before returning. Off by default because it is a lossy transform (e.g. float64 → float32).
1198
+ - **_verbose_** : Print the memory-optimization report and progress bar. Only applies when `optimize_memory=True`.
1199
+
1200
+ #### Returns
1201
+ - A pandas DataFrame
1202
+
1203
+ #### Requirements
1204
+ ```bash
1205
+ pip install gitlabds[snowflake]
1206
+ ```
1207
+
1208
+ #### Examples:
1209
+
1210
+ ```python
1211
+ df = gitlabds.ds_query_dataframe(conn, "select * from my_table")
1212
+
1213
+ # Large pull: downcast dtypes on the way out
1214
+ df = gitlabds.ds_query_dataframe(conn, training_query, optimize_memory=True, verbose=True)
1215
+ ```
1216
+ </details>
1217
+
1218
+ <details><summary> Query Executor </summary>
1219
+
1220
+ #### Description
1221
+ Execute a SQL statement and return its rows as a list of dicts keyed by column name, as Snowflake returns them. Use for DML/DDL/transaction-control statements, or reads where a DataFrame isn't needed.
1222
+
1223
+ `gitlabds.query_executor(conn, query):`
1224
+
1225
+ #### Parameters:
1226
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1227
+ - **_query_** : SQL statement to run
1228
+
1229
+ #### Returns
1230
+ - `List[Dict[str, Any]]` — rows as `{column: value}`. DML and DDL statements return Snowflake's own status row (e.g. `{"number of rows inserted": 12}`) rather than an empty list.
1231
+
1232
+ #### Requirements
1233
+ ```bash
1234
+ pip install gitlabds[snowflake]
1235
+ ```
1236
+
1237
+ #### Examples:
1238
+
1239
+ ```python
1240
+ rows = gitlabds.query_executor(conn, "select count(*) as n from my_table")
1241
+ print(rows[0]["N"])
1242
+ ```
1243
+ ```python
1244
+ # Connections autocommit by default, so group statements explicitly
1245
+ gitlabds.query_executor(conn, "begin;")
1246
+ gitlabds.query_executor(conn, "delete from my_table where run_date = '2026-01-01';")
1247
+ gitlabds.query_executor(conn, "commit;")
1248
+ ```
1249
+ </details>
1250
+
1251
+ <details><summary> Snowflake Stage Load Copy Remove </summary>
1252
+
1253
+ #### Description
1254
+ Load a local file into a Snowflake table via an internal stage: remove any leftover staged file, PUT the file, COPY INTO the target table, then REMOVE the staged file.
1255
+
1256
+ `gitlabds.snowflake_stage_load_copy_remove(file, stage, table_path, conn, type="json", on_error="abort_statement", file_format_options="", col_names=""):`
1257
+
1258
+ #### Parameters:
1259
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1260
+ - **_file_** : Path to the local file to upload
1261
+ - **_stage_** : Internal stage name **without** a leading `@` — the function adds it (e.g. `"my_stage"` or `"db.schema.my_stage"`)
1262
+ - **_table_path_** : Fully-qualified target table
1263
+ - **_type_** : File format type (e.g. `"json"`, `"csv"`)
1264
+ - **_on_error_** : `COPY INTO` `ON_ERROR` behavior
1265
+ - **_file_format_options_** : Extra `FILE_FORMAT` options appended after `TYPE`, e.g. `'FIELD_DELIMITER="," SKIP_HEADER=1'`
1266
+ - **_col_names_** : Explicit column list for `COPY INTO`, e.g. `"(first_name, last_name)"`. Defaults to `"(jsontext)"` for json, empty otherwise.
1267
+
1268
+ #### Returns
1269
+ - `None`. Progress is emitted at INFO level through the `logging` module; call `logging.basicConfig()` in your entry point to see it.
1270
+
1271
+ #### Requirements
1272
+ ```bash
1273
+ pip install gitlabds[snowflake]
1274
+ ```
1275
+
1276
+ #### Examples:
1277
+
1278
+ ```python
1279
+ gitlabds.snowflake_stage_load_copy_remove(
1280
+ file="scores.json", stage="my_stage", table_path="prod.data_science.scores", conn=conn
1281
+ )
1282
+ ```
1283
+ ```python
1284
+ # CSV with a header row and an explicit column list
1285
+ gitlabds.snowflake_stage_load_copy_remove(
1286
+ file="scores.csv",
1287
+ stage="my_stage",
1288
+ table_path="prod.data_science.scores",
1289
+ conn=conn,
1290
+ type="csv",
1291
+ file_format_options='FIELD_DELIMITER="," SKIP_HEADER=1',
1292
+ col_names="(account_id, score)",
1293
+ )
1294
+ ```
1295
+ </details>
1296
+
1297
+ <details><summary> Feature Store Factory </summary>
1298
+
1299
+ #### Description
1300
+ Create a Snowflake Feature Store connection (Snowpark Session + native `FeatureStore`), resolving credentials the same way as `snowflake_engine_factory`.
1301
+
1302
+ `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):`
1303
+
1304
+ #### Parameters:
1305
+ - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1306
+ - **_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.
1307
+ - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1308
+ - **_profile_target_** : dbt profile output name. Only used on the local path.
1309
+ - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path.
1310
+ - **_creation_mode_** : Either a `"CreationMode.X"` string or a `snowflake.ml.feature_store.CreationMode` member.
1311
+ - **_profiles_path_** : Override the dbt `profiles.yml` location. Only used on the local path.
1312
+
1313
+ #### Returns
1314
+ - A `snowflake.ml.feature_store.FeatureStore`. The underlying Snowpark session is reachable as `fs._session` and is not closed by this function.
1315
+
1316
+ #### Requirements
1317
+ ```bash
1318
+ pip install gitlabds[snowflake]
1319
+ ```
1320
+
1321
+ #### Examples:
1322
+
1323
+ ```python
1324
+ import gitlabds
1325
+
1326
+ fs = gitlabds.feature_store_factory()
1327
+
1328
+ df = gitlabds.serve_features(
1329
+ session=fs._session,
1330
+ feature_store=fs,
1331
+ feature_views_dict={"account_ci_usage": "1.0"},
1332
+ spine_df="select account_id from prod.data_science.my_spine",
1333
+ )
1334
+ ```
1335
+ </details>
1336
+
1337
+ <details><summary> Migrating from gitlabdata.orchestration_utils </summary>
1338
+
1339
+ These functions replace their `gitlabdata.orchestration_utils` equivalents, but the signatures intentionally differ — this is not a drop-in import swap. Changes that fail at runtime rather than at import:
1340
+
1341
+ | Was | Now |
1342
+ |---|---|
1343
+ | `snowflake_engine_factory(...)` returns a SQLAlchemy `Engine` | returns a native `SnowflakeConnection` |
1344
+ | `data_science_engine_factory(...)` | folded into `snowflake_engine_factory(...)`, parameterized by `role` / `profile_target` |
1345
+ | `pd.read_sql(query, engine)` | `gitlabds.ds_query_dataframe(conn, query)` |
1346
+ | `df.to_sql(..., engine)` | `snowflake.connector.pandas_tools.write_pandas(conn, df, ...)` |
1347
+ | `engine.connect()` / `engine.dispose()` | `conn` is already open; `conn.close()` |
1348
+ | `query_executor(engine, query)` returns tuples — `row[0]` | returns dicts — `row["COLUMN_NAME"]` |
1349
+ | `query_executor(..., dispose_engine=..., connection=...)` | dropped; the caller owns the connection and passes it in |
1350
+ | `snowflake_stage_load_copy_remove(..., engine, type=...)` | unchanged, except `engine` is now a connection rather than a SQLAlchemy Engine |
1351
+ | `in_docker=True` | `profiles_path="/usr/local/snowflake_profile/profiles.yml"` |
1352
+ | `run_target="local"` default | `run_target=None`, inferred from `GITLAB_CI` |
1353
+
1354
+ Also worth knowing:
1355
+ - `file_format_options` is now honored for `json` too; the original silently dropped it for that type.
1356
+ - `stage` must be passed **without** a leading `@` (the function adds it). The original behaved the same way but did not document or enforce it.
1357
+ - These functions log through `logging.getLogger(__name__)` and no longer call `logging.basicConfig()` for you.
1358
+ </details>
1359
+
1128
1360
  #### Feature Serving
1129
1361
 
1130
1362
  <details><summary> Serve Features </summary>
@@ -1157,7 +1389,7 @@ Serve features from Snowflake Feature Store with built-in distributed locking an
1157
1389
  #### Requirements
1158
1390
  This function requires Snowflake packages. Install with:
1159
1391
  ```bash
1160
- pip install snowflake-snowpark-python snowflake-ml-python
1392
+ pip install gitlabds[snowflake]
1161
1393
  ```
1162
1394
 
1163
1395
  #### Examples:
@@ -1252,7 +1484,7 @@ Look up each feature's source feature view and description from the feature stor
1252
1484
  #### Requirements
1253
1485
  This function requires Snowflake packages. Install with:
1254
1486
  ```bash
1255
- pip install snowflake-snowpark-python snowflake-ml-python
1487
+ pip install gitlabds[snowflake]
1256
1488
  ```
1257
1489
 
1258
1490
  #### Examples:
@@ -1,35 +1,3 @@
1
- Metadata-Version: 2.4
2
- Name: gitlabds
3
- Version: 2.2.3
4
- Summary: GitLab Data Science Tools
5
- Author-email: Kevin Dietz <kdietz@gitlab.com>
6
- License-Expression: MIT
7
- Project-URL: Homepage, https://gitlab.com/gitlab-data/gitlabds
8
- Classifier: Development Status :: 5 - Production/Stable
9
- Classifier: Intended Audience :: Developers
10
- Classifier: Intended Audience :: Science/Research
11
- Classifier: Programming Language :: Python :: 3
12
- Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
13
- Classifier: Topic :: Scientific/Engineering :: Information Analysis
14
- Classifier: Topic :: Software Development :: Libraries
15
- Classifier: Topic :: Utilities
16
- Requires-Python: >=3.11
17
- Description-Content-Type: text/markdown
18
- License-File: LICENSE
19
- Requires-Dist: pandas>=1.5.3
20
- Requires-Dist: numpy>=1.23.5
21
- Requires-Dist: scipy>=1.13.1
22
- Requires-Dist: scikit-learn>=1.1.1
23
- Requires-Dist: imbalanced-learn>=0.9.1
24
- Requires-Dist: seaborn>=0.13.2
25
- Requires-Dist: shap>=0.50.0
26
- Requires-Dist: tqdm>=4.66.2
27
- Requires-Dist: xgboost>=1.6.1
28
- Provides-Extra: feature-store
29
- Requires-Dist: snowflake-snowpark-python; extra == "feature-store"
30
- Requires-Dist: snowflake-ml-python; extra == "feature-store"
31
- Dynamic: license-file
32
-
33
1
  [![PyPI version](https://badge.fury.io/py/gitlabds.svg)](https://badge.fury.io/py/gitlabds)
34
2
  [![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
35
3
  [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
@@ -61,19 +29,23 @@ pip install gitlabds
61
29
 
62
30
  ### Optional Dependencies
63
31
 
64
- #### Snowflake Feature Store
32
+ #### Snowflake
65
33
 
66
- The `serve_features`, `clear_feature_serving_locks`, and `get_feature_descriptions` functions require additional Snowflake dependencies.
34
+ The Snowflake connectors (`snowflake_engine_factory`, `ds_query_dataframe`, `query_executor`, `snowflake_stage_load_copy_remove`, `feature_store_factory`) and the feature serving functions (`serve_features`, `clear_feature_serving_locks`, `get_feature_descriptions`) require additional dependencies.
67
35
 
68
36
  To use these functions, install with:
69
37
 
70
38
  ```bash
71
- pip install gitlabds[feature-store]
39
+ pip install gitlabds[snowflake]
72
40
  ```
73
41
 
74
42
  This will install:
43
+ - `snowflake-connector-python[pandas]`
75
44
  - `snowflake-snowpark-python`
76
45
  - `snowflake-ml-python`
46
+ - `cryptography`
47
+
48
+ These are optional rather than core because `snowflake-ml-python` constrains the numpy/scikit-learn/xgboost stack that the core package also depends on.
77
49
 
78
50
  If you try to use these functions without installing the optional dependencies, you'll get an error message with installation instructions.
79
51
 
@@ -1125,6 +1097,231 @@ feature_drift = results['feature_drift']
1125
1097
 
1126
1098
  ### Deployment
1127
1099
 
1100
+ #### Snowflake Connectors
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.
1103
+
1104
+ All of these require `pip install gitlabds[snowflake]`.
1105
+
1106
+ <details><summary> Snowflake Engine Factory </summary>
1107
+
1108
+ #### Description
1109
+ Open a Snowflake connection, resolving credentials from your dbt profile locally or from GitLab CI environment variables in CI.
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):`
1112
+
1113
+ #### Parameters:
1114
+ - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1115
+ - **_role_** : CI credential set to use. Currently only supports `DATA_SCIENCE_LOADER`
1116
+ - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1117
+ - **_profile_target_** : dbt profile output name. Only used on the local path.
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.
1120
+ - **_profiles_path_** : Override the dbt `profiles.yml` location (defaults to `~/.dbt/profiles.yml`). Only used on the local path.
1121
+ - **_\*\*connect_kwargs_** : Passed through to `snowflake.connector.connect()`, e.g. `autocommit=False`, `client_session_keep_alive=True`.
1122
+
1123
+ #### Returns
1124
+ - An open `snowflake.connector.SnowflakeConnection`. The caller owns its lifecycle.
1125
+
1126
+ #### Requirements
1127
+ ```bash
1128
+ pip install gitlabds[snowflake]
1129
+ ```
1130
+
1131
+ #### Examples:
1132
+
1133
+ ```python
1134
+ import gitlabds
1135
+
1136
+ # Local: reads the "dev" target from ~/.dbt/profiles.yml and authenticates via SSO
1137
+ conn = gitlabds.snowflake_engine_factory()
1138
+ try:
1139
+ df = gitlabds.ds_query_dataframe(conn, "select * from prod.data_science.my_table")
1140
+ finally:
1141
+ conn.close()
1142
+ ```
1143
+ ```python
1144
+ # CI: credentials come from os.environ; run_target is inferred from GITLAB_CI
1145
+ conn = gitlabds.snowflake_engine_factory(role="DATA_SCIENCE_LOADER", schema="data_science")
1146
+
1147
+ # Opt out of autocommit to group statements into one transaction
1148
+ conn = gitlabds.snowflake_engine_factory(autocommit=False)
1149
+ ```
1150
+ </details>
1151
+
1152
+ <details><summary> DS Query Dataframe </summary>
1153
+
1154
+ #### Description
1155
+ Run a SELECT and return a pandas DataFrame, using Snowflake's Arrow-based fetch. Prefer this over `query_executor` for large reads such as pulling training data.
1156
+
1157
+ `gitlabds.ds_query_dataframe(conn, query, optimize_memory=False, verbose=False):`
1158
+
1159
+ #### Parameters:
1160
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1161
+ - **_query_** : SELECT query to run
1162
+ - **_optimize_memory_** : Run the result through `reduce_memory_usage()` (dtype downcasting) before returning. Off by default because it is a lossy transform (e.g. float64 → float32).
1163
+ - **_verbose_** : Print the memory-optimization report and progress bar. Only applies when `optimize_memory=True`.
1164
+
1165
+ #### Returns
1166
+ - A pandas DataFrame
1167
+
1168
+ #### Requirements
1169
+ ```bash
1170
+ pip install gitlabds[snowflake]
1171
+ ```
1172
+
1173
+ #### Examples:
1174
+
1175
+ ```python
1176
+ df = gitlabds.ds_query_dataframe(conn, "select * from my_table")
1177
+
1178
+ # Large pull: downcast dtypes on the way out
1179
+ df = gitlabds.ds_query_dataframe(conn, training_query, optimize_memory=True, verbose=True)
1180
+ ```
1181
+ </details>
1182
+
1183
+ <details><summary> Query Executor </summary>
1184
+
1185
+ #### Description
1186
+ Execute a SQL statement and return its rows as a list of dicts keyed by column name, as Snowflake returns them. Use for DML/DDL/transaction-control statements, or reads where a DataFrame isn't needed.
1187
+
1188
+ `gitlabds.query_executor(conn, query):`
1189
+
1190
+ #### Parameters:
1191
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1192
+ - **_query_** : SQL statement to run
1193
+
1194
+ #### Returns
1195
+ - `List[Dict[str, Any]]` — rows as `{column: value}`. DML and DDL statements return Snowflake's own status row (e.g. `{"number of rows inserted": 12}`) rather than an empty list.
1196
+
1197
+ #### Requirements
1198
+ ```bash
1199
+ pip install gitlabds[snowflake]
1200
+ ```
1201
+
1202
+ #### Examples:
1203
+
1204
+ ```python
1205
+ rows = gitlabds.query_executor(conn, "select count(*) as n from my_table")
1206
+ print(rows[0]["N"])
1207
+ ```
1208
+ ```python
1209
+ # Connections autocommit by default, so group statements explicitly
1210
+ gitlabds.query_executor(conn, "begin;")
1211
+ gitlabds.query_executor(conn, "delete from my_table where run_date = '2026-01-01';")
1212
+ gitlabds.query_executor(conn, "commit;")
1213
+ ```
1214
+ </details>
1215
+
1216
+ <details><summary> Snowflake Stage Load Copy Remove </summary>
1217
+
1218
+ #### Description
1219
+ Load a local file into a Snowflake table via an internal stage: remove any leftover staged file, PUT the file, COPY INTO the target table, then REMOVE the staged file.
1220
+
1221
+ `gitlabds.snowflake_stage_load_copy_remove(file, stage, table_path, conn, type="json", on_error="abort_statement", file_format_options="", col_names=""):`
1222
+
1223
+ #### Parameters:
1224
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1225
+ - **_file_** : Path to the local file to upload
1226
+ - **_stage_** : Internal stage name **without** a leading `@` — the function adds it (e.g. `"my_stage"` or `"db.schema.my_stage"`)
1227
+ - **_table_path_** : Fully-qualified target table
1228
+ - **_type_** : File format type (e.g. `"json"`, `"csv"`)
1229
+ - **_on_error_** : `COPY INTO` `ON_ERROR` behavior
1230
+ - **_file_format_options_** : Extra `FILE_FORMAT` options appended after `TYPE`, e.g. `'FIELD_DELIMITER="," SKIP_HEADER=1'`
1231
+ - **_col_names_** : Explicit column list for `COPY INTO`, e.g. `"(first_name, last_name)"`. Defaults to `"(jsontext)"` for json, empty otherwise.
1232
+
1233
+ #### Returns
1234
+ - `None`. Progress is emitted at INFO level through the `logging` module; call `logging.basicConfig()` in your entry point to see it.
1235
+
1236
+ #### Requirements
1237
+ ```bash
1238
+ pip install gitlabds[snowflake]
1239
+ ```
1240
+
1241
+ #### Examples:
1242
+
1243
+ ```python
1244
+ gitlabds.snowflake_stage_load_copy_remove(
1245
+ file="scores.json", stage="my_stage", table_path="prod.data_science.scores", conn=conn
1246
+ )
1247
+ ```
1248
+ ```python
1249
+ # CSV with a header row and an explicit column list
1250
+ gitlabds.snowflake_stage_load_copy_remove(
1251
+ file="scores.csv",
1252
+ stage="my_stage",
1253
+ table_path="prod.data_science.scores",
1254
+ conn=conn,
1255
+ type="csv",
1256
+ file_format_options='FIELD_DELIMITER="," SKIP_HEADER=1',
1257
+ col_names="(account_id, score)",
1258
+ )
1259
+ ```
1260
+ </details>
1261
+
1262
+ <details><summary> Feature Store Factory </summary>
1263
+
1264
+ #### Description
1265
+ Create a Snowflake Feature Store connection (Snowpark Session + native `FeatureStore`), resolving credentials the same way as `snowflake_engine_factory`.
1266
+
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):`
1268
+
1269
+ #### Parameters:
1270
+ - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1271
+ - **_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
+ - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1273
+ - **_profile_target_** : dbt profile output name. Only used on the local path.
1274
+ - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path.
1275
+ - **_creation_mode_** : Either a `"CreationMode.X"` string or a `snowflake.ml.feature_store.CreationMode` member.
1276
+ - **_profiles_path_** : Override the dbt `profiles.yml` location. Only used on the local path.
1277
+
1278
+ #### Returns
1279
+ - A `snowflake.ml.feature_store.FeatureStore`. The underlying Snowpark session is reachable as `fs._session` and is not closed by this function.
1280
+
1281
+ #### Requirements
1282
+ ```bash
1283
+ pip install gitlabds[snowflake]
1284
+ ```
1285
+
1286
+ #### Examples:
1287
+
1288
+ ```python
1289
+ import gitlabds
1290
+
1291
+ fs = gitlabds.feature_store_factory()
1292
+
1293
+ df = gitlabds.serve_features(
1294
+ session=fs._session,
1295
+ feature_store=fs,
1296
+ feature_views_dict={"account_ci_usage": "1.0"},
1297
+ spine_df="select account_id from prod.data_science.my_spine",
1298
+ )
1299
+ ```
1300
+ </details>
1301
+
1302
+ <details><summary> Migrating from gitlabdata.orchestration_utils </summary>
1303
+
1304
+ These functions replace their `gitlabdata.orchestration_utils` equivalents, but the signatures intentionally differ — this is not a drop-in import swap. Changes that fail at runtime rather than at import:
1305
+
1306
+ | Was | Now |
1307
+ |---|---|
1308
+ | `snowflake_engine_factory(...)` returns a SQLAlchemy `Engine` | returns a native `SnowflakeConnection` |
1309
+ | `data_science_engine_factory(...)` | folded into `snowflake_engine_factory(...)`, parameterized by `role` / `profile_target` |
1310
+ | `pd.read_sql(query, engine)` | `gitlabds.ds_query_dataframe(conn, query)` |
1311
+ | `df.to_sql(..., engine)` | `snowflake.connector.pandas_tools.write_pandas(conn, df, ...)` |
1312
+ | `engine.connect()` / `engine.dispose()` | `conn` is already open; `conn.close()` |
1313
+ | `query_executor(engine, query)` returns tuples — `row[0]` | returns dicts — `row["COLUMN_NAME"]` |
1314
+ | `query_executor(..., dispose_engine=..., connection=...)` | dropped; the caller owns the connection and passes it in |
1315
+ | `snowflake_stage_load_copy_remove(..., engine, type=...)` | unchanged, except `engine` is now a connection rather than a SQLAlchemy Engine |
1316
+ | `in_docker=True` | `profiles_path="/usr/local/snowflake_profile/profiles.yml"` |
1317
+ | `run_target="local"` default | `run_target=None`, inferred from `GITLAB_CI` |
1318
+
1319
+ Also worth knowing:
1320
+ - `file_format_options` is now honored for `json` too; the original silently dropped it for that type.
1321
+ - `stage` must be passed **without** a leading `@` (the function adds it). The original behaved the same way but did not document or enforce it.
1322
+ - These functions log through `logging.getLogger(__name__)` and no longer call `logging.basicConfig()` for you.
1323
+ </details>
1324
+
1128
1325
  #### Feature Serving
1129
1326
 
1130
1327
  <details><summary> Serve Features </summary>
@@ -1157,7 +1354,7 @@ Serve features from Snowflake Feature Store with built-in distributed locking an
1157
1354
  #### Requirements
1158
1355
  This function requires Snowflake packages. Install with:
1159
1356
  ```bash
1160
- pip install snowflake-snowpark-python snowflake-ml-python
1357
+ pip install gitlabds[snowflake]
1161
1358
  ```
1162
1359
 
1163
1360
  #### Examples:
@@ -1252,7 +1449,7 @@ Look up each feature's source feature view and description from the feature stor
1252
1449
  #### Requirements
1253
1450
  This function requires Snowflake packages. Install with:
1254
1451
  ```bash
1255
- pip install snowflake-snowpark-python snowflake-ml-python
1452
+ pip install gitlabds[snowflake]
1256
1453
  ```
1257
1454
 
1258
1455
  #### Examples:
@@ -12,19 +12,29 @@ from .outliers import mad_outliers, apply_outliers
12
12
  from .split_data import split_data
13
13
  from .trends import generate_sql_trend_query, trend_analysis
14
14
 
15
- # Optional imports for Snowflake feature serving
15
+
16
+ # Snowflake connectors and feature serving. Optional so that a missing Snowflake
17
+ # install doesn't make the whole package unimportable.
16
18
  try:
19
+ from .snowflake_connectors import (
20
+ snowflake_engine_factory,
21
+ ds_query_dataframe,
22
+ query_executor,
23
+ snowflake_stage_load_copy_remove,
24
+ feature_store_factory,
25
+ )
17
26
  from .serving_features import clear_feature_serving_locks, serve_features
18
27
  from .feature_descriptions import get_feature_descriptions
19
- except ImportError as e:
20
- # Snowflake packages not installed - create placeholder functions with helpful error messages
28
+ except ImportError:
29
+
21
30
  def _snowflake_not_installed(*args, **kwargs):
22
31
  raise ImportError(
23
- "Snowflake feature serving functions require additional dependencies. "
24
- "Install them with: pip install snowflake-snowpark-python snowflake-ml-python\n\n"
32
+ "Snowflake functions require additional dependencies. "
33
+ "Install them with: pip install gitlabds[snowflake]"
25
34
  )
26
35
 
27
- clear_feature_serving_locks = _snowflake_not_installed
28
- serve_features = _snowflake_not_installed
36
+ snowflake_engine_factory = ds_query_dataframe = query_executor = _snowflake_not_installed
37
+ snowflake_stage_load_copy_remove = feature_store_factory = _snowflake_not_installed
38
+ clear_feature_serving_locks = serve_features = _snowflake_not_installed
29
39
  get_feature_descriptions = _snowflake_not_installed
30
40