Pydantic 2 (#3)

* [WIP] Started converting to pydantic 2 * [WIP] Started converting to pydantic 2 * Bump version: 0.3.0 → 1.0.0 * [WIP] Started converting to pydantic 2 --------- Co-authored-by: David Maxson <[email protected]>
rearc-data · Nov 21, 2023 · ce441cc · ce441cc · github-actions · Nov 21, 2023
1 parent 139fb0e
commit ce441cc
Show file tree

Hide file tree

Showing 24 changed files with 236 additions and 279 deletions.
diff --git a/.bumpversion.cfg b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.3.0
+current_version = 1.0.0
 commit = True
 tag = True
 

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -7,36 +7,16 @@ jobs:
   docs_build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: ammaraskar/sphinx-action@master
-        with:
-          docs-folder: "docs/"
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+        cache: 'pip' # caching pip dependencies
+    - run: pip install -r docs/requirements.txt
+    - run: cd docs && make html
 
-      - uses: actions/upload-artifact@v1
-        with:
-          name: DocumentationHTML
-          path: docs/_build/html/
-
-# Use RTFD instead
-#  publish_gh_pages:
-#    runs-on: ubuntu-latest
-#    needs: docs_build
-#    if: github.ref == 'refs/heads/master'
-#    steps:
-#      - name: Commit documentation changes
-#        run: |
-#          git clone https://github.com/your_git/repository.git --branch gh-pages --single-branch gh-pages
-#          cp -r docs/_build/html/* gh-pages/
-#          cd gh-pages
-#          git config --local user.email "[email protected]"
-#          git config --local user.name "GitHub Action"
-#          git add .
-#          git commit -m "Update documentation" -a || true
-#          # The above command will fail if no changes were present, so we ignore
-#          # the return code.
-#      - name: Push changes
-#        uses: ad-m/github-push-action@master
-#        with:
-#          branch: gh-pages
-#          directory: gh-pages
-#          github_token: ${{ secrets.GITHUB_TOKEN }}
+    - uses: actions/upload-artifact@v3
+      with:
+        name: DocumentationHTML
+        path: docs/_build/html/
+        retention-days: 30
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
@@ -11,7 +11,7 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
       with:
         fetch-depth: 0
     - uses: actions/setup-python@v4

diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
@@ -8,7 +8,7 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python
         uses: actions/setup-python@v4
         with:

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
@@ -8,17 +8,17 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.7', '3.8', '3.9', '3.10' ]
+        python-version: [ '3.8', '3.9', '3.10' ]
     name: Python ${{ matrix.python-version }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.x'
+          python-version: ${{ matrix.python-version }}
           cache: 'pip'
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
+          pip install pendulum
           pip install -r requirements.txt
       - name: Test with pytest
         run: |

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
     -   id: trailing-whitespace
     -   id: end-of-file-fixer
@@ -11,14 +11,14 @@ repos:
     -   id: check-added-large-files
 
 -   repo: https://github.com/psf/black
-    rev: 23.7.0
+    rev: 23.11.0
     hooks:
     -   id: black
         additional_dependencies: ['click!=8.1.0']
 
 -   repo: https://github.com/charliermarsh/ruff-pre-commit
     # Ruff version.
-    rev: 'v0.0.280'
+    rev: 'v0.1.6'
     hooks:
     -   id: ruff
         args: ["--fix"]
diff --git a/.ruff.toml b/.ruff.toml
@@ -42,7 +42,7 @@ line-length = 120
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 
 # Assume Python 3.7.
-target-version = "py37"
+target-version = "py38"
 
 [mccabe]
 # Unlike Flake8, default to a complexity level of 10.

diff --git a/README.md b/README.md
@@ -20,7 +20,10 @@ from data_dag.dag_factory import DagFactory
 
 from urllib.request import urlretrieve
 from typing import List
-from airflow.operators.dummy import DummyOperator
+try:
+    from airflow.operators.empty import EmptyOperator
+except ImportError:
+    from airflow.operators.dummy import DummyOperator as EmptyOperator
 from airflow.providers.http.sensors.http import HttpSensor
 from airflow.operators.python import PythonOperator
 from airflow.utils.task_group import TaskGroup
@@ -54,8 +57,8 @@ class DownloaderDag(DagFactory):
     downloads: List[DownloadOperator]
 
     def _make_dag(self):
-        start = DummyOperator(task_id='start')
-        end = DummyOperator(task_id='end')
+        start = EmptyOperator(task_id='start')
+        end = EmptyOperator(task_id='end')
 
         for download in self.downloads:
             start >> download.make_operator() >> end
@@ -68,7 +71,7 @@ Then a definition for a particular DAG can live in a data file:
 
 dag_id: sample_dag
 description: An example of how to write a data-driven DAG
-schedule_interval: '@daily'
+schedule: '@daily'
 start_date: '2020-01-01T00:00:00'
 downloads:
 - name: data
@@ -90,7 +93,7 @@ from my_factories.download import DownloaderDag
 with open('yaml/sample_dag.yaml', 'r') as f:
     dag_data = safe_load(f)
 
-dag = DownloaderDag.parse_obj(dag_data).make_dag()
+dag = DownloaderDag.model_validate(dag_data).make_dag()
 ```
 
 ![img.png](docs/_images/img.png)
@@ -115,7 +118,7 @@ for yaml_file_path in dag_dir.glob('typical_dags/**.yml'):
         dag_metadata = yaml.safe_load(f)
 
     # ... generate a DAG from that metadata
-    dag_metadata_obj = BaseDag.parse_obj(dag_metadata)
+    dag_metadata_obj = BaseDag.model_validate(dag_metadata)
     dag = dag_metadata_obj.make_dag()
 
     # See https://www.astronomer.io/guides/dynamically-generating-dags/

diff --git a/data_dag/__init__.py b/data_dag/__init__.py
@@ -1,3 +1,3 @@
 """Tooling to help build data-driven DAGs"""
 
-__version__ = "0.3.0"
+__version__ = "1.0.0"
diff --git a/data_dag/dag_factory.py b/data_dag/dag_factory.py
@@ -3,7 +3,7 @@
 from typing import Dict, Iterable, List, Optional, Union
 
 from airflow.models.dag import DAG
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict
 
 
 class DagFactory(BaseModel, abc.ABC):
@@ -16,9 +16,9 @@ class DagFactory(BaseModel, abc.ABC):
 
         class MyKindOfDag(DagFactory):
             def _make_dag(self):
-                start = DummyOperator(...)
+                start = EmptyOperator(...)
                 do_something = PythonOperator(...)
-                end = DummyOperator(...)
+                end = EmptyOperator(...)
 
                 start >> do_something >> end
 
@@ -38,7 +38,7 @@ def _make_dag(self):
 
     dag_id: str
     description: Optional[str] = None
-    schedule_interval: Union[timedelta, str, None] = None
+    schedule: Union[timedelta, str, None] = None
     start_date: Optional[datetime] = None
     end_date: Optional[datetime] = None
     full_filepath: Optional[str] = None
@@ -60,9 +60,7 @@ def _make_dag(self):
     jinja_environment_kwargs: Optional[Dict] = None
     render_template_as_native_obj: bool = False
     tags: Optional[List[str]] = None
-
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(extra="forbid")
 
     @property
     def default_dag_kwargs(self) -> Dict:
@@ -79,7 +77,7 @@ def make_dag_object(self, **overrides) -> DAG:
         kwargs.update(
             {
                 field: getattr(self, field, None)
-                for field in DagFactory.__fields__
+                for field in DagFactory.model_fields
                 if getattr(self, field, None) is not None
             }
         )

diff --git a/data_dag/operator_factory/__init__.py b/data_dag/operator_factory/__init__.py
@@ -3,6 +3,7 @@
     OperatorFactory,
 )
 from .dynamic import (
+    DynamicOperatorComponent,
     DynamicOperatorFactory,
 )
 from .simple import (
@@ -16,4 +17,5 @@
     "SimpleOperatorFactory",
     "SimpleOperatorComponent",
     "DynamicOperatorFactory",
+    "DynamicOperatorComponent",
 )
diff --git a/data_dag/operator_factory/base.py b/data_dag/operator_factory/base.py
@@ -4,12 +4,11 @@
 
 from airflow.models.taskmixin import TaskMixin
 from airflow.utils.task_group import TaskGroup
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict
 
 
 class BaseOperatorFactory(BaseModel, abc.ABC):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(extra="forbid")
 
     def make_operator(
         self, *args, **kwargs
@@ -25,7 +24,7 @@ def make_operator(
 class OperatorFactory(BaseOperatorFactory, abc.ABC):
     """An interface for writing operator factories."""
 
-    task_id: Optional[str]
+    task_id: Optional[str] = None
 
     @property
     def default_task_id(self) -> str:
@@ -68,5 +67,4 @@ def _make_operators(self, *args, **kwargs) -> None:
 class OperatorComponent(BaseModel, abc.ABC):
     """A non-operator component for use in other operator factories. Just a proxy for :py:class:`pydantic.BaseModel`."""
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(extra="forbid")
File	Stmts	Miss	Cover	Missing
data_dag
dag_factory.py	46	10	78%	69, 76–85, 89–93, 97
data_dag/operator_factory
__init__.py	4	2	50%	9–14
base.py	30	10	67%	21, 32, 36, 40–43, 56–59, 64
dynamic.py	37	24	35%	19–25, 29–95
simple.py	37	37	0%	1–110
TOTAL	155	83	46%