fullscript.spellbook 2.1.6__tar.gz → 2.2.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 (46) hide show
  1. fullscript_spellbook-2.2.0/.gitignore +9 -0
  2. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/PKG-INFO +16 -3
  3. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/README.md +11 -1
  4. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/documentation/db_magic.md +64 -1
  5. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/fullscript.spellbook.egg-info/PKG-INFO +16 -3
  6. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/fullscript.spellbook.egg-info/SOURCES.txt +4 -3
  7. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/fullscript.spellbook.egg-info/requires.txt +4 -0
  8. fullscript_spellbook-2.2.0/fullscript.spellbook.egg-info/scm_file_list.json +38 -0
  9. fullscript_spellbook-2.2.0/fullscript.spellbook.egg-info/scm_version.json +8 -0
  10. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/pyproject.toml +9 -1
  11. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/requirements.txt +3 -3
  12. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/db_magic.py +136 -6
  13. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook_config.yaml +11 -0
  14. fullscript_spellbook-2.2.0/tests/test_db_magic.py +316 -0
  15. fullscript_spellbook-2.1.6/.DS_Store +0 -0
  16. fullscript_spellbook-2.1.6/.gitignore +0 -3
  17. fullscript_spellbook-2.1.6/spellbook/.DS_Store +0 -0
  18. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/.gitlab-ci.yml +0 -0
  19. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/.python-version +0 -0
  20. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/LICENSE +0 -0
  21. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/MANIFEST.in +0 -0
  22. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/documentation/drive_magic.md +0 -0
  23. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/documentation/sheet_magic.md +0 -0
  24. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/documentation/snowflake_key_pair_setup.md +0 -0
  25. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/documentation/stats_magic.md +0 -0
  26. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/fullscript.spellbook.egg-info/dependency_links.txt +0 -0
  27. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/fullscript.spellbook.egg-info/top_level.txt +0 -0
  28. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/setup.cfg +0 -0
  29. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/__init__.py +0 -0
  30. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/drive_magic.py +0 -0
  31. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/sheet_magic.py +0 -0
  32. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/stats_magic.py +0 -0
  33. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/utils.py +0 -0
  34. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/__init__.py +0 -0
  35. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard1.txt +0 -0
  36. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard10.txt +0 -0
  37. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard11.txt +0 -0
  38. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard12.txt +0 -0
  39. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard2.txt +0 -0
  40. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard3.txt +0 -0
  41. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard4.txt +0 -0
  42. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard5.txt +0 -0
  43. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard6.txt +0 -0
  44. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard7.txt +0 -0
  45. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard8.txt +0 -0
  46. {fullscript_spellbook-2.1.6 → fullscript_spellbook-2.2.0}/spellbook/wizards/wizard9.txt +0 -0
@@ -0,0 +1,9 @@
1
+ *.whl
2
+ *.egg-info
3
+ *.tar.gz
4
+ __pycache__/
5
+ .pytest_cache/
6
+ .idea/
7
+ venv/
8
+ /.junie/
9
+ .DS_Store
@@ -1,8 +1,8 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fullscript.spellbook
3
- Version: 2.1.6
3
+ Version: 2.2.0
4
4
  Summary: A simple package to make your life easier.
5
- Author-email: Thomas Edwards <thomas.edwards@fullscript.com>, Gui Chagas <gui.chagas@fullscript.com>, David King <david.king@fullscript.com>, Meinhard Mare <meinhard.mare@fullscript.com>, Karsten Melchers <karsten.melchers@fullscript.com>
5
+ Author-email: Meinhard Mare <meinhard.mare@fullscript.com>, Thomas Edwards <thomas.edwards@fullscript.com>, Gui Chagas <gui.chagas@fullscript.com>, David King <david.king@fullscript.com>, Karsten Melchers <karsten.melchers@fullscript.com>
6
6
  License: MIT
7
7
  Project-URL: homepage, https://github.com/Fullscript/spellbook
8
8
  Requires-Python: <3.12,>=3.10
@@ -46,6 +46,7 @@ Requires-Dist: pycparser==2.22
46
46
  Requires-Dist: Pygments==2.18.0
47
47
  Requires-Dist: PyJWT==2.9.0
48
48
  Requires-Dist: pyOpenSSL==24.2.1
49
+ Requires-Dist: PyMySQL==1.2.3
49
50
  Requires-Dist: python-dateutil==2.9.0.post0
50
51
  Requires-Dist: pytz==2024.2
51
52
  Requires-Dist: PyYAML==6.0.2
@@ -116,6 +117,8 @@ Requires-Dist: scipy==1.14.1; extra == "statsmagic"
116
117
  Requires-Dist: seaborn==0.13.2; extra == "statsmagic"
117
118
  Requires-Dist: six==1.16.0; extra == "statsmagic"
118
119
  Requires-Dist: statsmodels==0.14.4; extra == "statsmagic"
120
+ Provides-Extra: test
121
+ Requires-Dist: pytest==8.3.4; extra == "test"
119
122
  Dynamic: license-file
120
123
 
121
124
  # Spellbook
@@ -159,7 +162,7 @@ Spellbook requires Python 3.11 or higher. Use the following steps to install:
159
162
  1. Pip install:
160
163
 
161
164
  ```bash
162
- pip install git+https://git.fullscript.io/data/spellbook.git
165
+ pip install fullscript.spellbook
163
166
  ```
164
167
 
165
168
  ## Configuration
@@ -186,6 +189,16 @@ databases:
186
189
  database: snowflake_database
187
190
  schema: snowflake_schema
188
191
 
192
+ # MySQL Database Connection using AWS RDS IAM authentication (no stored password).
193
+ # Fullscript internal only: requires the `rx` CLI. See documentation/db_magic.md.
194
+ - name: alias_of_your_connection
195
+ type: mysql
196
+ host: your_host.rds.amazonaws.com
197
+ port: 3306
198
+ user: ${DB_USER}
199
+ iam_auth: true
200
+ database: ${MYSQL_DATABASE}
201
+
189
202
  google_accounts:
190
203
  # Google Sheets Service Account 1
191
204
  - name: google_sheets_service_account_1
@@ -39,7 +39,7 @@ Spellbook requires Python 3.11 or higher. Use the following steps to install:
39
39
  1. Pip install:
40
40
 
41
41
  ```bash
42
- pip install git+https://git.fullscript.io/data/spellbook.git
42
+ pip install fullscript.spellbook
43
43
  ```
44
44
 
45
45
  ## Configuration
@@ -66,6 +66,16 @@ databases:
66
66
  database: snowflake_database
67
67
  schema: snowflake_schema
68
68
 
69
+ # MySQL Database Connection using AWS RDS IAM authentication (no stored password).
70
+ # Fullscript internal only: requires the `rx` CLI. See documentation/db_magic.md.
71
+ - name: alias_of_your_connection
72
+ type: mysql
73
+ host: your_host.rds.amazonaws.com
74
+ port: 3306
75
+ user: ${DB_USER}
76
+ iam_auth: true
77
+ database: ${MYSQL_DATABASE}
78
+
69
79
  google_accounts:
70
80
  # Google Sheets Service Account 1
71
81
  - name: google_sheets_service_account_1
@@ -10,6 +10,69 @@ A magical wrapper to perform work with databases.
10
10
  from spellbook import db_magic
11
11
  ```
12
12
 
13
+ ### MySQL Authentication
14
+
15
+ MySQL connections support two authentication methods, chosen per database in `spellbook_config.yaml`.
16
+
17
+ **Password authentication** is the default and needs no changes to an existing config:
18
+
19
+ ```yaml
20
+ - name: legacy_mysql
21
+ type: mysql
22
+ host: your_host
23
+ port: 3306
24
+ user: mysql_user_name
25
+ password: ${MYSQL_PASSWORD}
26
+ database: mysql_database
27
+ ```
28
+
29
+ **AWS RDS IAM authentication** stores no password at all. Add `iam_auth: true` and Spellbook
30
+ calls `rx rds password <rds_db_name>` for a temporary token each time it connects, because the
31
+ tokens expire after about 15 minutes:
32
+
33
+ > **Internal use only.** `iam_auth` depends on the `rx` CLI, which is Fullscript internal
34
+ > tooling and is not publicly available. Password authentication is the only MySQL option
35
+ > outside Fullscript, and it remains fully supported.
36
+
37
+ ```yaml
38
+ - name: reporting_db
39
+ type: mysql
40
+ host: ${MYSQL_HOST} # e.g. your-instance.abc123.us-east-1.rds.amazonaws.com
41
+ port: 3306
42
+ user: ${DB_USER}
43
+ iam_auth: true
44
+ rds_db_name: your_rds_instance_name # Optional, defaults to the value of `name`
45
+ database: ${MYSQL_DATABASE}
46
+ ```
47
+
48
+ | Field | Required | Notes |
49
+ | --- | --- | --- |
50
+ | `iam_auth` | No | `true` switches the database to IAM tokens. Omitting it keeps password authentication. |
51
+ | `user` | Yes | Your IAM database user. |
52
+ | `rds_db_name` | No | The remote database as named by `rx rds list`. Defaults to `name`, so set it whenever your config alias differs from the RDS name. |
53
+ | `ssl_ca` | No | Not needed in normal use; see below. |
54
+ | `password` | No | Ignored when `iam_auth` is true. |
55
+
56
+ Requirements for the IAM path:
57
+ * The `rx` CLI must be installed and on your `PATH`.
58
+ * You must be signed in to AWS. If a connection hangs or fails on authentication, run
59
+ `rx rds password <rds_db_name>` in a terminal once to complete the AWS SSO prompt.
60
+ * IAM connections use the `pymysql` driver (installed with Spellbook). Password connections
61
+ continue to use `mysql-connector`.
62
+
63
+ ##### A note on TLS and `ssl_ca`
64
+
65
+ AWS requires TLS for IAM authentication, so Spellbook always negotiates an encrypted
66
+ connection. You do **not** need to set `ssl_ca` — leave it out and you get TLS without
67
+ verifying which server you reached, which is the normal setup.
68
+
69
+ Set `ssl_ca` only if you also want the server certificate verified. If you do:
70
+ * Point it at the [AWS RDS CA bundle](https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem).
71
+ Your operating system's bundle will not work, because RDS certificates do not chain to a
72
+ publicly trusted root; using it fails with `CERTIFICATE_VERIFY_FAILED`.
73
+ * Use the instance's own RDS endpoint as `host`. Verification also checks the hostname, so a
74
+ CNAME that points at the instance is likely to be rejected even with the right bundle.
75
+
13
76
  ### Methods Overview
14
77
 
15
78
 
@@ -95,4 +158,4 @@ update_query = "UPDATE users SET is_active = 1 WHERE last_login > '2024-01-01';"
95
158
  db_magic.run_sql("mysql_db", update_query)
96
159
 
97
160
  print("SQL executed successfully!")
98
- ```
161
+ ```
@@ -1,8 +1,8 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fullscript.spellbook
3
- Version: 2.1.6
3
+ Version: 2.2.0
4
4
  Summary: A simple package to make your life easier.
5
- Author-email: Thomas Edwards <thomas.edwards@fullscript.com>, Gui Chagas <gui.chagas@fullscript.com>, David King <david.king@fullscript.com>, Meinhard Mare <meinhard.mare@fullscript.com>, Karsten Melchers <karsten.melchers@fullscript.com>
5
+ Author-email: Meinhard Mare <meinhard.mare@fullscript.com>, Thomas Edwards <thomas.edwards@fullscript.com>, Gui Chagas <gui.chagas@fullscript.com>, David King <david.king@fullscript.com>, Karsten Melchers <karsten.melchers@fullscript.com>
6
6
  License: MIT
7
7
  Project-URL: homepage, https://github.com/Fullscript/spellbook
8
8
  Requires-Python: <3.12,>=3.10
@@ -46,6 +46,7 @@ Requires-Dist: pycparser==2.22
46
46
  Requires-Dist: Pygments==2.18.0
47
47
  Requires-Dist: PyJWT==2.9.0
48
48
  Requires-Dist: pyOpenSSL==24.2.1
49
+ Requires-Dist: PyMySQL==1.2.3
49
50
  Requires-Dist: python-dateutil==2.9.0.post0
50
51
  Requires-Dist: pytz==2024.2
51
52
  Requires-Dist: PyYAML==6.0.2
@@ -116,6 +117,8 @@ Requires-Dist: scipy==1.14.1; extra == "statsmagic"
116
117
  Requires-Dist: seaborn==0.13.2; extra == "statsmagic"
117
118
  Requires-Dist: six==1.16.0; extra == "statsmagic"
118
119
  Requires-Dist: statsmodels==0.14.4; extra == "statsmagic"
120
+ Provides-Extra: test
121
+ Requires-Dist: pytest==8.3.4; extra == "test"
119
122
  Dynamic: license-file
120
123
 
121
124
  # Spellbook
@@ -159,7 +162,7 @@ Spellbook requires Python 3.11 or higher. Use the following steps to install:
159
162
  1. Pip install:
160
163
 
161
164
  ```bash
162
- pip install git+https://git.fullscript.io/data/spellbook.git
165
+ pip install fullscript.spellbook
163
166
  ```
164
167
 
165
168
  ## Configuration
@@ -186,6 +189,16 @@ databases:
186
189
  database: snowflake_database
187
190
  schema: snowflake_schema
188
191
 
192
+ # MySQL Database Connection using AWS RDS IAM authentication (no stored password).
193
+ # Fullscript internal only: requires the `rx` CLI. See documentation/db_magic.md.
194
+ - name: alias_of_your_connection
195
+ type: mysql
196
+ host: your_host.rds.amazonaws.com
197
+ port: 3306
198
+ user: ${DB_USER}
199
+ iam_auth: true
200
+ database: ${MYSQL_DATABASE}
201
+
189
202
  google_accounts:
190
203
  # Google Sheets Service Account 1
191
204
  - name: google_sheets_service_account_1
@@ -1,4 +1,3 @@
1
- .DS_Store
2
1
  .gitignore
3
2
  .gitlab-ci.yml
4
3
  .python-version
@@ -17,8 +16,9 @@ fullscript.spellbook.egg-info/PKG-INFO
17
16
  fullscript.spellbook.egg-info/SOURCES.txt
18
17
  fullscript.spellbook.egg-info/dependency_links.txt
19
18
  fullscript.spellbook.egg-info/requires.txt
19
+ fullscript.spellbook.egg-info/scm_file_list.json
20
+ fullscript.spellbook.egg-info/scm_version.json
20
21
  fullscript.spellbook.egg-info/top_level.txt
21
- spellbook/.DS_Store
22
22
  spellbook/__init__.py
23
23
  spellbook/db_magic.py
24
24
  spellbook/drive_magic.py
@@ -37,4 +37,5 @@ spellbook/wizards/wizard5.txt
37
37
  spellbook/wizards/wizard6.txt
38
38
  spellbook/wizards/wizard7.txt
39
39
  spellbook/wizards/wizard8.txt
40
- spellbook/wizards/wizard9.txt
40
+ spellbook/wizards/wizard9.txt
41
+ tests/test_db_magic.py
@@ -36,6 +36,7 @@ pycparser==2.22
36
36
  Pygments==2.18.0
37
37
  PyJWT==2.9.0
38
38
  pyOpenSSL==24.2.1
39
+ PyMySQL==1.2.3
39
40
  python-dateutil==2.9.0.post0
40
41
  pytz==2024.2
41
42
  PyYAML==6.0.2
@@ -109,3 +110,6 @@ scipy==1.14.1
109
110
  seaborn==0.13.2
110
111
  six==1.16.0
111
112
  statsmodels==0.14.4
113
+
114
+ [test]
115
+ pytest==8.3.4
@@ -0,0 +1,38 @@
1
+ {
2
+ "files": [
3
+ ".gitignore",
4
+ ".gitlab-ci.yml",
5
+ ".python-version",
6
+ "LICENSE",
7
+ "MANIFEST.in",
8
+ "README.md",
9
+ "documentation/db_magic.md",
10
+ "documentation/drive_magic.md",
11
+ "documentation/sheet_magic.md",
12
+ "documentation/snowflake_key_pair_setup.md",
13
+ "documentation/stats_magic.md",
14
+ "pyproject.toml",
15
+ "requirements.txt",
16
+ "spellbook/__init__.py",
17
+ "spellbook/db_magic.py",
18
+ "spellbook/drive_magic.py",
19
+ "spellbook/sheet_magic.py",
20
+ "spellbook/stats_magic.py",
21
+ "spellbook/utils.py",
22
+ "spellbook/wizards/__init__.py",
23
+ "spellbook/wizards/wizard1.txt",
24
+ "spellbook/wizards/wizard10.txt",
25
+ "spellbook/wizards/wizard11.txt",
26
+ "spellbook/wizards/wizard12.txt",
27
+ "spellbook/wizards/wizard2.txt",
28
+ "spellbook/wizards/wizard3.txt",
29
+ "spellbook/wizards/wizard4.txt",
30
+ "spellbook/wizards/wizard5.txt",
31
+ "spellbook/wizards/wizard6.txt",
32
+ "spellbook/wizards/wizard7.txt",
33
+ "spellbook/wizards/wizard8.txt",
34
+ "spellbook/wizards/wizard9.txt",
35
+ "spellbook_config.yaml",
36
+ "tests/test_db_magic.py"
37
+ ]
38
+ }
@@ -0,0 +1,8 @@
1
+ {
2
+ "tag": "2.2.0",
3
+ "distance": 0,
4
+ "node": "gb79ed51cfa3c9db792124b375e8edb615250590b",
5
+ "dirty": false,
6
+ "branch": "HEAD",
7
+ "node_date": "2026-10-22"
8
+ }
@@ -10,10 +10,10 @@ readme = "README.md"
10
10
  license = {text = "MIT"}
11
11
  requires-python = ">=3.10, <3.12"
12
12
  authors = [
13
+ {name = "Meinhard Mare", email = "meinhard.mare@fullscript.com"},
13
14
  {name = "Thomas Edwards", email = "thomas.edwards@fullscript.com"},
14
15
  {name = "Gui Chagas", email = "gui.chagas@fullscript.com"},
15
16
  {name = "David King", email = "david.king@fullscript.com"},
16
- {name = "Meinhard Mare", email = "meinhard.mare@fullscript.com"},
17
17
  {name = "Karsten Melchers", email = "karsten.melchers@fullscript.com"}
18
18
  ]
19
19
  urls = {homepage = "https://github.com/Fullscript/spellbook"}
@@ -56,6 +56,7 @@ dependencies = [
56
56
  "Pygments==2.18.0",
57
57
  "PyJWT==2.9.0",
58
58
  "pyOpenSSL==24.2.1",
59
+ "PyMySQL==1.2.3",
59
60
  "python-dateutil==2.9.0.post0",
60
61
  "pytz==2024.2",
61
62
  "PyYAML==6.0.2",
@@ -133,6 +134,13 @@ statsmagic = [
133
134
  "statsmodels==0.14.4"
134
135
  ]
135
136
 
137
+ test = [
138
+ "pytest==8.3.4"
139
+ ]
140
+
141
+ [tool.pytest.ini_options]
142
+ testpaths = ["tests"]
143
+
136
144
  [tool.setuptools.packages.find]
137
145
  where = ["."]
138
146
  include = ["spellbook"]
@@ -183,14 +183,14 @@ six==1.16.0
183
183
  # asttokens
184
184
  # fullscript.spellbook (pyproject.toml)
185
185
  # python-dateutil
186
- snowflake-connector-python==3.15.0
186
+ snowflake-connector-python==4.2.0
187
187
  # via
188
188
  # fullscript.spellbook (pyproject.toml)
189
189
  # snowflake-snowpark-python
190
190
  # snowflake-sqlalchemy
191
- snowflake-snowpark-python==1.25.0
191
+ snowflake-snowpark-python==1.44.0
192
192
  # via fullscript.spellbook (pyproject.toml)
193
- snowflake-sqlalchemy==1.7.0
193
+ snowflake-sqlalchemy==1.8.2
194
194
  # via fullscript.spellbook (pyproject.toml)
195
195
  sortedcontainers==2.4.0
196
196
  # via
@@ -1,16 +1,38 @@
1
1
  import os
2
+ import re
3
+ import subprocess
2
4
  from snowflake.snowpark import Session
3
5
  from sqlalchemy import create_engine
4
6
  import pandas as pd
5
7
  from cryptography.hazmat.backends import default_backend
6
8
  from cryptography.hazmat.primitives import serialization
7
9
  import base64
10
+ from urllib.parse import quote_plus
8
11
  from spellbook.utils import load_and_parse_config, get_wizard, load_sql_file
9
12
 
10
13
 
11
14
  # Needs a very specific pip install
12
15
  # pip install "snowflake-connector-python[pandas]"
13
16
 
17
+ # How long to wait for `rx` to mint an RDS IAM token. The first call of the day can
18
+ # be slow because `rx` has to walk through AWS SSO before it reaches RDS.
19
+ RDS_TOKEN_TIMEOUT = 120
20
+
21
+ # `rx` colourises its progress output, which has to come off before we can read the token.
22
+ _ANSI_ESCAPE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]')
23
+
24
+ # Built on first use so that pymysql is only imported when an IAM connection is made.
25
+ _IAM_CURSOR_CLASS = None
26
+
27
+
28
+ def _strip_ansi(text):
29
+ """
30
+ Remove ANSI colour and cursor codes from CLI output.
31
+ :param text: The raw text written by the CLI.
32
+ :return: The same text with escape codes removed.
33
+ """
34
+ return _ANSI_ESCAPE.sub('', text or '')
35
+
14
36
 
15
37
  def get_db_config(db_name):
16
38
  # load the configuration
@@ -82,18 +104,126 @@ def snowflake_connect(db_name, **kwargs):
82
104
  return None
83
105
 
84
106
 
107
+ def _iam_cursor_class():
108
+ """
109
+ Build the pymysql cursor used for IAM connections.
110
+
111
+ pymysql substitutes parameters whenever `args is not None`, so the empty parameter set
112
+ pandas passes turns a `like '%value%'` clause into a format string and raises
113
+ "unsupported format character". mysql-connector guards on `if params:` instead, which
114
+ leaves such queries alone, so this keeps both drivers behaving the same way.
115
+
116
+ :return: A pymysql cursor class that treats an empty parameter set as no parameters.
117
+ """
118
+ global _IAM_CURSOR_CLASS
119
+ if _IAM_CURSOR_CLASS is None:
120
+ import pymysql.cursors
121
+
122
+ class LiteralPercentCursor(pymysql.cursors.Cursor):
123
+ def execute(self, query, args=None):
124
+ return super().execute(query, args if args else None)
125
+
126
+ _IAM_CURSOR_CLASS = LiteralPercentCursor
127
+ return _IAM_CURSOR_CLASS
128
+
129
+
130
+ def get_rds_iam_token(rds_db_name, timeout=RDS_TOKEN_TIMEOUT):
131
+ """
132
+ Generate a temporary AWS RDS IAM authentication token using the `rx` CLI.
133
+
134
+ The token stands in for a stored password and expires after about 15 minutes, so it
135
+ has to be fetched fresh for every connection rather than cached in the config file.
136
+
137
+ :param rds_db_name: The remote database name, as listed by `rx rds list`.
138
+ :param timeout: Seconds to wait for `rx` before giving up.
139
+ :return: The token to use as the connection password.
140
+ :raises RuntimeError: If `rx` is unavailable, times out, fails, or returns no token.
141
+ """
142
+ command = ['rx', 'rds', 'password', rds_db_name]
143
+ printable = ' '.join(command)
144
+
145
+ try:
146
+ result = subprocess.run(command, capture_output=True, text=True, timeout=timeout)
147
+ except FileNotFoundError:
148
+ raise RuntimeError(
149
+ "AWS RDS IAM authentication failed: the 'rx' command is not installed or not on "
150
+ "your PATH. Install the Fullscript developer tooling, or drop 'iam_auth' from the "
151
+ "database configuration to go back to password authentication."
152
+ ) from None
153
+ except subprocess.TimeoutExpired:
154
+ raise RuntimeError(
155
+ f"AWS RDS IAM authentication failed: '{printable}' did not finish within {timeout} "
156
+ f"seconds. It is probably waiting on AWS SSO - run '{printable}' in a terminal once "
157
+ "to sign in, then try again."
158
+ ) from None
159
+
160
+ if result.returncode != 0:
161
+ detail = _strip_ansi(result.stderr or result.stdout or '').strip()
162
+ raise RuntimeError(
163
+ f"AWS RDS IAM authentication failed: '{printable}' exited with status "
164
+ f"{result.returncode}. Check that '{rds_db_name}' appears in 'rx rds list' and that "
165
+ f"you are signed in to AWS. Output from rx: {detail or '<none>'}"
166
+ )
167
+
168
+ # `rx` writes progress lines before the token, so the token is the last thing it prints.
169
+ lines = [line.strip() for line in _strip_ansi(result.stdout).splitlines() if line.strip()]
170
+ if not lines:
171
+ raise RuntimeError(
172
+ f"AWS RDS IAM authentication failed: '{printable}' succeeded but returned no token."
173
+ )
174
+
175
+ return lines[-1]
176
+
177
+
85
178
  def mysql_connect(db_name, **kwargs):
179
+ """
180
+ Create a MySQL connection using either AWS RDS IAM authentication or a stored password.
181
+
182
+ Set `iam_auth: true` on the database in the config file to authenticate with a temporary
183
+ IAM token from `rx rds password`. Without it, the `password` field is used as before.
184
+
185
+ :param db_name: The name of the database as defined in the config file.
186
+ :return: SQLAlchemy connection object, or None if the connection could not be made.
187
+ """
86
188
  db_config = get_db_config(db_name, **kwargs)
87
189
  try:
88
- connection_string = (
89
- f"mysql+mysqlconnector://{db_config['user']}:{db_config['password']}"
90
- f"@{db_config['host']}:{db_config.get('port', 3306)}/{db_config['database']}"
91
- )
92
- engine = create_engine(connection_string)
190
+ host = db_config['host']
191
+ port = db_config.get('port', 3306)
192
+ database = db_config['database']
193
+
194
+ if db_config.get('iam_auth'):
195
+ # RDS instance names rarely match the alias used in the config file, so allow an
196
+ # override and fall back to the alias when the two happen to line up.
197
+ token = get_rds_iam_token(db_config.get('rds_db_name', db_config['name']))
198
+ # IAM tokens are URL shaped, so they have to be escaped inside the connection
199
+ # string. AWS also requires TLS, and only pymysql completes the token handshake.
200
+ ssl_ca = db_config.get('ssl_ca')
201
+ connection_string = (
202
+ f"mysql+pymysql://{quote_plus(db_config['user'])}:{quote_plus(token)}"
203
+ f"@{host}:{port}/{database}"
204
+ )
205
+ # A CA bundle gets the certificate verified too; without one we still insist on
206
+ # an encrypted connection, we just cannot prove who is on the other end.
207
+ connect_args = {
208
+ 'ssl': {'ca': ssl_ca} if ssl_ca else {'check_hostname': False},
209
+ 'cursorclass': _iam_cursor_class(),
210
+ }
211
+ else:
212
+ connection_string = (
213
+ f"mysql+mysqlconnector://{db_config['user']}:{db_config['password']}"
214
+ f"@{host}:{port}/{database}"
215
+ )
216
+ connect_args = {}
217
+
218
+ engine = create_engine(connection_string, connect_args=connect_args)
93
219
  con = engine.connect()
94
220
  return con
95
221
  except Exception as e:
96
- print(f"An error occurred creating the MySQL connection: {e}")
222
+ auth_method = 'AWS RDS IAM' if db_config.get('iam_auth') else 'password'
223
+ print(
224
+ f"An error occurred creating the MySQL connection to '{db_name}' "
225
+ f"using {auth_method} authentication: {e}"
226
+ )
97
227
  return None
98
228
 
99
229
 
@@ -33,6 +33,17 @@ databases:
33
33
  password: mysql_password
34
34
  database: mysql_database
35
35
 
36
+ # MySQL Database Connection using AWS RDS IAM authentication.
37
+ # Requires the `rx` CLI, which is Fullscript internal tooling.
38
+ - name: alias_of_your_connection_3
39
+ type: mysql
40
+ host: ${MYSQL_HOST} # e.g. your-instance.abc123.us-east-1.rds.amazonaws.com
41
+ port: 3306
42
+ user: ${DB_USER} # Your IAM database user
43
+ iam_auth: true # Use an AWS RDS IAM token instead of a password
44
+ rds_db_name: your_rds_instance_name # Optional, as named by `rx rds list`; defaults to `name`
45
+ database: ${MYSQL_DATABASE}
46
+
36
47
  google_accounts:
37
48
  # Google Sheets Service Account 1
38
49
  - name: google_sheets_service_account_1
@@ -0,0 +1,316 @@
1
+ import subprocess
2
+
3
+ import pytest
4
+
5
+ from spellbook import db_magic
6
+
7
+
8
+ PASSWORD_DB = {
9
+ 'name': 'legacy_mysql',
10
+ 'type': 'mysql',
11
+ 'host': 'legacy.example.com',
12
+ 'port': 3306,
13
+ 'user': 'mysql_user_name',
14
+ 'password': 'mysql_password',
15
+ 'database': 'mysql_database',
16
+ }
17
+
18
+ IAM_DB = {
19
+ 'name': 'reporting_db',
20
+ 'type': 'mysql',
21
+ 'host': 'example-db.abc123.us-east-1.rds.amazonaws.com',
22
+ 'port': 3306,
23
+ 'user': 'iam_user',
24
+ 'iam_auth': True,
25
+ 'rds_db_name': 'example-rds-instance',
26
+ 'database': 'mysql_database',
27
+ }
28
+
29
+ # Real IAM tokens are URL shaped: the query string has to survive the trip into the
30
+ # connection string without being mistaken for part of the URL.
31
+ FAKE_TOKEN = (
32
+ 'example-db.abc123.us-east-1.rds.amazonaws.com:3306/'
33
+ '?Action=connect&DBUser=iam_user&X-Amz-Signature=abc123'
34
+ )
35
+
36
+
37
+ @pytest.fixture
38
+ def engines(monkeypatch):
39
+ """Capture the arguments mysql_connect passes to create_engine."""
40
+ calls = []
41
+
42
+ class FakeEngine:
43
+ def connect(self):
44
+ return 'connection'
45
+
46
+ def fake_create_engine(connection_string, **kwargs):
47
+ calls.append((connection_string, kwargs))
48
+ return FakeEngine()
49
+
50
+ monkeypatch.setattr(db_magic, 'create_engine', fake_create_engine)
51
+ return calls
52
+
53
+
54
+ @pytest.fixture
55
+ def config(monkeypatch):
56
+ """Point get_db_config at an in-test database list."""
57
+ def use(*databases):
58
+ by_name = {db['name']: db for db in databases}
59
+ monkeypatch.setattr(db_magic, 'get_db_config', lambda name, **kwargs: by_name[name])
60
+ return use
61
+
62
+
63
+ def fake_rx(stdout='', stderr='', returncode=0, expect=None, raises=None):
64
+ """Build a stand-in for subprocess.run that mimics one `rx rds password` outcome."""
65
+ def run(command, **kwargs):
66
+ if expect is not None:
67
+ assert command == expect
68
+ if raises is not None:
69
+ raise raises
70
+ return subprocess.CompletedProcess(command, returncode, stdout=stdout, stderr=stderr)
71
+ return run
72
+
73
+
74
+ # --- Password authentication stays exactly as it was -------------------------------------
75
+
76
+ def test_password_auth_uses_mysqlconnector(monkeypatch, config, engines):
77
+ config(PASSWORD_DB)
78
+ # Any attempt to shell out to rx on the password path is a backward compatibility bug.
79
+ monkeypatch.setattr(subprocess, 'run', fake_rx(raises=AssertionError('rx must not run')))
80
+
81
+ con = db_magic.mysql_connect('legacy_mysql')
82
+
83
+ assert con == 'connection'
84
+ connection_string, kwargs = engines[0]
85
+ assert connection_string == (
86
+ 'mysql+mysqlconnector://mysql_user_name:mysql_password'
87
+ '@legacy.example.com:3306/mysql_database'
88
+ )
89
+ assert kwargs == {'connect_args': {}}
90
+
91
+
92
+ def test_password_auth_defaults_the_port(config, engines):
93
+ config({key: value for key, value in PASSWORD_DB.items() if key != 'port'})
94
+
95
+ db_magic.mysql_connect('legacy_mysql')
96
+
97
+ assert ':3306/' in engines[0][0]
98
+
99
+
100
+ def test_iam_auth_false_behaves_like_password_auth(config, engines):
101
+ config({**PASSWORD_DB, 'iam_auth': False})
102
+
103
+ db_magic.mysql_connect('legacy_mysql')
104
+
105
+ assert engines[0][0].startswith('mysql+mysqlconnector://')
106
+
107
+
108
+ # --- IAM authentication ------------------------------------------------------------------
109
+
110
+ def test_iam_auth_builds_a_pymysql_connection_string(monkeypatch, config, engines):
111
+ config(IAM_DB)
112
+ monkeypatch.setattr(subprocess, 'run', fake_rx(
113
+ stdout=FAKE_TOKEN + '\n',
114
+ expect=['rx', 'rds', 'password', 'example-rds-instance'],
115
+ ))
116
+
117
+ con = db_magic.mysql_connect('reporting_db')
118
+
119
+ assert con == 'connection'
120
+ connection_string, kwargs = engines[0]
121
+ assert connection_string.startswith('mysql+pymysql://iam_user:')
122
+ assert connection_string.endswith(
123
+ '@example-db.abc123.us-east-1.rds.amazonaws.com:3306/mysql_database'
124
+ )
125
+ # The token's own '?' and '&' must be escaped, or SQLAlchemy reads them as URL structure.
126
+ assert '?Action=connect' not in connection_string
127
+ assert '%3FAction%3Dconnect' in connection_string
128
+ # TLS is mandatory for RDS IAM, so the ssl options must be truthy for pymysql to require it.
129
+ assert kwargs['connect_args']['ssl']
130
+
131
+
132
+ def test_iam_auth_connection_string_round_trips_the_token(monkeypatch, config, engines):
133
+ from sqlalchemy.engine import make_url
134
+
135
+ config(IAM_DB)
136
+ monkeypatch.setattr(subprocess, 'run', fake_rx(stdout=FAKE_TOKEN))
137
+
138
+ db_magic.mysql_connect('reporting_db')
139
+
140
+ url = make_url(engines[0][0])
141
+ assert url.password == FAKE_TOKEN
142
+ assert url.username == 'iam_user'
143
+ assert url.database == 'mysql_database'
144
+
145
+
146
+ def test_iam_auth_falls_back_to_the_config_alias_for_the_rds_name(monkeypatch, config, engines):
147
+ without_override = {key: value for key, value in IAM_DB.items() if key != 'rds_db_name'}
148
+ config(without_override)
149
+ monkeypatch.setattr(subprocess, 'run', fake_rx(
150
+ stdout=FAKE_TOKEN,
151
+ expect=['rx', 'rds', 'password', 'reporting_db'],
152
+ ))
153
+
154
+ assert db_magic.mysql_connect('reporting_db') == 'connection'
155
+
156
+
157
+ def test_iam_auth_verifies_the_certificate_when_a_ca_bundle_is_given(monkeypatch, config, engines):
158
+ config({**IAM_DB, 'ssl_ca': '/certs/rds-ca.pem'})
159
+ monkeypatch.setattr(subprocess, 'run', fake_rx(stdout=FAKE_TOKEN))
160
+
161
+ db_magic.mysql_connect('reporting_db')
162
+
163
+ assert engines[0][1]['connect_args']['ssl'] == {'ca': '/certs/rds-ca.pem'}
164
+
165
+
166
+ def test_iam_auth_requires_tls_without_a_ca_bundle(monkeypatch, config, engines):
167
+ config(IAM_DB)
168
+ monkeypatch.setattr(subprocess, 'run', fake_rx(stdout=FAKE_TOKEN))
169
+
170
+ db_magic.mysql_connect('reporting_db')
171
+
172
+ assert engines[0][1]['connect_args']['ssl'] == {'check_hostname': False}
173
+
174
+
175
+ def test_iam_auth_fetches_a_fresh_token_per_connection(monkeypatch, config, engines):
176
+ config(IAM_DB)
177
+ tokens = iter(['token-one', 'token-two'])
178
+ monkeypatch.setattr(
179
+ subprocess, 'run',
180
+ lambda command, **kwargs: subprocess.CompletedProcess(
181
+ command, 0, stdout=next(tokens), stderr=''
182
+ ),
183
+ )
184
+
185
+ db_magic.mysql_connect('reporting_db')
186
+ db_magic.mysql_connect('reporting_db')
187
+
188
+ assert 'token-one' in engines[0][0]
189
+ assert 'token-two' in engines[1][0]
190
+
191
+
192
+ # --- Literal '%' must survive on the pymysql path ----------------------------------------
193
+ # pymysql substitutes whenever args is not None, so the empty parameter set pandas passes
194
+ # would turn `like '%x%'` into a format string. mysql-connector guards on `if params:`,
195
+ # so without this the IAM path would break queries the password path handles fine.
196
+
197
+ def test_iam_auth_uses_a_cursor_that_leaves_literal_percent_alone(monkeypatch, config, engines):
198
+ config(IAM_DB)
199
+ monkeypatch.setattr(subprocess, 'run', fake_rx(stdout=FAKE_TOKEN))
200
+
201
+ db_magic.mysql_connect('reporting_db')
202
+
203
+ cursor_class = engines[0][1]['connect_args']['cursorclass']
204
+ import pymysql.cursors
205
+ assert issubclass(cursor_class, pymysql.cursors.Cursor)
206
+
207
+
208
+ @pytest.fixture
209
+ def iam_cursor(monkeypatch):
210
+ """An IAM cursor whose pymysql parent records the args it is handed."""
211
+ import pymysql.cursors
212
+
213
+ recorded = []
214
+ monkeypatch.setattr(
215
+ pymysql.cursors.Cursor, 'execute',
216
+ lambda self, query, args=None: recorded.append(args),
217
+ )
218
+ return db_magic._iam_cursor_class()(None), recorded
219
+
220
+
221
+ @pytest.mark.parametrize('empty', [(), [], {}])
222
+ def test_empty_parameters_reach_pymysql_as_none(iam_cursor, empty):
223
+ """An empty parameter set must arrive as None, or '%' gets read as a format character."""
224
+ cursor, recorded = iam_cursor
225
+
226
+ cursor.execute("select 1 from t where name like '%x%'", empty)
227
+
228
+ assert recorded == [None]
229
+
230
+
231
+ def test_omitted_parameters_reach_pymysql_as_none(iam_cursor):
232
+ cursor, recorded = iam_cursor
233
+
234
+ cursor.execute("select 1 from t where name like '%x%'")
235
+
236
+ assert recorded == [None]
237
+
238
+
239
+ def test_real_parameters_are_still_forwarded(iam_cursor):
240
+ cursor, recorded = iam_cursor
241
+
242
+ cursor.execute('select 1 from t where id = %s', (42,))
243
+
244
+ assert recorded == [(42,)]
245
+
246
+
247
+ # --- Token retrieval ---------------------------------------------------------------------
248
+
249
+ def test_token_is_read_from_the_last_line_of_noisy_output(monkeypatch):
250
+ monkeypatch.setattr(subprocess, 'run', fake_rx(
251
+ stdout='\x1b[32m⠿ Waiting for AWS authentication\x1b[0m\n\n' + FAKE_TOKEN + '\n',
252
+ ))
253
+
254
+ assert db_magic.get_rds_iam_token('example-rds-instance') == FAKE_TOKEN
255
+
256
+
257
+ def test_missing_rx_raises_a_clear_error(monkeypatch):
258
+ monkeypatch.setattr(subprocess, 'run', fake_rx(raises=FileNotFoundError('rx')))
259
+
260
+ with pytest.raises(RuntimeError, match="'rx' command is not installed"):
261
+ db_magic.get_rds_iam_token('example-rds-instance')
262
+
263
+
264
+ def test_rx_timeout_raises_a_clear_error(monkeypatch):
265
+ monkeypatch.setattr(subprocess, 'run', fake_rx(
266
+ raises=subprocess.TimeoutExpired(cmd='rx', timeout=120),
267
+ ))
268
+
269
+ with pytest.raises(RuntimeError, match='did not finish within'):
270
+ db_magic.get_rds_iam_token('example-rds-instance')
271
+
272
+
273
+ def test_rx_failure_surfaces_its_output(monkeypatch):
274
+ monkeypatch.setattr(subprocess, 'run', fake_rx(
275
+ returncode=1, stderr='\x1b[;31m✘ database not found\x1b[;m',
276
+ ))
277
+
278
+ with pytest.raises(RuntimeError) as error:
279
+ db_magic.get_rds_iam_token('typo-db')
280
+
281
+ message = str(error.value)
282
+ assert 'exited with status 1' in message
283
+ assert '✘ database not found' in message
284
+ assert '\x1b' not in message # colour codes stripped so the message stays readable
285
+ assert 'typo-db' in message
286
+
287
+
288
+ def test_empty_rx_output_raises_rather_than_connecting_with_a_blank_password(monkeypatch):
289
+ monkeypatch.setattr(subprocess, 'run', fake_rx(stdout=' \n\n'))
290
+
291
+ with pytest.raises(RuntimeError, match='returned no token'):
292
+ db_magic.get_rds_iam_token('example-rds-instance')
293
+
294
+
295
+ def test_iam_failure_reports_which_auth_method_failed(monkeypatch, config, capsys):
296
+ config(IAM_DB)
297
+ monkeypatch.setattr(subprocess, 'run', fake_rx(raises=FileNotFoundError('rx')))
298
+
299
+ assert db_magic.mysql_connect('reporting_db') is None
300
+
301
+ output = capsys.readouterr().out
302
+ assert 'AWS RDS IAM authentication' in output
303
+ assert 'reporting_db' in output
304
+
305
+
306
+ def test_password_failure_reports_which_auth_method_failed(monkeypatch, config, capsys):
307
+ config(PASSWORD_DB)
308
+ monkeypatch.setattr(db_magic, 'create_engine', lambda *a, **k: (_ for _ in ()).throw(
309
+ RuntimeError('driver missing')
310
+ ))
311
+
312
+ assert db_magic.mysql_connect('legacy_mysql') is None
313
+
314
+ output = capsys.readouterr().out
315
+ assert 'password authentication' in output
316
+ assert 'driver missing' in output
Binary file
@@ -1,3 +0,0 @@
1
- *.whl
2
- *.egg-info
3
- *.tar.gz