airflow-dbt-python 2.1.0__tar.gz → 3.0.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 (24) hide show
  1. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/PKG-INFO +25 -25
  2. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/README.md +3 -3
  3. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/__version__.py +1 -1
  4. airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/__init__.py +1 -0
  5. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/hooks/dbt.py +48 -259
  6. airflow_dbt_python-2.1.0/airflow_dbt_python/hooks/remote.py → airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/fs/__init__.py +25 -19
  7. airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/fs/gcs.py +228 -0
  8. {airflow_dbt_python-2.1.0/airflow_dbt_python/hooks → airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/fs}/git.py +21 -14
  9. airflow_dbt_python-2.1.0/airflow_dbt_python/hooks/localfs.py → airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/fs/local.py +10 -38
  10. {airflow_dbt_python-2.1.0/airflow_dbt_python/hooks → airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/fs}/s3.py +8 -8
  11. airflow_dbt_python-3.0.0/airflow_dbt_python/hooks/target.py +464 -0
  12. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/operators/dbt.py +120 -6
  13. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/utils/configs.py +94 -85
  14. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/utils/url.py +3 -3
  15. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/utils/version.py +5 -6
  16. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/pyproject.toml +65 -55
  17. airflow_dbt_python-2.1.0/airflow_dbt_python/hooks/__init__.py +0 -1
  18. airflow_dbt_python-2.1.0/setup.py +0 -51
  19. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/LICENSE +0 -0
  20. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/__init__.py +0 -0
  21. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/operators/__init__.py +0 -0
  22. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/utils/__init__.py +0 -0
  23. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/utils/enums.py +0 -0
  24. {airflow_dbt_python-2.1.0 → airflow_dbt_python-3.0.0}/airflow_dbt_python/utils/env.py +0 -0
@@ -1,48 +1,48 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.3
2
2
  Name: airflow-dbt-python
3
- Version: 2.1.0
3
+ Version: 3.0.0
4
4
  Summary: A collection of Airflow operators, hooks, and utilities to execute dbt commands
5
- Home-page: https://github.com/tomasfarias/airflow-dbt-python
6
5
  License: MIT
7
6
  Author: Tomás Farías Santana
8
7
  Author-email: tomas@tomasfarias.dev
9
- Requires-Python: >=3.8,<3.13
8
+ Requires-Python: >=3.9,<3.13
10
9
  Classifier: Development Status :: 5 - Production/Stable
11
10
  Classifier: Intended Audience :: Developers
12
11
  Classifier: License :: OSI Approved :: MIT License
13
- Classifier: Programming Language :: Python :: 3
14
- Classifier: Programming Language :: Python :: 3.8
15
12
  Classifier: Programming Language :: Python :: 3.9
16
13
  Classifier: Programming Language :: Python :: 3.10
17
14
  Classifier: Programming Language :: Python :: 3.11
18
- Classifier: Programming Language :: Python :: 3.10
19
- Classifier: Programming Language :: Python :: 3.11
20
15
  Classifier: Programming Language :: Python :: 3.12
21
- Classifier: Programming Language :: Python :: 3.8
22
- Classifier: Programming Language :: Python :: 3.9
23
16
  Provides-Extra: adapters
24
17
  Provides-Extra: airflow-providers
25
18
  Provides-Extra: bigquery
26
- Provides-Extra: docs
19
+ Provides-Extra: gcs
27
20
  Provides-Extra: git
28
21
  Provides-Extra: postgres
29
22
  Provides-Extra: redshift
30
23
  Provides-Extra: s3
31
24
  Provides-Extra: snowflake
32
25
  Provides-Extra: spark
33
- Requires-Dist: apache-airflow (>=2.7)
34
- Requires-Dist: apache-airflow-providers-amazon (>=3.0.0) ; extra == "airflow-providers" or extra == "s3"
35
- Requires-Dist: apache-airflow-providers-ssh (>=3.0.0) ; extra == "airflow-providers" or extra == "git"
36
- Requires-Dist: connexion (<3.0)
26
+ Requires-Dist: apache-airflow (>=2.8)
27
+ Requires-Dist: apache-airflow-providers-amazon (>=3.0.0) ; extra == "airflow-providers"
28
+ Requires-Dist: apache-airflow-providers-amazon (>=3.0.0) ; extra == "s3"
29
+ Requires-Dist: apache-airflow-providers-google (>=11.0.0) ; extra == "airflow-providers"
30
+ Requires-Dist: apache-airflow-providers-google (>=11.0.0) ; extra == "gcs"
31
+ Requires-Dist: apache-airflow-providers-ssh (>=3.0.0) ; extra == "airflow-providers"
32
+ Requires-Dist: apache-airflow-providers-ssh (>=3.0.0) ; extra == "git"
37
33
  Requires-Dist: contextlib-chdir (==1.0.2) ; python_version < "3.11"
38
- Requires-Dist: dbt-bigquery (>=1.7.0,<1.9.0) ; extra == "adapters" or extra == "bigquery"
39
- Requires-Dist: dbt-core (>=1.7.5,<1.9.0)
40
- Requires-Dist: dbt-postgres (>=1.7.0,<1.9.0) ; extra == "adapters" or extra == "postgres"
41
- Requires-Dist: dbt-redshift (>=1.7.0,<1.9.0) ; extra == "adapters" or extra == "redshift"
42
- Requires-Dist: dbt-snowflake (>=1.7.0,<1.9.0) ; extra == "adapters" or extra == "snowflake"
43
- Requires-Dist: dbt-spark (>=1.7.0,<1.9.0) ; extra == "adapters" or extra == "spark"
34
+ Requires-Dist: dbt-bigquery (>=1.8.0,<2.0.0) ; extra == "adapters"
35
+ Requires-Dist: dbt-bigquery (>=1.8.0,<2.0.0) ; extra == "bigquery"
36
+ Requires-Dist: dbt-core (>=1.8.0,<2.0.0)
37
+ Requires-Dist: dbt-postgres (>=1.8.0,<2.0.0) ; extra == "adapters"
38
+ Requires-Dist: dbt-postgres (>=1.8.0,<2.0.0) ; extra == "postgres"
39
+ Requires-Dist: dbt-redshift (>=1.8.0,<2.0.0) ; extra == "adapters"
40
+ Requires-Dist: dbt-redshift (>=1.8.0,<2.0.0) ; extra == "redshift"
41
+ Requires-Dist: dbt-snowflake (>=1.8.0,<2.0.0) ; extra == "adapters"
42
+ Requires-Dist: dbt-snowflake (>=1.8.0,<2.0.0) ; extra == "snowflake"
43
+ Requires-Dist: dbt-spark (>=1.8.0,<2.0.0) ; extra == "spark"
44
+ Requires-Dist: dbt-spark[all] (>=1.8.0,<2.0.0) ; extra == "adapters"
44
45
  Requires-Dist: dulwich (>=0.21) ; extra == "git"
45
- Project-URL: Repository, https://github.com/tomasfarias/airflow-dbt-python
46
46
  Description-Content-Type: text/markdown
47
47
 
48
48
  # airflow-dbt-python
@@ -62,13 +62,13 @@ Read the [documentation](https://airflow-dbt-python.readthedocs.io) for examples
62
62
  ## Requirements
63
63
 
64
64
  Before using *airflow-dbt-python*, ensure you meet the following requirements:
65
- * A *dbt* project using [dbt-core](https://pypi.org/project/dbt-core/) version 1.7.5 or later.
66
- * An Airflow environment using version 2.7 or later.
65
+ * A *dbt* project using [dbt-core](https://pypi.org/project/dbt-core/) version 1.8 or later.
66
+ * An Airflow environment using version 2.8 or later.
67
67
 
68
68
  * If using any managed service, like AWS MWAA or GCP Cloud Composer 2/3, ensure your environment is created with a supported version of Airflow.
69
69
  * If self-hosting, Airflow installation instructions can be found in their [official documentation](https://airflow.apache.org/docs/apache-airflow/stable/installation/index.html).
70
70
 
71
- * Running Python 3.8 or later in your Airflow environment.
71
+ * Running Python 3.9 or later in your Airflow environment.
72
72
 
73
73
  > **Warning**
74
74
  >
@@ -15,13 +15,13 @@ Read the [documentation](https://airflow-dbt-python.readthedocs.io) for examples
15
15
  ## Requirements
16
16
 
17
17
  Before using *airflow-dbt-python*, ensure you meet the following requirements:
18
- * A *dbt* project using [dbt-core](https://pypi.org/project/dbt-core/) version 1.7.5 or later.
19
- * An Airflow environment using version 2.7 or later.
18
+ * A *dbt* project using [dbt-core](https://pypi.org/project/dbt-core/) version 1.8 or later.
19
+ * An Airflow environment using version 2.8 or later.
20
20
 
21
21
  * If using any managed service, like AWS MWAA or GCP Cloud Composer 2/3, ensure your environment is created with a supported version of Airflow.
22
22
  * If self-hosting, Airflow installation instructions can be found in their [official documentation](https://airflow.apache.org/docs/apache-airflow/stable/installation/index.html).
23
23
 
24
- * Running Python 3.8 or later in your Airflow environment.
24
+ * Running Python 3.9 or later in your Airflow environment.
25
25
 
26
26
  > **Warning**
27
27
  >
@@ -3,4 +3,4 @@
3
3
  __author__ = "Tomás Farías Santana"
4
4
  __copyright__ = "Copyright 2021 Tomás Farías Santana"
5
5
  __title__ = "airflow-dbt-python"
6
- __version__ = "2.1.0"
6
+ __version__ = "3.0.0"
@@ -0,0 +1 @@
1
+ """Hooks module provides DbtHooks and DbtFSHooks."""
@@ -5,6 +5,7 @@ from __future__ import annotations
5
5
  import json
6
6
  import logging
7
7
  import sys
8
+ from abc import ABC
8
9
  from contextlib import contextmanager
9
10
  from pathlib import Path
10
11
  from tempfile import TemporaryDirectory
@@ -17,15 +18,11 @@ from typing import (
17
18
  NamedTuple,
18
19
  Optional,
19
20
  Tuple,
20
- Union,
21
21
  )
22
22
  from urllib.parse import urlparse
23
23
 
24
24
  from airflow.exceptions import AirflowException
25
- from airflow.hooks.base import BaseHook
26
- from airflow.models.connection import Connection
27
-
28
- from airflow_dbt_python.utils.version import DBT_INSTALLED_GTE_1_8
25
+ from airflow.utils.log.logging_mixin import LoggingMixin
29
26
 
30
27
  if sys.version_info >= (3, 11):
31
28
  from contextlib import chdir as chdir_ctx
@@ -34,13 +31,13 @@ else:
34
31
 
35
32
  if TYPE_CHECKING:
36
33
  from dbt.contracts.results import RunResult
37
- from dbt.task.base import BaseTask
38
34
 
39
- from airflow_dbt_python.hooks.remote import DbtRemoteHook
35
+ from airflow_dbt_python.hooks.fs import DbtFSHook
36
+ from airflow_dbt_python.hooks.target import DbtConnectionHook
40
37
  from airflow_dbt_python.utils.configs import BaseConfig
41
38
  from airflow_dbt_python.utils.url import URLLike
42
39
 
43
- DbtRemoteHooksDict = Dict[Tuple[str, Optional[str]], DbtRemoteHook]
40
+ DbtFSHooksDict = Dict[Tuple[str, Optional[str]], DbtFSHook]
44
41
 
45
42
 
46
43
  class DbtTaskResult(NamedTuple):
@@ -57,35 +54,6 @@ class DbtTaskResult(NamedTuple):
57
54
  artifacts: dict[str, Any]
58
55
 
59
56
 
60
- class DbtConnectionParam(NamedTuple):
61
- """A tuple indicating connection parameters relevant to dbt.
62
-
63
- Attributes:
64
- name: The name of the connection parameter. This name will be used to get the
65
- parameter from an Airflow Connection or its extras.
66
- store_override_name: A new name for the connection parameter. If not None, this
67
- is the name used in a dbt profiles.
68
- default: A default value if the parameter is not found.
69
- """
70
-
71
- name: str
72
- store_override_name: Optional[str] = None
73
- default: Optional[Any] = None
74
-
75
- @property
76
- def override_name(self):
77
- """Returns the override_name if defined, otherwise defaults to name.
78
-
79
- >>> DbtConnectionParam("login", "user").override_name
80
- 'user'
81
- >>> DbtConnectionParam("port").override_name
82
- 'port'
83
- """
84
- if self.store_override_name is None:
85
- return self.name
86
- return self.store_override_name
87
-
88
-
89
57
  class DbtTemporaryDirectory(TemporaryDirectory):
90
58
  """A wrapper on TemporaryDirectory for older versions of Python.
91
59
 
@@ -107,57 +75,43 @@ class DbtTemporaryDirectory(TemporaryDirectory):
107
75
  )
108
76
 
109
77
 
110
- class DbtHook(BaseHook):
78
+ class DbtHook(ABC, LoggingMixin):
111
79
  """A hook to interact with dbt.
112
80
 
113
81
  Allows for running dbt tasks and provides required configurations for each task.
114
82
  """
115
83
 
116
- conn_name_attr = "dbt_conn_id"
117
- default_conn_name = "dbt_default"
118
- conn_type = "dbt"
119
- hook_name = "dbt Hook"
120
-
121
- conn_params: list[Union[DbtConnectionParam, str]] = [
122
- DbtConnectionParam("conn_type", "type"),
123
- "host",
124
- DbtConnectionParam("conn_id", "dbname"),
125
- "schema",
126
- DbtConnectionParam("login", "user"),
127
- "password",
128
- "port",
129
- ]
130
- conn_extra_params: list[Union[DbtConnectionParam, str]] = []
131
-
132
84
  def __init__(
133
85
  self,
134
86
  *args,
135
- dbt_conn_id: Optional[str] = default_conn_name,
87
+ dbt_conn_id: Optional[str] = None,
136
88
  project_conn_id: Optional[str] = None,
137
89
  profiles_conn_id: Optional[str] = None,
138
90
  **kwargs,
139
91
  ):
140
- self.remotes: DbtRemoteHooksDict = {}
141
92
  self.dbt_conn_id = dbt_conn_id
142
93
  self.project_conn_id = project_conn_id
143
94
  self.profiles_conn_id = profiles_conn_id
144
95
  super().__init__(*args, **kwargs)
145
96
 
146
- def get_remote(self, scheme: str, conn_id: Optional[str]) -> DbtRemoteHook:
147
- """Get a remote to interact with dbt files.
97
+ @staticmethod
98
+ def get_dbt_target_hook(conn_id: str) -> DbtConnectionHook:
99
+ """Get a hook to get a dbt profile based on the Airflow connection."""
100
+ from .target import DbtConnectionHook
148
101
 
149
- RemoteHooks are defined by the scheme we are looking for and an optional
102
+ return DbtConnectionHook.get_db_conn_hook(conn_id)
103
+
104
+ @staticmethod
105
+ def get_fs_hook(scheme: str, conn_id: Optional[str]) -> DbtFSHook:
106
+ """Get a fs_hook to interact with dbt files.
107
+
108
+ FSHooks are defined by the scheme we are looking for and an optional
150
109
  connection id if we are looking to interface with any Airflow hook that
151
110
  uses a connection.
152
111
  """
153
- from .remote import get_remote
112
+ from .fs import get_fs_hook
154
113
 
155
- try:
156
- return self.remotes[(scheme, conn_id)]
157
- except KeyError:
158
- remote = get_remote(scheme, conn_id)
159
- self.remotes[(scheme, conn_id)] = remote
160
- return remote
114
+ return get_fs_hook(scheme, conn_id)
161
115
 
162
116
  def download_dbt_profiles(
163
117
  self,
@@ -166,13 +120,13 @@ class DbtHook(BaseHook):
166
120
  ) -> Path:
167
121
  """Pull a dbt profiles.yml file from a given profiles_dir.
168
122
 
169
- This operation is delegated to a DbtRemoteHook. An optional connection id is
123
+ This operation is delegated to a DbtFSHook. An optional connection id is
170
124
  supported for remotes that require it.
171
125
  """
172
126
  scheme = urlparse(str(profiles_dir)).scheme
173
- remote = self.get_remote(scheme, self.project_conn_id)
127
+ fs_hook = self.get_fs_hook(scheme, self.profiles_conn_id)
174
128
 
175
- return remote.download_dbt_profiles(profiles_dir, destination)
129
+ return fs_hook.download_dbt_profiles(profiles_dir, destination)
176
130
 
177
131
  def download_dbt_project(
178
132
  self,
@@ -181,13 +135,13 @@ class DbtHook(BaseHook):
181
135
  ) -> Path:
182
136
  """Pull a dbt project from a given project_dir.
183
137
 
184
- This operation is delegated to a DbtRemoteHook. An optional connection id is
138
+ This operation is delegated to a DbtFSHook. An optional connection id is
185
139
  supported for remotes that require it.
186
140
  """
187
141
  scheme = urlparse(str(project_dir)).scheme
188
- remote = self.get_remote(scheme, self.project_conn_id)
142
+ fs_hook = self.get_fs_hook(scheme, self.project_conn_id)
189
143
 
190
- return remote.download_dbt_project(project_dir, destination)
144
+ return fs_hook.download_dbt_project(project_dir, destination)
191
145
 
192
146
  def upload_dbt_project(
193
147
  self,
@@ -198,13 +152,13 @@ class DbtHook(BaseHook):
198
152
  ) -> None:
199
153
  """Push a dbt project from a given project_dir.
200
154
 
201
- This operation is delegated to a DbtRemoteHook. An optional connection id is
155
+ This operation is delegated to a DbtFSHook. An optional connection id is
202
156
  supported for remotes that require it.
203
157
  """
204
158
  scheme = urlparse(str(destination)).scheme
205
- remote = self.get_remote(scheme, self.project_conn_id)
159
+ fs_hook = self.get_fs_hook(scheme, self.project_conn_id)
206
160
 
207
- return remote.upload_dbt_project(
161
+ return fs_hook.upload_dbt_project(
208
162
  project_dir, destination, replace=replace, delete_before=delete_before
209
163
  )
210
164
 
@@ -233,8 +187,14 @@ class DbtHook(BaseHook):
233
187
  from dbt.task.deps import DepsTask
234
188
  from dbt.tracking import track_run
235
189
 
190
+ if self.dbt_conn_id:
191
+ kwargs["target"] = self.dbt_conn_id
192
+ target_hook = self.get_dbt_target_hook(self.dbt_conn_id)
193
+ extra_target = target_hook.get_dbt_target_from_connection()
194
+ else:
195
+ extra_target = None
196
+
236
197
  config = self.get_dbt_task_config(command, **kwargs)
237
- extra_target = self.get_dbt_target_from_connection(config.target)
238
198
 
239
199
  with self.dbt_directory(
240
200
  config,
@@ -248,15 +208,13 @@ class DbtHook(BaseHook):
248
208
  nearest_project_dir = get_nearest_project_dir(config.project_dir)
249
209
 
250
210
  with chdir_ctx(nearest_project_dir):
251
- self.ensure_profiles(config)
252
-
253
211
  with adapter_management():
254
212
  task, runtime_config = config.create_dbt_task(
255
213
  extra_target, write_perf_info
256
214
  )
257
215
  requires_profile = isinstance(task, (CleanTask, DepsTask))
258
216
 
259
- self.setup_dbt_logging(task, config.debug)
217
+ self.setup_dbt_logging(config.debug)
260
218
 
261
219
  if runtime_config is not None and not requires_profile:
262
220
  # The deps command installs the dependencies, which means they
@@ -375,38 +333,31 @@ class DbtHook(BaseHook):
375
333
  project_dir,
376
334
  tmp_dir,
377
335
  )
378
- new_project_dir = str(project_dir_path) + "/"
379
-
380
- if (project_dir_path / "profiles.yml").exists():
381
- # We may have downloaded the profiles.yml file together
382
- # with the project.
383
- return new_project_dir, new_project_dir
384
336
 
385
337
  if profiles_dir is not None:
386
338
  profiles_file_path = self.download_dbt_profiles(
387
339
  profiles_dir,
388
340
  tmp_dir,
389
341
  )
390
- new_profiles_dir = str(profiles_file_path.parent) + "/"
342
+ profiles_dir_path = profiles_file_path.parent
343
+ elif (project_dir_path / "profiles.yml").exists():
344
+ profiles_dir_path = project_dir_path
391
345
  else:
392
- new_profiles_dir = None
346
+ profiles_dir_path = None
393
347
 
394
- return new_project_dir, new_profiles_dir
348
+ return (
349
+ str(project_dir_path),
350
+ str(profiles_dir_path) if profiles_dir_path is not None else None,
351
+ )
395
352
 
396
- def setup_dbt_logging(self, task: BaseTask, debug: Optional[bool]):
353
+ def setup_dbt_logging(self, debug: Optional[bool]):
397
354
  """Setup dbt logging.
398
355
 
399
356
  Starting with dbt v1, dbt initializes two loggers: default_file and
400
357
  default_stdout. As these are initialized by the CLI app, we need to
401
358
  initialize them here.
402
359
  """
403
- if DBT_INSTALLED_GTE_1_8:
404
- from dbt.events.logging import setup_event_logger
405
- else:
406
- from dbt.events.functions import ( # type: ignore[no-redef]
407
- setup_event_logger,
408
- )
409
-
360
+ from dbt.events.logging import setup_event_logger
410
361
  from dbt.flags import get_flags
411
362
 
412
363
  flags = get_flags()
@@ -415,6 +366,7 @@ class DbtHook(BaseHook):
415
366
  configured_file = logging.getLogger("configured_file")
416
367
  file_log = logging.getLogger("file_log")
417
368
  stdout_log = logging.getLogger("stdout_log")
369
+ stdout_log.handlers.clear()
418
370
  stdout_log.propagate = True
419
371
 
420
372
  if not debug:
@@ -425,166 +377,3 @@ class DbtHook(BaseHook):
425
377
  file_log.propagate = False
426
378
  configured_file.setLevel("INFO")
427
379
  configured_file.propagate = False
428
-
429
- def ensure_profiles(self, config: BaseConfig):
430
- """Ensure a profiles file exists."""
431
- if config.profiles_dir is not None:
432
- # We expect one to exist given that we have passed a profiles_dir.
433
- return
434
-
435
- profiles_path = Path.home() / ".dbt/profiles.yml"
436
- config.profiles_dir = str(profiles_path.parent)
437
- if not profiles_path.exists():
438
- profiles_path.parent.mkdir(exist_ok=True)
439
- with profiles_path.open("w", encoding="utf-8") as f:
440
- f.write("flags:\n send_anonymous_usage_stats: false\n")
441
-
442
- def get_dbt_target_from_connection(
443
- self, target: Optional[str]
444
- ) -> Optional[dict[str, Any]]:
445
- """Return a dictionary of connection details to use as a dbt target.
446
-
447
- The connection details are fetched from an Airflow connection identified by
448
- target or self.dbt_conn_id.
449
-
450
- Args:
451
- target: The target name to use as an Airflow connection ID. If ommitted, we
452
- will use self.dbt_conn_id.
453
-
454
- Returns:
455
- A dictionary with a configuration for a dbt target, or None if a matching
456
- Airflow connection is not found for given dbt target.
457
- """
458
- conn_id = target or self.dbt_conn_id
459
-
460
- if conn_id is None:
461
- return None
462
-
463
- try:
464
- conn = self.get_connection(conn_id)
465
- except AirflowException:
466
- self.log.debug(
467
- "No Airflow connection matching dbt target %s was found.", target
468
- )
469
- return None
470
-
471
- details = self.get_dbt_details_from_connection(conn)
472
-
473
- return {conn_id: details}
474
-
475
- def get_dbt_details_from_connection(self, conn: Connection) -> dict[str, Any]:
476
- """Extract dbt connection details from Airflow Connection.
477
-
478
- dbt connection details may be present as Airflow Connection attributes or in the
479
- Connection's extras. This class' conn_params and conn_extra_params will be used
480
- to fetch required attributes from attributes and extras respectively. If
481
- conn_extra_params is empty, we merge parameters with all extras.
482
-
483
- Subclasses may override this class attributes to narrow down the connection
484
- details for a specific dbt target (like Postgres, or Redshift).
485
-
486
- Args:
487
- conn: The Airflow Connection to extract dbt connection details from.
488
-
489
- Returns:
490
- A dictionary of dbt connection details.
491
- """
492
- dbt_details = {}
493
- for param in self.conn_params:
494
- if isinstance(param, DbtConnectionParam):
495
- key = param.override_name
496
- value = getattr(conn, param.name, param.default)
497
- else:
498
- key = param
499
- value = getattr(conn, key, None)
500
-
501
- if value is None:
502
- continue
503
-
504
- dbt_details[key] = value
505
-
506
- extra = conn.extra_dejson
507
-
508
- if not self.conn_extra_params:
509
- return {**dbt_details, **extra}
510
-
511
- for param in self.conn_extra_params:
512
- if isinstance(param, DbtConnectionParam):
513
- key = param.override_name
514
- value = extra.get(param.name, param.default)
515
- else:
516
- key = param
517
- value = extra.get(key, None)
518
-
519
- if value is None:
520
- continue
521
-
522
- dbt_details[key] = value
523
-
524
- return dbt_details
525
-
526
-
527
- class DbtPostgresHook(DbtHook):
528
- """A hook to interact with dbt using a Postgres connection."""
529
-
530
- conn_type = "postgres"
531
- hook_name = "dbt Postgres Hook"
532
- conn_params = [
533
- DbtConnectionParam("conn_type", "type", "postgres"),
534
- "host",
535
- "schema",
536
- DbtConnectionParam("login", "user"),
537
- "password",
538
- "port",
539
- ]
540
- conn_extra_params = [
541
- "dbname",
542
- "threads",
543
- "keepalives_idle",
544
- "connect_timeout",
545
- "retries",
546
- "search_path",
547
- "role",
548
- "sslmode",
549
- ]
550
-
551
-
552
- class DbtRedshiftHook(DbtPostgresHook):
553
- """A hook to interact with dbt using a Redshift connection."""
554
-
555
- conn_type = "redshift"
556
- hook_name = "dbt Redshift Hook"
557
- conn_extra_params = DbtPostgresHook.conn_extra_params + [
558
- "ra3_node",
559
- "iam_profile",
560
- "iam_duration_secons",
561
- "autocreate",
562
- "db_groups",
563
- ]
564
-
565
-
566
- class DbtSnowflakeHook(DbtHook):
567
- """A hook to interact with dbt using a Snowflake connection."""
568
-
569
- conn_type = "snowflake"
570
- hook_name = "dbt Snowflake Hook"
571
- conn_params = [
572
- DbtConnectionParam("conn_type", "type", "postgres"),
573
- "host",
574
- "schema",
575
- DbtConnectionParam("login", "user"),
576
- "password",
577
- ]
578
- conn_extra_params = [
579
- "account",
580
- "role",
581
- "database",
582
- "warehouse",
583
- "threads",
584
- "client_session_keep_alive",
585
- "query_tag",
586
- "connect_retries",
587
- "connect_timeout",
588
- "retry_on_database_errors",
589
- "retry_all",
590
- ]
@@ -1,11 +1,12 @@
1
- """The DbtRemoteHook interface includes methods for downloading and uploading files.
1
+ """The DbtFSHook interface includes methods for downloading and uploading files.
2
2
 
3
- Internally, DbtRemoteHooks can use Airflow hooks to execute the actual operations.
3
+ Internally, DbtFSHooks can use Airflow hooks to execute the actual operations.
4
4
 
5
5
  Currently, only AWS S3 and the local filesystem are supported as remotes.
6
6
  """
7
7
 
8
8
  from abc import ABC, abstractmethod
9
+ from functools import cache
9
10
  from pathlib import Path
10
11
  from typing import Optional, Type
11
12
 
@@ -16,7 +17,7 @@ from airflow_dbt_python.utils.url import URL, URLLike
16
17
  StrPath = str
17
18
 
18
19
 
19
- class DbtRemoteHook(ABC, LoggingMixin):
20
+ class DbtFSHook(ABC, LoggingMixin):
20
21
  """Represents a dbt project storing any dbt files.
21
22
 
22
23
  A concrete backend class should implement the push and pull methods to fetch one
@@ -32,7 +33,7 @@ class DbtRemoteHook(ABC, LoggingMixin):
32
33
  """
33
34
 
34
35
  @abstractmethod
35
- def download(
36
+ def _download(
36
37
  self,
37
38
  source: URL,
38
39
  destination: URL,
@@ -43,7 +44,7 @@ class DbtRemoteHook(ABC, LoggingMixin):
43
44
  return NotImplemented
44
45
 
45
46
  @abstractmethod
46
- def upload(
47
+ def _upload(
47
48
  self,
48
49
  source: URL,
49
50
  destination: URL,
@@ -71,7 +72,7 @@ class DbtRemoteHook(ABC, LoggingMixin):
71
72
  if source_url.is_archive():
72
73
  destination_url = destination_url / source_url.name
73
74
 
74
- self.download(source_url, destination_url)
75
+ self._download(source_url, destination_url)
75
76
 
76
77
  if destination_url.exists() and destination_url.is_archive():
77
78
  destination_url.extract()
@@ -103,7 +104,7 @@ class DbtRemoteHook(ABC, LoggingMixin):
103
104
  if destination_url.is_dir() or destination_url.name != "profiles.yml":
104
105
  destination_url = destination_url / "profiles.yml"
105
106
 
106
- self.download(source_url, destination_url)
107
+ self._download(source_url, destination_url)
107
108
 
108
109
  return destination_url.path
109
110
 
@@ -133,35 +134,40 @@ class DbtRemoteHook(ABC, LoggingMixin):
133
134
  source_url.archive(zip_url)
134
135
  source_url = zip_url
135
136
 
136
- self.upload(source_url, destination_url, replace, delete_before)
137
+ self._upload(source_url, destination_url, replace, delete_before)
137
138
 
138
139
  if destination_url.is_archive():
139
140
  source_url.unlink()
140
141
 
141
142
 
142
- def get_remote(scheme: str, conn_id: Optional[str] = None) -> DbtRemoteHook:
143
- """Get a DbtRemoteHook as long as the scheme is supported.
143
+ @cache
144
+ def get_fs_hook(scheme: str, conn_id: Optional[str] = None) -> DbtFSHook:
145
+ """Get a DbtFSHook as long as the scheme is supported.
144
146
 
145
147
  In the future we should make our hooks discoverable and package ourselves as a
146
148
  proper Airflow providers package.
147
149
  """
148
150
  if scheme == "s3":
149
- from .s3 import DbtS3RemoteHook
151
+ from .s3 import DbtS3FSHook
150
152
 
151
- remote_cls: Type[DbtRemoteHook] = DbtS3RemoteHook
153
+ fs_hook_cls: Type[DbtFSHook] = DbtS3FSHook
154
+ elif scheme == "gs":
155
+ from .gcs import DbtGCSFSHook
156
+
157
+ fs_hook_cls = DbtGCSFSHook
152
158
  elif scheme in ("https", "git", "git+ssh", "ssh", "http"):
153
- from .git import DbtGitRemoteHook
159
+ from .git import DbtGitFSHook
154
160
 
155
- remote_cls = DbtGitRemoteHook
161
+ fs_hook_cls = DbtGitFSHook
156
162
  elif scheme == "":
157
- from .localfs import DbtLocalFsRemoteHook
163
+ from .local import DbtLocalFsHook
158
164
 
159
- remote_cls = DbtLocalFsRemoteHook
165
+ fs_hook_cls = DbtLocalFsHook
160
166
  else:
161
167
  raise NotImplementedError(f"Backend {scheme} is not supported")
162
168
 
163
169
  if conn_id is not None:
164
- remote = remote_cls(conn_id)
170
+ fs_hook = fs_hook_cls(conn_id)
165
171
  else:
166
- remote = remote_cls()
167
- return remote
172
+ fs_hook = fs_hook_cls()
173
+ return fs_hook