gitlabds 2.2.3__tar.gz → 2.3.1__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.1}/PKG-INFO +242 -9
  2. gitlabds-2.2.3/PKG-INFO → gitlabds-2.3.1/README.md +234 -37
  3. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/__init__.py +17 -7
  4. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/feature_descriptions.py +3 -22
  5. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/model_evaluator.py +1 -1
  6. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/serving_features.py +4 -26
  7. gitlabds-2.3.1/gitlabds/snowflake_connectors.py +450 -0
  8. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/trends.py +0 -1
  9. gitlabds-2.2.3/README.md → gitlabds-2.3.1/gitlabds.egg-info/PKG-INFO +270 -5
  10. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds.egg-info/SOURCES.txt +3 -1
  11. gitlabds-2.3.1/gitlabds.egg-info/requires.txt +17 -0
  12. {gitlabds-2.2.3 → gitlabds-2.3.1}/pyproject.toml +7 -3
  13. gitlabds-2.3.1/tests/test_snowflake_connectors.py +769 -0
  14. gitlabds-2.2.3/gitlabds.egg-info/requires.txt +0 -13
  15. {gitlabds-2.2.3 → gitlabds-2.3.1}/LICENSE +0 -0
  16. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/baselines.py +0 -0
  17. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/config_generator.py +0 -0
  18. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/dummy.py +0 -0
  19. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/feature_reduction.py +0 -0
  20. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/insights.py +0 -0
  21. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/memory_optimization.py +0 -0
  22. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/missing.py +0 -0
  23. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/monitoring_metrics.py +0 -0
  24. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/outliers.py +0 -0
  25. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/record_benchmarking.py +0 -0
  26. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds/split_data.py +0 -0
  27. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds.egg-info/dependency_links.txt +0 -0
  28. {gitlabds-2.2.3 → gitlabds-2.3.1}/gitlabds.egg-info/top_level.txt +0 -0
  29. {gitlabds-2.2.3 → gitlabds-2.3.1}/setup.cfg +0 -0
  30. {gitlabds-2.2.3 → gitlabds-2.3.1}/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.1
4
4
  Summary: GitLab Data Science Tools
5
5
  Author-email: Kevin Dietz <kdietz@gitlab.com>
6
6
  License-Expression: MIT
@@ -21,13 +21,17 @@ Requires-Dist: numpy>=1.23.5
21
21
  Requires-Dist: scipy>=1.13.1
22
22
  Requires-Dist: scikit-learn>=1.1.1
23
23
  Requires-Dist: imbalanced-learn>=0.9.1
24
+ Requires-Dist: matplotlib>=3.5.0
24
25
  Requires-Dist: seaborn>=0.13.2
25
26
  Requires-Dist: shap>=0.50.0
26
27
  Requires-Dist: tqdm>=4.66.2
27
28
  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"
29
+ Requires-Dist: pyyaml>=6.0.0
30
+ Provides-Extra: snowflake
31
+ Requires-Dist: cryptography>=41.0.0; extra == "snowflake"
32
+ Requires-Dist: snowflake-connector-python[pandas]>=3.0.0; extra == "snowflake"
33
+ Requires-Dist: snowflake-snowpark-python>=1.37.0; extra == "snowflake"
34
+ Requires-Dist: snowflake-ml-python>=1.11.0; extra == "snowflake"
31
35
  Dynamic: license-file
32
36
 
33
37
  [![PyPI version](https://badge.fury.io/py/gitlabds.svg)](https://badge.fury.io/py/gitlabds)
@@ -61,19 +65,23 @@ pip install gitlabds
61
65
 
62
66
  ### Optional Dependencies
63
67
 
64
- #### Snowflake Feature Store
68
+ #### Snowflake
65
69
 
66
- The `serve_features`, `clear_feature_serving_locks`, and `get_feature_descriptions` functions require additional Snowflake dependencies.
70
+ 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
71
 
68
72
  To use these functions, install with:
69
73
 
70
74
  ```bash
71
- pip install gitlabds[feature-store]
75
+ pip install gitlabds[snowflake]
72
76
  ```
73
77
 
74
78
  This will install:
79
+ - `snowflake-connector-python[pandas]`
75
80
  - `snowflake-snowpark-python`
76
81
  - `snowflake-ml-python`
82
+ - `cryptography`
83
+
84
+ 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
85
 
78
86
  If you try to use these functions without installing the optional dependencies, you'll get an error message with installation instructions.
79
87
 
@@ -1125,6 +1133,231 @@ feature_drift = results['feature_drift']
1125
1133
 
1126
1134
  ### Deployment
1127
1135
 
1136
+ #### Snowflake Connectors
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.
1139
+
1140
+ All of these require `pip install gitlabds[snowflake]`.
1141
+
1142
+ <details><summary> Snowflake Engine Factory </summary>
1143
+
1144
+ #### Description
1145
+ Open a Snowflake connection, resolving credentials from your dbt profile locally or from GitLab CI environment variables in CI.
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):`
1148
+
1149
+ #### Parameters:
1150
+ - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1151
+ - **_role_** : CI credential set to use. Currently only supports `DATA_SCIENCE_LOADER`
1152
+ - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1153
+ - **_profile_target_** : dbt profile output name. Only used on the local path.
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.
1156
+ - **_profiles_path_** : Override the dbt `profiles.yml` location (defaults to `~/.dbt/profiles.yml`). Only used on the local path.
1157
+ - **_\*\*connect_kwargs_** : Passed through to `snowflake.connector.connect()`, e.g. `autocommit=False`, `client_session_keep_alive=True`.
1158
+
1159
+ #### Returns
1160
+ - An open `snowflake.connector.SnowflakeConnection`. The caller owns its lifecycle.
1161
+
1162
+ #### Requirements
1163
+ ```bash
1164
+ pip install gitlabds[snowflake]
1165
+ ```
1166
+
1167
+ #### Examples:
1168
+
1169
+ ```python
1170
+ import gitlabds
1171
+
1172
+ # Local: reads the "dev" target from ~/.dbt/profiles.yml and authenticates via SSO
1173
+ conn = gitlabds.snowflake_engine_factory()
1174
+ try:
1175
+ df = gitlabds.ds_query_dataframe(conn, "select * from prod.data_science.my_table")
1176
+ finally:
1177
+ conn.close()
1178
+ ```
1179
+ ```python
1180
+ # CI: credentials come from os.environ; run_target is inferred from GITLAB_CI
1181
+ conn = gitlabds.snowflake_engine_factory(role="DATA_SCIENCE_LOADER", schema="data_science")
1182
+
1183
+ # Opt out of autocommit to group statements into one transaction
1184
+ conn = gitlabds.snowflake_engine_factory(autocommit=False)
1185
+ ```
1186
+ </details>
1187
+
1188
+ <details><summary> DS Query Dataframe </summary>
1189
+
1190
+ #### Description
1191
+ 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.
1192
+
1193
+ `gitlabds.ds_query_dataframe(conn, query, optimize_memory=False, verbose=False):`
1194
+
1195
+ #### Parameters:
1196
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1197
+ - **_query_** : SELECT query to run
1198
+ - **_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).
1199
+ - **_verbose_** : Print the memory-optimization report and progress bar. Only applies when `optimize_memory=True`.
1200
+
1201
+ #### Returns
1202
+ - A pandas DataFrame
1203
+
1204
+ #### Requirements
1205
+ ```bash
1206
+ pip install gitlabds[snowflake]
1207
+ ```
1208
+
1209
+ #### Examples:
1210
+
1211
+ ```python
1212
+ df = gitlabds.ds_query_dataframe(conn, "select * from my_table")
1213
+
1214
+ # Large pull: downcast dtypes on the way out
1215
+ df = gitlabds.ds_query_dataframe(conn, training_query, optimize_memory=True, verbose=True)
1216
+ ```
1217
+ </details>
1218
+
1219
+ <details><summary> Query Executor </summary>
1220
+
1221
+ #### Description
1222
+ 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.
1223
+
1224
+ `gitlabds.query_executor(conn, query):`
1225
+
1226
+ #### Parameters:
1227
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1228
+ - **_query_** : SQL statement to run
1229
+
1230
+ #### Returns
1231
+ - `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.
1232
+
1233
+ #### Requirements
1234
+ ```bash
1235
+ pip install gitlabds[snowflake]
1236
+ ```
1237
+
1238
+ #### Examples:
1239
+
1240
+ ```python
1241
+ rows = gitlabds.query_executor(conn, "select count(*) as n from my_table")
1242
+ print(rows[0]["N"])
1243
+ ```
1244
+ ```python
1245
+ # Connections autocommit by default, so group statements explicitly
1246
+ gitlabds.query_executor(conn, "begin;")
1247
+ gitlabds.query_executor(conn, "delete from my_table where run_date = '2026-01-01';")
1248
+ gitlabds.query_executor(conn, "commit;")
1249
+ ```
1250
+ </details>
1251
+
1252
+ <details><summary> Snowflake Stage Load Copy Remove </summary>
1253
+
1254
+ #### Description
1255
+ 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.
1256
+
1257
+ `gitlabds.snowflake_stage_load_copy_remove(file, stage, table_path, conn, type="json", on_error="abort_statement", file_format_options="", col_names=""):`
1258
+
1259
+ #### Parameters:
1260
+ - **_conn_** : An open connection, e.g. from `snowflake_engine_factory()`
1261
+ - **_file_** : Path to the local file to upload
1262
+ - **_stage_** : Internal stage name **without** a leading `@` — the function adds it (e.g. `"my_stage"` or `"db.schema.my_stage"`)
1263
+ - **_table_path_** : Fully-qualified target table
1264
+ - **_type_** : File format type (e.g. `"json"`, `"csv"`)
1265
+ - **_on_error_** : `COPY INTO` `ON_ERROR` behavior
1266
+ - **_file_format_options_** : Extra `FILE_FORMAT` options appended after `TYPE`, e.g. `'FIELD_DELIMITER="," SKIP_HEADER=1'`
1267
+ - **_col_names_** : Explicit column list for `COPY INTO`, e.g. `"(first_name, last_name)"`. Defaults to `"(jsontext)"` for json, empty otherwise.
1268
+
1269
+ #### Returns
1270
+ - `None`. Progress is emitted at INFO level through the `logging` module; call `logging.basicConfig()` in your entry point to see it.
1271
+
1272
+ #### Requirements
1273
+ ```bash
1274
+ pip install gitlabds[snowflake]
1275
+ ```
1276
+
1277
+ #### Examples:
1278
+
1279
+ ```python
1280
+ gitlabds.snowflake_stage_load_copy_remove(
1281
+ file="scores.json", stage="my_stage", table_path="prod.data_science.scores", conn=conn
1282
+ )
1283
+ ```
1284
+ ```python
1285
+ # CSV with a header row and an explicit column list
1286
+ gitlabds.snowflake_stage_load_copy_remove(
1287
+ file="scores.csv",
1288
+ stage="my_stage",
1289
+ table_path="prod.data_science.scores",
1290
+ conn=conn,
1291
+ type="csv",
1292
+ file_format_options='FIELD_DELIMITER="," SKIP_HEADER=1',
1293
+ col_names="(account_id, score)",
1294
+ )
1295
+ ```
1296
+ </details>
1297
+
1298
+ <details><summary> Feature Store Factory </summary>
1299
+
1300
+ #### Description
1301
+ Create a Snowflake Feature Store connection (Snowpark Session + native `FeatureStore`), resolving credentials the same way as `snowflake_engine_factory`.
1302
+
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):`
1304
+
1305
+ #### Parameters:
1306
+ - **_args_** : Environment mapping to read CI credentials from. Defaults to `os.environ`. Ignored on the local path.
1307
+ - **_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
+ - **_run_target_** : `"local"` or `"ci"`. Inferred from the `GITLAB_CI` environment variable when not given.
1309
+ - **_profile_target_** : dbt profile output name. Only used on the local path.
1310
+ - **_auth_method_** : `"password"` or `"keypair"`. Only used on the CI path.
1311
+ - **_creation_mode_** : Either a `"CreationMode.X"` string or a `snowflake.ml.feature_store.CreationMode` member.
1312
+ - **_profiles_path_** : Override the dbt `profiles.yml` location. Only used on the local path.
1313
+
1314
+ #### Returns
1315
+ - A `snowflake.ml.feature_store.FeatureStore`. The underlying Snowpark session is reachable as `fs._session` and is not closed by this function.
1316
+
1317
+ #### Requirements
1318
+ ```bash
1319
+ pip install gitlabds[snowflake]
1320
+ ```
1321
+
1322
+ #### Examples:
1323
+
1324
+ ```python
1325
+ import gitlabds
1326
+
1327
+ fs = gitlabds.feature_store_factory()
1328
+
1329
+ df = gitlabds.serve_features(
1330
+ session=fs._session,
1331
+ feature_store=fs,
1332
+ feature_views_dict={"account_ci_usage": "1.0"},
1333
+ spine_df="select account_id from prod.data_science.my_spine",
1334
+ )
1335
+ ```
1336
+ </details>
1337
+
1338
+ <details><summary> Migrating from gitlabdata.orchestration_utils </summary>
1339
+
1340
+ 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:
1341
+
1342
+ | Was | Now |
1343
+ |---|---|
1344
+ | `snowflake_engine_factory(...)` returns a SQLAlchemy `Engine` | returns a native `SnowflakeConnection` |
1345
+ | `data_science_engine_factory(...)` | folded into `snowflake_engine_factory(...)`, parameterized by `role` / `profile_target` |
1346
+ | `pd.read_sql(query, engine)` | `gitlabds.ds_query_dataframe(conn, query)` |
1347
+ | `df.to_sql(..., engine)` | `snowflake.connector.pandas_tools.write_pandas(conn, df, ...)` |
1348
+ | `engine.connect()` / `engine.dispose()` | `conn` is already open; `conn.close()` |
1349
+ | `query_executor(engine, query)` returns tuples — `row[0]` | returns dicts — `row["COLUMN_NAME"]` |
1350
+ | `query_executor(..., dispose_engine=..., connection=...)` | dropped; the caller owns the connection and passes it in |
1351
+ | `snowflake_stage_load_copy_remove(..., engine, type=...)` | unchanged, except `engine` is now a connection rather than a SQLAlchemy Engine |
1352
+ | `in_docker=True` | `profiles_path="/usr/local/snowflake_profile/profiles.yml"` |
1353
+ | `run_target="local"` default | `run_target=None`, inferred from `GITLAB_CI` |
1354
+
1355
+ Also worth knowing:
1356
+ - `file_format_options` is now honored for `json` too; the original silently dropped it for that type.
1357
+ - `stage` must be passed **without** a leading `@` (the function adds it). The original behaved the same way but did not document or enforce it.
1358
+ - These functions log through `logging.getLogger(__name__)` and no longer call `logging.basicConfig()` for you.
1359
+ </details>
1360
+
1128
1361
  #### Feature Serving
1129
1362
 
1130
1363
  <details><summary> Serve Features </summary>
@@ -1157,7 +1390,7 @@ Serve features from Snowflake Feature Store with built-in distributed locking an
1157
1390
  #### Requirements
1158
1391
  This function requires Snowflake packages. Install with:
1159
1392
  ```bash
1160
- pip install snowflake-snowpark-python snowflake-ml-python
1393
+ pip install gitlabds[snowflake]
1161
1394
  ```
1162
1395
 
1163
1396
  #### Examples:
@@ -1252,7 +1485,7 @@ Look up each feature's source feature view and description from the feature stor
1252
1485
  #### Requirements
1253
1486
  This function requires Snowflake packages. Install with:
1254
1487
  ```bash
1255
- pip install snowflake-snowpark-python snowflake-ml-python
1488
+ pip install gitlabds[snowflake]
1256
1489
  ```
1257
1490
 
1258
1491
  #### 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