[dagster-airlift] filter (#28606)

## Summary & Motivation
Write up a simple filter abstraction which limits the number of dags
that we retrieve up front. This would be for perf reasons.
## How I Tested These Changes

## Changelog
- Added `AirflowFilter` API for use with `dagster-airlift`, allows you
to filter down the set of dags retrieved up front for perf improvements.

Author: dpeng817

Diff for: docs/docs/guides/migrate/airflow-to-dagster/federation/observe.md

@@ -22,7 +22,7 @@ For a full list of `dagster-airlift` classes and methods, see the [API docs](htt
 
 ## Observe the `warehouse` Airflow instance
 
-Next, declare a reference to the `warehouse` Airflow instance, which is running at `http://localhost:8081`:
+Next, in your `airlift_federation_tutorial/dagster_defs/definitions.py` file, declare a reference to the `warehouse` Airflow instance, which is running at `http://localhost:8081`:
 
 <CodeExample
   path="airlift-federation-tutorial/snippets/observe.py"

Diff for: examples/airlift-federation-tutorial/Makefile

@@ -27,7 +27,8 @@ help:
 airflow_install:
 	pip install uv && \
 	uv pip install dagster-airlift[tutorial] && \
-	uv pip install -e $(MAKEFILE_DIR)
+	uv pip install -e $(MAKEFILE_DIR) && \
+	uv pip install duckdb pandas
 
 airflow_setup: wipe
 	mkdir -p $$WAREHOUSE_AIRFLOW_HOME

Diff for: examples/airlift-federation-tutorial/scripts/airflow_setup.sh

@@ -17,6 +17,9 @@ if [[ "$DAGS_FOLDER" != /* ]] || [[ "$AIRFLOW_HOME_DIR" != /* ]]; then
   exit 1
 fi
 
+# Generate a unique secret key
+SECRET_KEY=$(openssl rand -hex 30)
+
 # Create the airflow.cfg file in the specified AIRFLOW_HOME_DIR
 cat <<EOL > $AIRFLOW_HOME_DIR/airflow.cfg
 [core]
@@ -28,7 +31,106 @@ auth_backend = airflow.api.auth.backend.basic_auth
 [webserver]
 expose_config = True
 web_server_port = $PORT
+secret_key = $SECRET_KEY
+
+EOL
+
+# Create the webserver_config.py file
+cat <<EOL > $AIRFLOW_HOME_DIR/webserver_config.py
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+from flask_appbuilder.security.manager import AUTH_DB
+
+# from airflow.www.fab_security.manager import AUTH_LDAP
+
+basedir = os.path.abspath(os.path.dirname(__file__))
+
+# The SQLAlchemy connection string.
+SQLALCHEMY_DATABASE_URI = 'sqlite:///' + os.path.join(basedir, 'webserver.db')
+
+# Flask-WTF flag for CSRF
+WTF_CSRF_ENABLED = True
+
+# Use a unique cookie name for this instance based on port
+SESSION_COOKIE_NAME = "airflow_session_${PORT}"
+
+# ----------------------------------------------------
+# AUTHENTICATION CONFIG
+# ----------------------------------------------------
+# For details on how to set up each of the following authentication, see
+# http://flask-appbuilder.readthedocs.io/en/latest/security.html
+
+# The authentication type
+# AUTH_OID : Is for OpenID
+# AUTH_DB : Is for database
+# AUTH_LDAP : Is for LDAP
+# AUTH_REMOTE_USER : Is for using REMOTE_USER from web server
+# AUTH_OAUTH : Is for OAuth
+AUTH_TYPE = AUTH_DB
+
+# When using LDAP Auth, setup the ldap server
+# LDAP_SERVER = "ldap://ldapserver.new"
+# LDAP_PORT = 389
+# LDAP_USE_TLS = False
+# LDAP_SEARCH_SCOPE = "LEVEL"
+# LDAP_BIND_USER = "uid=admin,ou=users,dc=example,dc=com"
+# LDAP_BIND_PASSWORD = "admin_password"
+# LDAP_BASEDN = "dc=example,dc=com"
+# LDAP_USER_DN = "ou=users"
+# LDAP_USER_FILTER = "(uid=%s)"
+# LDAP_GROUP_DN = "ou=groups"
+# LDAP_GROUP_FILTER = "(member=%s)"
+# LDAP_USER_NAME_FORMAT = "uid=%s,ou=users,dc=example,dc=com"
+# LDAP_GROUP_NAME_FORMAT = "cn=%s,ou=groups,dc=example,dc=com"
+
+# Uncomment to setup Full admin role name
+# AUTH_ROLE_ADMIN = 'Admin'
+
+# Uncomment to setup Public role name, no authentication needed
+# AUTH_ROLE_PUBLIC = 'Public'
+
+# Will allow user self registration
+# AUTH_USER_REGISTRATION = True
+
+# The default user self registration role
+# AUTH_USER_REGISTRATION_ROLE = "Public"
 
+# When using OAuth Auth, uncomment to setup provider(s) info
+# Google OAuth example:
+# OAUTH_PROVIDERS = [{
+#   'name':'google',
+#     'token_key':'access_token',
+#     'icon':'fa-google',
+#         'remote_app': {
+#             'api_base_url':'https://www.googleapis.com/oauth2/v2/',
+#             'client_kwargs':{
+#                 'scope': 'email profile'
+#             },
+#             'access_token_url':'https://accounts.google.com/o/oauth2/token',
+#             'authorize_url':'https://accounts.google.com/o/oauth2/auth',
+#             'request_token_url': None,
+#             'client_id': GOOGLE_KEY,
+#             'client_secret': GOOGLE_SECRET_KEY,
+#         }
+# }]
 EOL
 
 # call airflow command to create the default user

Diff for: python_modules/libraries/dagster-airlift/dagster_airlift/core/airflow_instance.py

@@ -12,6 +12,7 @@
 from dagster._record import record
 from dagster._time import get_current_datetime
 
+from dagster_airlift.core.filter import AirflowFilter
 from dagster_airlift.core.serialization.serialized_data import DagInfo, TaskInfo
 
 TERMINAL_STATES = {"success", "failed", "canceled"}
@@ -79,13 +80,21 @@ def normalized_name(self) -> str:
     def get_api_url(self) -> str:
         return f"{self.auth_backend.get_webserver_url()}/api/v1"
 
-    def list_dags(self) -> list["DagInfo"]:
+    def list_dags(self, retrieval_filter: Optional[AirflowFilter] = None) -> list["DagInfo"]:
+        retrieval_filter = retrieval_filter or AirflowFilter()
         dag_responses = []
         webserver_url = self.auth_backend.get_webserver_url()
+
         while True:
+            params = retrieval_filter.augment_request_params(
+                {
+                    "limit": self.dag_list_limit,
+                    "offset": len(dag_responses),
+                }
+            )
             response = self.auth_backend.get_session().get(
                 f"{self.get_api_url()}/dags",
-                params={"limit": self.dag_list_limit, "offset": len(dag_responses)},
+                params=params,
             )
             if response.status_code == 200:
                 dags = response.json()

Diff for: python_modules/libraries/dagster-airlift/dagster_airlift/core/filter.py

@@ -0,0 +1,25 @@
+from collections.abc import Sequence
+from typing import Optional
+
+from dagster._record import record
+
+
+@record
+class AirflowFilter:
+    """Filters the set of Airflow objects to fetch.
+
+    Args:
+        dag_id_ilike (Optional[str]): A pattern used to match the set of dag_ids to retrieve. Uses the sql ILIKE operator Airflow-side.
+        airflow_tags (Optional[Sequence[str]]): Filters down to the set of Airflow DAGs whcih contain the particular tags provided.
+    """
+
+    dag_id_ilike: Optional[str] = None
+    airflow_tags: Optional[Sequence[str]] = None
+
+    def augment_request_params(self, request_params: dict) -> dict:
+        new_request_params = request_params.copy()
+        if self.dag_id_ilike is not None:
+            new_request_params["dag_id_pattern"] = self.dag_id_ilike
+        if self.airflow_tags is not None:
+            new_request_params["tags"] = self.airflow_tags
+        return new_request_params

Diff for: python_modules/libraries/dagster-airlift/dagster_airlift/core/load_defs.py

@@ -15,6 +15,7 @@
 
 from dagster_airlift.core.airflow_defs_data import MappedAsset
 from dagster_airlift.core.airflow_instance import AirflowInstance
+from dagster_airlift.core.filter import AirflowFilter
 from dagster_airlift.core.sensor.event_translation import (
     DagsterEventTransformerFn,
     default_event_transformer,
@@ -38,6 +39,7 @@
 @dataclass
 class AirflowInstanceDefsLoader(StateBackedDefinitionsLoader[SerializedAirflowDefinitionsData]):
     airflow_instance: AirflowInstance
+    retrieval_filter: AirflowFilter
     mapped_assets: Sequence[MappedAsset]
     source_code_retrieval_enabled: Optional[bool]
     sensor_minimum_interval_seconds: int = DEFAULT_AIRFLOW_SENSOR_INTERVAL_SECONDS
@@ -54,6 +56,7 @@ def fetch_state(self) -> SerializedAirflowDefinitionsData:
             dag_selector_fn=self.dag_selector_fn,
             automapping_enabled=False,
             source_code_retrieval_enabled=self.source_code_retrieval_enabled,
+            retrieval_filter=self.retrieval_filter,
         )
 
     def defs_from_state(  # pyright: ignore[reportIncompatibleMethodOverride]
@@ -77,6 +80,7 @@ def build_defs_from_airflow_instance(
     dag_selector_fn: Optional[Callable[[DagInfo], bool]] = None,
     source_code_retrieval_enabled: Optional[bool] = None,
     default_sensor_status: Optional[DefaultSensorStatus] = None,
+    retrieval_filter: Optional[AirflowFilter] = None,
 ) -> Definitions:
     """Builds a :py:class:`dagster.Definitions` object from an Airflow instance.
 
@@ -221,6 +225,7 @@ def only_include_dag(dag_info: DagInfo) -> bool:
         mapped_assets=mapped_assets,
         dag_selector_fn=dag_selector_fn,
         source_code_retrieval_enabled=source_code_retrieval_enabled,
+        retrieval_filter=retrieval_filter or AirflowFilter(),
     ).get_or_fetch_state()
     mapped_and_constructed_assets = [
         *_apply_airflow_data_to_specs(mapped_assets, serialized_airflow_data),
@@ -292,12 +297,14 @@ def enrich_airflow_mapped_assets(
     mapped_assets: Sequence[MappedAsset],
     airflow_instance: AirflowInstance,
     source_code_retrieval_enabled: Optional[bool],
+    retrieval_filter: Optional[AirflowFilter] = None,
 ) -> Sequence[AssetsDefinition]:
     """Enrich Airflow-mapped assets with metadata from the provided :py:class:`AirflowInstance`."""
     serialized_data = AirflowInstanceDefsLoader(
         airflow_instance=airflow_instance,
         mapped_assets=mapped_assets,
         source_code_retrieval_enabled=source_code_retrieval_enabled,
+        retrieval_filter=retrieval_filter or AirflowFilter(),
     ).get_or_fetch_state()
     return list(_apply_airflow_data_to_specs(mapped_assets, serialized_data))
 
@@ -308,12 +315,14 @@ def load_airflow_dag_asset_specs(
     mapped_assets: Optional[Sequence[MappedAsset]] = None,
     dag_selector_fn: Optional[Callable[[DagInfo], bool]] = None,
     source_code_retrieval_enabled: Optional[bool] = None,
+    retrieval_filter: Optional[AirflowFilter] = None,
 ) -> Sequence[AssetSpec]:
     """Load asset specs for Airflow DAGs from the provided :py:class:`AirflowInstance`, and link upstreams from mapped assets."""
     serialized_data = AirflowInstanceDefsLoader(
         airflow_instance=airflow_instance,
         mapped_assets=mapped_assets or [],
         dag_selector_fn=dag_selector_fn,
         source_code_retrieval_enabled=source_code_retrieval_enabled,
+        retrieval_filter=retrieval_filter or AirflowFilter(),
     ).get_or_fetch_state()
     return list(spec_iterator(construct_dag_assets_defs(serialized_data)))

Diff for: python_modules/libraries/dagster-airlift/dagster_airlift/core/serialization/compute.py

@@ -8,6 +8,7 @@
 
 from dagster_airlift.core.airflow_instance import AirflowInstance, DagInfo
 from dagster_airlift.core.dag_asset import get_leaf_assets_for_dag
+from dagster_airlift.core.filter import AirflowFilter
 from dagster_airlift.core.serialization.serialized_data import (
     DagHandle,
     KeyScopedDagHandles,
@@ -144,10 +145,11 @@ def fetch_all_airflow_data(
     mapping_info: AirliftMetadataMappingInfo,
     dag_selector_fn: Optional[DagSelectorFn],
     automapping_enabled: bool,
+    retrieval_filter: AirflowFilter,
 ) -> FetchedAirflowData:
     dag_infos = {
         dag.dag_id: dag
-        for dag in airflow_instance.list_dags()
+        for dag in airflow_instance.list_dags(retrieval_filter=retrieval_filter)
         if dag_selector_fn is None or dag_selector_fn(dag)
     }
     # To limit the number of API calls, only fetch task infos for the dags that we absolutely have to.
@@ -185,10 +187,15 @@ def compute_serialized_data(
     dag_selector_fn: Optional[DagSelectorFn],
     automapping_enabled: bool,
     source_code_retrieval_enabled: Optional[bool],
+    retrieval_filter: AirflowFilter,
 ) -> "SerializedAirflowDefinitionsData":
     mapping_info = build_airlift_metadata_mapping_info(mapped_assets)
     fetched_airflow_data = fetch_all_airflow_data(
-        airflow_instance, mapping_info, dag_selector_fn, automapping_enabled=automapping_enabled
+        airflow_instance,
+        mapping_info,
+        dag_selector_fn,
+        automapping_enabled=automapping_enabled,
+        retrieval_filter=retrieval_filter,
     )
     source_code_retrieval_enabled = infer_code_retrieval_enabled(
         source_code_retrieval_enabled, fetched_airflow_data

Diff for: python_modules/libraries/dagster-airlift/dagster_airlift/test/airflow_test_instance.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from datetime import datetime, timedelta
 from typing import Any, Optional
 
@@ -8,6 +8,7 @@
 from dagster_airlift.core import AirflowInstance
 from dagster_airlift.core.airflow_instance import DagInfo, DagRun, TaskInfo, TaskInstance
 from dagster_airlift.core.basic_auth import AirflowAuthBackend
+from dagster_airlift.core.filter import AirflowFilter
 
 
 class DummyAuthBackend(AirflowAuthBackend):
@@ -18,6 +19,9 @@ def get_webserver_url(self) -> str:
         return "http://dummy.domain"
 
 
+DEFAULT_FAKE_INSTANCE_NAME = "test_instance"
+
+
 class AirflowInstanceFake(AirflowInstance):
     """Loads a set of provided DagInfo, TaskInfo, and TaskInstance objects for testing."""
 
@@ -50,11 +54,29 @@ def __init__(
         self._max_runs_per_batch = max_runs_per_batch
         super().__init__(
             auth_backend=DummyAuthBackend(),
-            name="test_instance" if instance_name is None else instance_name,
+            name=DEFAULT_FAKE_INSTANCE_NAME if instance_name is None else instance_name,
         )
 
-    def list_dags(self) -> list[DagInfo]:
-        return list(self._dag_infos_by_dag_id.values())
+    def list_dags(self, retrieval_filter: Optional[AirflowFilter] = None) -> list[DagInfo]:
+        retrieval_filter = retrieval_filter or AirflowFilter()
+        # Very basic filtering for testing purposes
+        dags_to_retrieve = list(self._dag_infos_by_dag_id.values())
+        if retrieval_filter.dag_id_ilike:
+            dags_to_retrieve = [
+                dag_info
+                for dag_info in dags_to_retrieve
+                if retrieval_filter.dag_id_ilike in dag_info.dag_id
+            ]
+        if retrieval_filter.airflow_tags:
+            dags_to_retrieve = [
+                dag_info
+                for dag_info in dags_to_retrieve
+                if all(
+                    tag in dag_info.metadata.get("tags", [])
+                    for tag in retrieval_filter.airflow_tags
+                )
+            ]
+        return dags_to_retrieve
 
     def list_variables(self) -> list[dict[str, Any]]:
         return self._variables
@@ -150,11 +172,13 @@ def get_dag_source_code(self, file_token: str) -> str:
         return "indicates found source code"
 
 
-def make_dag_info(dag_id: str, file_token: Optional[str]) -> DagInfo:
+def make_dag_info(
+    instance_name: str, dag_id: str, file_token: Optional[str], dag_props: Mapping[str, Any]
+) -> DagInfo:
     return DagInfo(
         webserver_url="http://dummy.domain",
         dag_id=dag_id,
-        metadata={"file_token": file_token if file_token else "dummy_file_token"},
+        metadata={"file_token": file_token if file_token else "dummy_file_token", **dag_props},
     )
 
 
@@ -220,6 +244,7 @@ def make_instance(
     task_deps: dict[str, list[str]] = {},
     instance_name: Optional[str] = None,
     max_runs_per_batch: Optional[int] = None,
+    dag_props: dict[str, Any] = {},
 ) -> AirflowInstanceFake:
     """Constructs DagInfo, TaskInfo, and TaskInstance objects from provided data.
 
@@ -231,7 +256,12 @@ def make_instance(
     dag_infos = []
     task_infos = []
     for dag_id, task_ids in dag_and_task_structure.items():
-        dag_info = make_dag_info(dag_id=dag_id, file_token=dag_id)
+        dag_info = make_dag_info(
+            instance_name=instance_name or DEFAULT_FAKE_INSTANCE_NAME,
+            dag_id=dag_id,
+            file_token=dag_id,
+            dag_props=dag_props.get(dag_id, {}),
+        )
         dag_infos.append(dag_info)
         task_infos.extend(
             [