al-data-utils 0.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 (65) hide show
  1. al_data_utils-0.3.1/LICENSE +21 -0
  2. al_data_utils-0.3.1/PKG-INFO +157 -0
  3. al_data_utils-0.3.1/README.md +141 -0
  4. al_data_utils-0.3.1/pyproject.toml +58 -0
  5. al_data_utils-0.3.1/setup.cfg +4 -0
  6. al_data_utils-0.3.1/src/al_data_utils/__init__.py +10 -0
  7. al_data_utils-0.3.1/src/al_data_utils/airflow_utils/__init__.py +13 -0
  8. al_data_utils-0.3.1/src/al_data_utils/airflow_utils/airflow_utils.py +83 -0
  9. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/__init__.py +31 -0
  10. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/auth.py +82 -0
  11. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/azure_as_utils.py +223 -0
  12. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/client.py +158 -0
  13. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/errors.py +78 -0
  14. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/models.py +265 -0
  15. al_data_utils-0.3.1/src/al_data_utils/azure_as_utils/poller.py +121 -0
  16. al_data_utils-0.3.1/src/al_data_utils/backup_utils/__init__.py +9 -0
  17. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/dacpac/__init__.py +1 -0
  18. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/dacpac/sqlpackage.py +207 -0
  19. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/__init__.py +8 -0
  20. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/config.py +150 -0
  21. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/emit/__init__.py +11 -0
  22. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/emit/foreign_keys.py +38 -0
  23. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/emit/indexes.py +37 -0
  24. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/emit/sql_objects.py +17 -0
  25. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/emit/tables.py +139 -0
  26. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/__init__.py +19 -0
  27. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/foreign_keys.py +173 -0
  28. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/functions.py +133 -0
  29. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/indexes.py +190 -0
  30. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/sequences.py +180 -0
  31. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/stored_procedures.py +313 -0
  32. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/synonyms.py +135 -0
  33. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/tables.py +251 -0
  34. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/introspect/views.py +244 -0
  35. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/io/__init__.py +10 -0
  36. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/io/layout.py +28 -0
  37. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/io/writer.py +40 -0
  38. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/models/__init__.py +14 -0
  39. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/models/foreign_key_def.py +19 -0
  40. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/models/index_def.py +19 -0
  41. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/models/sql_object_def.py +11 -0
  42. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/models/table_column_def.py +27 -0
  43. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/models/table_def.py +22 -0
  44. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/schema_snapshot/snapshot_processor.py +754 -0
  45. al_data_utils-0.3.1/src/al_data_utils/backup_utils/mssql/scripter/__init__.py +1 -0
  46. al_data_utils-0.3.1/src/al_data_utils/config_utils/__init__.py +6 -0
  47. al_data_utils-0.3.1/src/al_data_utils/config_utils/config_utils.py +31 -0
  48. al_data_utils-0.3.1/src/al_data_utils/files_sync_utils/__init__.py +15 -0
  49. al_data_utils-0.3.1/src/al_data_utils/files_sync_utils/files_sync_utils.py +567 -0
  50. al_data_utils-0.3.1/src/al_data_utils/geo_utils/__init__.py +0 -0
  51. al_data_utils-0.3.1/src/al_data_utils/git_utils/__init__.py +27 -0
  52. al_data_utils-0.3.1/src/al_data_utils/git_utils/git_utils.py +142 -0
  53. al_data_utils-0.3.1/src/al_data_utils/logging_utils/__init__.py +3 -0
  54. al_data_utils-0.3.1/src/al_data_utils/logging_utils/logging_utils.py +39 -0
  55. al_data_utils-0.3.1/src/al_data_utils/ones_api_utils/__init__.py +3 -0
  56. al_data_utils-0.3.1/src/al_data_utils/ones_api_utils/ones_api_utils.py +61 -0
  57. al_data_utils-0.3.1/src/al_data_utils/sql_utils/__init__.py +3 -0
  58. al_data_utils-0.3.1/src/al_data_utils/sql_utils/sql_utils.py +158 -0
  59. al_data_utils-0.3.1/src/al_data_utils/tg_utils/__init__.py +3 -0
  60. al_data_utils-0.3.1/src/al_data_utils/tg_utils/tg_utils.py +45 -0
  61. al_data_utils-0.3.1/src/al_data_utils.egg-info/PKG-INFO +157 -0
  62. al_data_utils-0.3.1/src/al_data_utils.egg-info/SOURCES.txt +63 -0
  63. al_data_utils-0.3.1/src/al_data_utils.egg-info/dependency_links.txt +1 -0
  64. al_data_utils-0.3.1/src/al_data_utils.egg-info/top_level.txt +1 -0
  65. al_data_utils-0.3.1/tests/test_azure_as_utils.py +677 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 AL Data Labs
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,157 @@
1
+ Metadata-Version: 2.4
2
+ Name: al-data-utils
3
+ Version: 0.3.1
4
+ Summary: Reusable utilities for data engineering (Airflow, Azure, Snowflake, SQL, config).
5
+ Author-email: AL Data Labs <pypi@aldatalabs.com>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://aldatalabs.com
8
+ Keywords: airflow,data-engineering,azure,snowflake,sql
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Operating System :: OS Independent
11
+ Classifier: Topic :: Software Development :: Libraries
12
+ Requires-Python: >=3.10
13
+ Description-Content-Type: text/markdown
14
+ License-File: LICENSE
15
+ Dynamic: license-file
16
+
17
+ # al-data-utils
18
+
19
+ Reusable Python utilities for data and automation workflows.
20
+
21
+ This repository contains independent utility modules for sync operations, git automation, SQL work, logging helpers, and MSSQL schema snapshot generation.
22
+
23
+ ## Python and Setup
24
+
25
+ - Python: 3.10+
26
+ - Packaging: setuptools (src layout)
27
+
28
+ Install locally in editable mode:
29
+
30
+ ```bash
31
+ pip install -e .
32
+ ```
33
+
34
+ ## Project Structure
35
+
36
+ - src/al_data_utils
37
+ - examples
38
+ - development_guides
39
+ - tools
40
+
41
+ Main utility groups under src/al_data_utils:
42
+
43
+ - airflow_utils
44
+ - backup_utils
45
+ - config_utils
46
+ - files_sync_utils
47
+ - geo_utils
48
+ - git_utils
49
+ - logging_utils
50
+ - ones_api_utils
51
+ - sql_utils
52
+ - tg_utils
53
+
54
+ ## Examples
55
+
56
+ Runnable examples are available in the `examples` folder for the main utility groups.
57
+ Use them as reference implementations for local scripts and automation jobs.
58
+
59
+ ## MSSQL Schema Snapshot
60
+
61
+ The snapshot tool exports SQL objects from MSSQL into a filesystem layout and syncs staging output to destination folders.
62
+
63
+ Current behavior highlights:
64
+
65
+ - Supports single DB and multi-DB runs in one execution.
66
+ - Always writes outputs into per-DB folders.
67
+ - Uses staging first, then syncs to final output.
68
+ - Supports cleanup of staging folders.
69
+ - Supports sync log retention cleanup.
70
+ - Supports include and exclude filters per object type.
71
+ - Supports DataGrip template override files.
72
+
73
+ Output shape:
74
+
75
+ ```text
76
+ <MSSQL_OUT_DIR>/<DatabaseName>/schemas/<schema>/
77
+ tables/
78
+ views/
79
+ routines/
80
+ sequences/
81
+ synonyms/
82
+ ```
83
+
84
+ ### Snapshot Environment Variables
85
+
86
+ Core connection and output:
87
+
88
+ - MSSQL_OUT_DIR
89
+ - MSSQL_SERVER
90
+ - MSSQL_DATABASE
91
+ - MSSQL_DATABASES
92
+ - MSSQL_AUTH_MODE
93
+ - MSSQL_USER
94
+ - MSSQL_PASSWORD
95
+ - MSSQL_DRIVER
96
+
97
+ Snapshot behavior:
98
+
99
+ - MSSQL_ENABLE_TABLE_DEFINITIONS
100
+ - MSSQL_ENABLE_INDEXES
101
+ - MSSQL_ENABLE_FOREIGN_KEYS
102
+ - MSSQL_ENABLE_VIEWS
103
+ - MSSQL_ENABLE_STORED_PROCEDURES
104
+ - MSSQL_ENABLE_FUNCTIONS
105
+ - MSSQL_ENABLE_SEQUENCES
106
+ - MSSQL_ENABLE_SYNONYMS
107
+ - MSSQL_CLEANUP_STAGING
108
+ - MSSQL_SYNC_LOG_RETENTION_DAYS
109
+
110
+ Filters:
111
+
112
+ - MSSQL_INCLUDE_SCHEMAS
113
+ - MSSQL_EXCLUDE_SCHEMAS
114
+ - MSSQL_INCLUDE_TABLES
115
+ - MSSQL_EXCLUDE_TABLES
116
+ - MSSQL_INCLUDE_VIEWS
117
+ - MSSQL_EXCLUDE_VIEWS
118
+ - MSSQL_INCLUDE_STORED_PROCEDURES
119
+ - MSSQL_EXCLUDE_STORED_PROCEDURES
120
+ - MSSQL_INCLUDE_FUNCTIONS
121
+ - MSSQL_EXCLUDE_FUNCTIONS
122
+ - MSSQL_INCLUDE_SEQUENCES
123
+ - MSSQL_EXCLUDE_SEQUENCES
124
+ - MSSQL_INCLUDE_SYNONYMS
125
+ - MSSQL_EXCLUDE_SYNONYMS
126
+
127
+ Notes:
128
+
129
+ - If MSSQL_DATABASES is set, it is treated as a comma-separated list.
130
+ - If MSSQL_DATABASES is empty, MSSQL_DATABASE is used.
131
+ - If MSSQL_DATABASE contains commas, it is also treated as a list for backward compatibility.
132
+
133
+ ## File Sync Utilities
134
+
135
+ `files_sync_utils` provides:
136
+
137
+ - Inventory building
138
+ - Plan generation
139
+ - Copy and update execution
140
+ - Extra destination cleanup rules
141
+ - Activity CSV logging
142
+ - Log retention cleanup
143
+
144
+ ## Git Utilities
145
+
146
+ `git_utils` wraps common git automation operations with consistent behavior and error handling.
147
+
148
+ ## Development Guides
149
+
150
+ - development_guides/dev_guide.txt
151
+ - development_guides/auth_key_guide.txt
152
+
153
+ These guides contain release flow notes and auth-related instructions.
154
+
155
+ ## Changelog
156
+
157
+ See CHANGELOG.md for release notes.
@@ -0,0 +1,141 @@
1
+ # al-data-utils
2
+
3
+ Reusable Python utilities for data and automation workflows.
4
+
5
+ This repository contains independent utility modules for sync operations, git automation, SQL work, logging helpers, and MSSQL schema snapshot generation.
6
+
7
+ ## Python and Setup
8
+
9
+ - Python: 3.10+
10
+ - Packaging: setuptools (src layout)
11
+
12
+ Install locally in editable mode:
13
+
14
+ ```bash
15
+ pip install -e .
16
+ ```
17
+
18
+ ## Project Structure
19
+
20
+ - src/al_data_utils
21
+ - examples
22
+ - development_guides
23
+ - tools
24
+
25
+ Main utility groups under src/al_data_utils:
26
+
27
+ - airflow_utils
28
+ - backup_utils
29
+ - config_utils
30
+ - files_sync_utils
31
+ - geo_utils
32
+ - git_utils
33
+ - logging_utils
34
+ - ones_api_utils
35
+ - sql_utils
36
+ - tg_utils
37
+
38
+ ## Examples
39
+
40
+ Runnable examples are available in the `examples` folder for the main utility groups.
41
+ Use them as reference implementations for local scripts and automation jobs.
42
+
43
+ ## MSSQL Schema Snapshot
44
+
45
+ The snapshot tool exports SQL objects from MSSQL into a filesystem layout and syncs staging output to destination folders.
46
+
47
+ Current behavior highlights:
48
+
49
+ - Supports single DB and multi-DB runs in one execution.
50
+ - Always writes outputs into per-DB folders.
51
+ - Uses staging first, then syncs to final output.
52
+ - Supports cleanup of staging folders.
53
+ - Supports sync log retention cleanup.
54
+ - Supports include and exclude filters per object type.
55
+ - Supports DataGrip template override files.
56
+
57
+ Output shape:
58
+
59
+ ```text
60
+ <MSSQL_OUT_DIR>/<DatabaseName>/schemas/<schema>/
61
+ tables/
62
+ views/
63
+ routines/
64
+ sequences/
65
+ synonyms/
66
+ ```
67
+
68
+ ### Snapshot Environment Variables
69
+
70
+ Core connection and output:
71
+
72
+ - MSSQL_OUT_DIR
73
+ - MSSQL_SERVER
74
+ - MSSQL_DATABASE
75
+ - MSSQL_DATABASES
76
+ - MSSQL_AUTH_MODE
77
+ - MSSQL_USER
78
+ - MSSQL_PASSWORD
79
+ - MSSQL_DRIVER
80
+
81
+ Snapshot behavior:
82
+
83
+ - MSSQL_ENABLE_TABLE_DEFINITIONS
84
+ - MSSQL_ENABLE_INDEXES
85
+ - MSSQL_ENABLE_FOREIGN_KEYS
86
+ - MSSQL_ENABLE_VIEWS
87
+ - MSSQL_ENABLE_STORED_PROCEDURES
88
+ - MSSQL_ENABLE_FUNCTIONS
89
+ - MSSQL_ENABLE_SEQUENCES
90
+ - MSSQL_ENABLE_SYNONYMS
91
+ - MSSQL_CLEANUP_STAGING
92
+ - MSSQL_SYNC_LOG_RETENTION_DAYS
93
+
94
+ Filters:
95
+
96
+ - MSSQL_INCLUDE_SCHEMAS
97
+ - MSSQL_EXCLUDE_SCHEMAS
98
+ - MSSQL_INCLUDE_TABLES
99
+ - MSSQL_EXCLUDE_TABLES
100
+ - MSSQL_INCLUDE_VIEWS
101
+ - MSSQL_EXCLUDE_VIEWS
102
+ - MSSQL_INCLUDE_STORED_PROCEDURES
103
+ - MSSQL_EXCLUDE_STORED_PROCEDURES
104
+ - MSSQL_INCLUDE_FUNCTIONS
105
+ - MSSQL_EXCLUDE_FUNCTIONS
106
+ - MSSQL_INCLUDE_SEQUENCES
107
+ - MSSQL_EXCLUDE_SEQUENCES
108
+ - MSSQL_INCLUDE_SYNONYMS
109
+ - MSSQL_EXCLUDE_SYNONYMS
110
+
111
+ Notes:
112
+
113
+ - If MSSQL_DATABASES is set, it is treated as a comma-separated list.
114
+ - If MSSQL_DATABASES is empty, MSSQL_DATABASE is used.
115
+ - If MSSQL_DATABASE contains commas, it is also treated as a list for backward compatibility.
116
+
117
+ ## File Sync Utilities
118
+
119
+ `files_sync_utils` provides:
120
+
121
+ - Inventory building
122
+ - Plan generation
123
+ - Copy and update execution
124
+ - Extra destination cleanup rules
125
+ - Activity CSV logging
126
+ - Log retention cleanup
127
+
128
+ ## Git Utilities
129
+
130
+ `git_utils` wraps common git automation operations with consistent behavior and error handling.
131
+
132
+ ## Development Guides
133
+
134
+ - development_guides/dev_guide.txt
135
+ - development_guides/auth_key_guide.txt
136
+
137
+ These guides contain release flow notes and auth-related instructions.
138
+
139
+ ## Changelog
140
+
141
+ See CHANGELOG.md for release notes.
@@ -0,0 +1,58 @@
1
+ [project]
2
+ name = "al-data-utils"
3
+ version = "0.3.1"
4
+ description = "Reusable utilities for data engineering (Airflow, Azure, Snowflake, SQL, config)."
5
+ authors = [{ name = "AL Data Labs", email = "pypi@aldatalabs.com" }]
6
+ license = "MIT"
7
+ license-files = ["LICENSE"]
8
+ keywords = ["airflow", "data-engineering", "azure", "snowflake", "sql"]
9
+ classifiers = [
10
+ "Programming Language :: Python :: 3",
11
+ "Operating System :: OS Independent",
12
+ "Topic :: Software Development :: Libraries",
13
+ ]
14
+ readme = "README.md"
15
+ requires-python = ">=3.10"
16
+ dependencies = [] # keep empty for now, using airflow environment's dependencies
17
+
18
+ [project.urls]
19
+ Homepage = "https://aldatalabs.com"
20
+
21
+
22
+ [build-system]
23
+ requires = ["setuptools>=77"]
24
+ build-backend = "setuptools.build_meta"
25
+
26
+ [tool.setuptools]
27
+ package-dir = {"" = "src"}
28
+
29
+ [tool.setuptools.packages.find]
30
+ where = ["src"]
31
+
32
+ # Ruff configuration
33
+ [tool.ruff]
34
+ # Exclude a variety of commonly ignored directories.
35
+ exclude = [
36
+ ".git",
37
+ ".venv",
38
+ ]
39
+
40
+ # Assume Python 3.10
41
+ target-version = "py310"
42
+
43
+ # Same as Black.
44
+ line-length = 88
45
+ indent-width = 4
46
+
47
+
48
+ [tool.ruff.format]
49
+ quote-style = "single"
50
+
51
+ # Like Black, indent with spaces, rather than tabs.
52
+ indent-style = "space"
53
+
54
+ # Like Black, respect magic trailing commas.
55
+ skip-magic-trailing-comma = true
56
+
57
+ # Like Black, automatically detect the appropriate line ending.
58
+ line-ending = "auto"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,10 @@
1
+ __all__ = [
2
+ 'airflow_utils',
3
+ 'config_utils',
4
+ 'geo_utils',
5
+ 'logging_utils',
6
+ 'azure_as_utils',
7
+ 'sql_utils',
8
+ 'tg_utils',
9
+ 'ones_api_utils',
10
+ ]
@@ -0,0 +1,13 @@
1
+ from .airflow_utils import (
2
+ get_variable_value,
3
+ get_variable_value_dict,
4
+ get_variable_value_int,
5
+ get_variable_value_str,
6
+ )
7
+
8
+ __all__ = [
9
+ "get_variable_value",
10
+ "get_variable_value_dict",
11
+ "get_variable_value_int",
12
+ "get_variable_value_str",
13
+ ]
@@ -0,0 +1,83 @@
1
+ from __future__ import annotations
2
+
3
+ import functools
4
+ import json
5
+ from pathlib import Path
6
+ from sys import platform
7
+ from typing import Any
8
+
9
+ import yaml
10
+
11
+
12
+ def _get_dags_folder() -> str:
13
+ try:
14
+ from airflow import settings # type: ignore
15
+ return settings.DAGS_FOLDER
16
+ except Exception:
17
+ return "./dags"
18
+
19
+
20
+ def _airflow_variable_get(name: str) -> str:
21
+ from airflow.sdk import Variable # type: ignore
22
+ return Variable.get(name)
23
+
24
+
25
+ @functools.cache
26
+ def get_variable_value(
27
+ airflow_variable_name: str,
28
+ pack_to_json: bool = False,
29
+ source: str = "airflow",
30
+ source_section: str = "variables",
31
+ ) -> Any:
32
+ """Get variable value different ways for windows and linux."""
33
+
34
+ if platform == "win32":
35
+ base = Path(".local/variables")
36
+ data = json.loads((base / "variables.json").read_text(encoding="utf-8"))
37
+ environment_code = data["ENVIRONMENT"]
38
+
39
+ env_data = json.loads(
40
+ (base / f"{environment_code.lower()}_variables.json").read_text(encoding="utf-8")
41
+ )
42
+ variable_value = env_data[airflow_variable_name]
43
+
44
+ else:
45
+ if source == "config":
46
+ dags_folder = _get_dags_folder()
47
+ p = Path(dags_folder) / ".configs" / "_variables.yaml"
48
+ variables = yaml.safe_load(p.read_text(encoding="utf-8")) or {}
49
+ variable_value = variables[source_section][airflow_variable_name]
50
+ else:
51
+ variable_value = _airflow_variable_get(airflow_variable_name)
52
+
53
+ if pack_to_json and not isinstance(variable_value, dict):
54
+ try:
55
+ variable_value = json.loads(variable_value)
56
+ except Exception:
57
+ pass
58
+
59
+ return variable_value
60
+
61
+
62
+ def get_variable_value_str(
63
+ airflow_variable_name: str,
64
+ source: str = "airflow",
65
+ source_section: str = "variables",
66
+ ) -> str:
67
+ return str(get_variable_value(airflow_variable_name, False, source, source_section))
68
+
69
+
70
+ def get_variable_value_int(
71
+ airflow_variable_name: str,
72
+ source: str = "airflow",
73
+ source_section: str = "variables",
74
+ ) -> int:
75
+ return int(get_variable_value(airflow_variable_name, False, source, source_section))
76
+
77
+
78
+ def get_variable_value_dict(
79
+ airflow_variable_name: str,
80
+ source: str = "airflow",
81
+ source_section: str = "variables",
82
+ ) -> dict:
83
+ return dict(get_variable_value(airflow_variable_name, True, source, source_section))
@@ -0,0 +1,31 @@
1
+ from .auth import ServicePrincipalTokenProvider, StaticTokenProvider, TokenProvider
2
+ from .azure_as_utils import SemanticModels
3
+ from .client import AasRestClient
4
+ from .errors import (
5
+ AasAuthError,
6
+ AasError,
7
+ AasHttpError,
8
+ AasRefreshError,
9
+ AasRefreshTimeout,
10
+ AasTransportError,
11
+ )
12
+ from .models import RefreshObject, RefreshRequest, RefreshStatus
13
+ from .poller import RefreshPoller
14
+
15
+ __all__ = [
16
+ 'AasAuthError',
17
+ 'AasError',
18
+ 'AasHttpError',
19
+ 'AasRefreshError',
20
+ 'AasRefreshTimeout',
21
+ 'AasRestClient',
22
+ 'AasTransportError',
23
+ 'RefreshObject',
24
+ 'RefreshPoller',
25
+ 'RefreshRequest',
26
+ 'RefreshStatus',
27
+ 'SemanticModels',
28
+ 'ServicePrincipalTokenProvider',
29
+ 'StaticTokenProvider',
30
+ 'TokenProvider',
31
+ ]
@@ -0,0 +1,82 @@
1
+ """Access tokens for the AAS REST API via an Azure AD service principal."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Mapping
6
+ from typing import Any, Protocol
7
+
8
+ import msal
9
+
10
+ from .errors import AasAuthError
11
+
12
+ #: Resource identifier of Azure Analysis Services; the wildcard is literal.
13
+ AAS_SCOPE = 'https://*.asazure.windows.net/.default'
14
+ DEFAULT_AUTHORITY_HOST = 'https://login.microsoftonline.com'
15
+
16
+
17
+ class TokenProvider(Protocol):
18
+ """Anything that can hand out a bearer token for the AAS resource."""
19
+
20
+ def get_token(self) -> str: ...
21
+
22
+
23
+ class StaticTokenProvider:
24
+ """Returns a fixed token. Useful for tests and pre-issued tokens."""
25
+
26
+ def __init__(self, token: str) -> None:
27
+ if not token:
28
+ raise ValueError('token must be a non-empty string')
29
+ self._token = token
30
+
31
+ def get_token(self) -> str:
32
+ return self._token
33
+
34
+
35
+ class ServicePrincipalTokenProvider:
36
+ """Client-credentials flow through MSAL.
37
+
38
+ MSAL keeps the issued token in its in-memory cache and only goes back to
39
+ Azure AD when it is about to expire, so calling :meth:`get_token` before
40
+ every request is cheap.
41
+ """
42
+
43
+ def __init__(
44
+ self,
45
+ tenant_id: str,
46
+ client_id: str,
47
+ client_secret: str,
48
+ *,
49
+ authority_host: str = DEFAULT_AUTHORITY_HOST,
50
+ scope: str = AAS_SCOPE,
51
+ ) -> None:
52
+ for name, value in (
53
+ ('tenant_id', tenant_id),
54
+ ('client_id', client_id),
55
+ ('client_secret', client_secret),
56
+ ):
57
+ if not isinstance(value, str) or not value.strip():
58
+ raise ValueError(f'{name} must be a non-empty string')
59
+ self._scopes = [scope]
60
+ self._authority = f'{authority_host.rstrip("/")}/{tenant_id}'
61
+ self._app = msal.ConfidentialClientApplication(
62
+ client_id, authority=self._authority, client_credential=client_secret
63
+ )
64
+
65
+ def get_token(self) -> str:
66
+ result: Mapping[str, Any] | None = self._app.acquire_token_for_client(
67
+ scopes=self._scopes
68
+ )
69
+ token = result.get('access_token') if isinstance(result, Mapping) else None
70
+ if not token:
71
+ raise AasAuthError(_describe_failure(result))
72
+ return token
73
+
74
+
75
+ def _describe_failure(result: Mapping[str, Any] | None) -> str:
76
+ if not isinstance(result, Mapping):
77
+ return f'token request returned no result ({result!r})'
78
+ error = result.get('error') or 'unknown_error'
79
+ description = result.get('error_description') or 'no description'
80
+ correlation = result.get('correlation_id')
81
+ suffix = f' (correlation_id={correlation})' if correlation else ''
82
+ return f'token request failed: {error}: {description}{suffix}'