dagster-census 0.28.0__py3-none-any.whl → 0.28.2__py3-none-any.whl

dagster_census/__init__.py

@@ -1,7 +1,8 @@
 from dagster_shared.libraries import DagsterLibraryRegistry
 
+from dagster_census.components.census_component import CensusComponent as CensusComponent
 from dagster_census.ops import census_trigger_sync_op
-from dagster_census.resources import CensusResource, census_resource
+from dagster_census.resources import CensusResource
 from dagster_census.types import CensusOutput
 from dagster_census.version import __version__
 
@@ -10,6 +11,5 @@ DagsterLibraryRegistry.register("dagster-census", __version__)
 __all__ = [
     "CensusOutput",
     "CensusResource",
-    "census_resource",
     "census_trigger_sync_op",
 ]

dagster_census/components/census_component.py

@@ -0,0 +1,219 @@
+from collections.abc import Callable, Iterable, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Annotated, Optional, Union
+
+import dagster as dg
+import pydantic
+from dagster._annotations import public
+from dagster._core.definitions.metadata import TableMetadataSet
+from dagster._utils.names import clean_name
+from dagster.components.component.state_backed_component import StateBackedComponent
+from dagster.components.resolved.base import resolve_fields
+from dagster.components.utils.defs_state import (
+    DefsStateConfig,
+    DefsStateConfigArgs,
+    ResolvedDefsStateConfig,
+)
+from dagster_shared import check
+from dagster_shared.serdes.serdes import deserialize_value
+
+from dagster_census.components.census_scaffolder import CensusComponentScaffolder
+from dagster_census.resources import CensusResource
+from dagster_census.translator import (
+    CensusMetadataSet,
+    CensusSync,
+    CensusWorkspaceData,
+    generate_table_schema,
+)
+
+
+class CensusSyncSelectorByName(dg.Resolvable, dg.Model):
+    by_name: Annotated[
+        Sequence[str],
+        pydantic.Field(..., description="A list of sync names to include in the collection."),
+    ]
+
+
+class CensusSyncSelectorById(dg.Resolvable, dg.Model):
+    by_id: Annotated[
+        Sequence[Union[int, str]],  # Allow strings for use-cases like '{{ env.ENV_VAR }}'
+        pydantic.Field(..., description="A list of sync IDs to include in the collection."),
+    ]
+
+
+def resolve_sync_selector(
+    context: dg.ResolutionContext, model
+) -> Optional[Callable[[CensusSync], bool]]:
+    if isinstance(model, str):
+        resolved = context.resolve_value(model)
+        resolved = check.callable_param(resolved, "unknown")  # pyright: ignore[reportArgumentType]
+        return resolved
+    if isinstance(model, CensusSyncSelectorByName.model()):
+        resolved = resolve_fields(model, CensusSyncSelectorByName.model(), context)
+        return lambda sync: sync.name in resolved["by_name"]
+    elif isinstance(model, CensusSyncSelectorById.model()):
+        resolved = resolve_fields(model, CensusSyncSelectorById.model(), context)
+        return lambda sync: sync.id in resolved["by_id"]
+    else:
+        check.failed(f"Unknown sync target type: {type(model)}")
+
+
+@dataclass
+class CensusResourceArgs(dg.Resolvable, dg.Model):
+    """The fields are analogous to the fields at `CensusResource`."""
+
+    api_key: Annotated[str, pydantic.Field(...)]
+    request_max_retries: Annotated[int, pydantic.Field(None)]
+    request_retry_delay: Annotated[float, pydantic.Field(None)]
+    request_timeout: Annotated[int, pydantic.Field(None)]
+
+
+def resolve_census_workspace(
+    context: dg.ResolutionContext, model: CensusResourceArgs
+) -> CensusResource:
+    return CensusResource(**resolve_fields(model, CensusResourceArgs, context))
+
+
+@public
+@dataclass
+@dg.scaffold_with(CensusComponentScaffolder)
+class CensusComponent(StateBackedComponent, dg.Resolvable):
+    """Loads Census syncs from a Census workspace as Dagster assets.
+    Materializing these assets will trigger the Census sync, enabling
+    you to schedule Census syncs using Dagster.
+
+    Example:
+
+        .. code-block:: yaml
+
+            # defs.yaml
+
+            type: dagster_census.CensusComponent
+            attributes:
+              workspace:
+                api_key: "{{ env.CENSUS_API_KEY }}"
+              sync_selector:
+                by_name:
+                  - my_first_sync
+                  - my_second_sync
+    """
+
+    workspace: Annotated[
+        CensusResource,
+        dg.Resolver(resolve_census_workspace, model_field_type=CensusResourceArgs),
+    ]
+
+    sync_selector: Annotated[
+        Optional[Callable[[CensusSync], bool]],
+        dg.Resolver(
+            resolve_sync_selector,
+            model_field_type=Union[
+                str, CensusSyncSelectorByName.model(), CensusSyncSelectorById.model()
+            ],
+            description="Function used to select Census syncs to pull into Dagster.",
+        ),
+    ] = None
+
+    defs_state: ResolvedDefsStateConfig = field(
+        default_factory=DefsStateConfigArgs.local_filesystem
+    )
+
+    @property
+    def defs_state_config(self) -> DefsStateConfig:
+        default_key = f"{self.__class__.__name__}"
+        return DefsStateConfig.from_args(self.defs_state, default_key=default_key)
+
+    def write_state_to_path(self, state_path: Path) -> None:
+        state = self.workspace.fetch_census_workspace_data()
+        state_path.write_text(dg.serialize_value(state))
+
+    @public
+    def get_asset_spec(self, sync: CensusSync) -> dg.AssetSpec:
+        metadata = {
+            **TableMetadataSet(
+                column_schema=generate_table_schema(sync.mappings),
+                table_name=sync.name,
+            ),
+            **CensusMetadataSet(
+                sync_id=sync.id,
+                sync_name=sync.name,
+                source_id=sync.source_id,
+                destination_id=sync.destination_id,
+            ),
+        }
+
+        return dg.AssetSpec(
+            key=dg.AssetKey(clean_name(sync.name)),
+            metadata=metadata,
+            description=f"Asset generated from Census sync {sync.id}",
+            kinds={"census"},
+        )
+
+    @public
+    def execute(
+        self, context: dg.AssetExecutionContext, census: CensusResource
+    ) -> Iterable[Union[dg.AssetMaterialization, dg.MaterializeResult]]:
+        """Executes a Census sync for the selected sync.
+
+        This method can be overridden in a subclass to customize the sync execution behavior,
+        such as adding custom logging or handling sync results differently.
+
+        Args:
+            context: The asset execution context provided by Dagster
+            census: The CensusResource used to trigger and monitor syncs
+
+        Returns:
+            MaterializeResult event from the Census sync
+
+        Example:
+            Override this method to add custom logging during sync execution:
+
+            .. code-block:: python
+
+                from dagster_census import CensusComponent
+                import dagster as dg
+
+                class CustomCensusComponent(CensusComponent):
+                    def execute(self, context, census):
+                        context.log.info(f"Starting Census sync for {context.asset_key}")
+                        result = super().execute(context, census)
+                        context.log.info("Census sync completed successfully")
+                        return result
+        """
+        # Select the first asset-spec (since there is only one)
+        spec = next(iter(context.assets_def.specs))
+
+        census_metadataset = CensusMetadataSet.extract(spec.metadata)
+        census.trigger_sync_and_poll(sync_id=census_metadataset.sync_id)
+        yield dg.AssetMaterialization(asset_key=clean_name(census_metadataset.sync_name))
+
+    def _load_asset_specs(self, state: CensusWorkspaceData) -> Sequence[dg.AssetSpec]:
+        connection_selector_fn = self.sync_selector or (lambda sync: True)
+        return [self.get_asset_spec(sync) for sync in state.syncs if connection_selector_fn(sync)]
+
+    def _get_census_assets_def(self, sync_name: str, spec: dg.AssetSpec) -> dg.AssetsDefinition:
+        @dg.multi_asset(
+            name=f"census_{clean_name(sync_name)}",
+            can_subset=True,
+            specs=[spec],
+        )
+        def _asset(context: dg.AssetExecutionContext):
+            yield from self.execute(context, self.workspace)
+
+        return _asset
+
+    def build_defs_from_state(
+        self, context: dg.ComponentLoadContext, state_path: Optional[Path]
+    ) -> dg.Definitions:
+        if state_path is None:
+            return dg.Definitions()
+
+        state = deserialize_value(state_path.read_text(), CensusWorkspaceData)
+
+        assets = [
+            self._get_census_assets_def(CensusMetadataSet.extract(spec.metadata).sync_name, spec)
+            for spec in self._load_asset_specs(state)
+        ]
+
+        return dg.Definitions(assets=assets)

dagster_census/components/census_scaffolder.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from dagster.components.component.component_scaffolder import Scaffolder
+from dagster.components.component_scaffolding import scaffold_component
+from dagster.components.scaffold.scaffold import ScaffoldRequest
+from pydantic import BaseModel
+
+
+class CensusScaffolderParams(BaseModel):
+    api_key: Optional[str] = None
+
+
+class CensusComponentScaffolder(Scaffolder[CensusScaffolderParams]):
+    @classmethod
+    def get_scaffold_params(cls) -> type[CensusScaffolderParams]:
+        return CensusScaffolderParams
+
+    def scaffold(self, request: ScaffoldRequest[CensusScaffolderParams]) -> None:
+        scaffold_component(
+            request,
+            {
+                "workspace": {
+                    "api_key": "{{ env.CENSUS_API_KEY }}",
+                }
+            },
+        )

dagster_census/resources.py

@@ -5,12 +5,14 @@ import time
 from collections.abc import Mapping
 from typing import Any, Optional
 
+import pydantic
 import requests
-from dagster import Failure, Field, StringSource, __version__, get_dagster_logger, resource
-from dagster._core.definitions.resource_definition import dagster_maintained_resource
+from dagster import ConfigurableResource, Failure, __version__, get_dagster_logger
+from dagster_shared.utils.cached_method import cached_method
 from requests.auth import HTTPBasicAuth
 from requests.exceptions import RequestException
 
+from dagster_census.translator import CensusSync, CensusWorkspaceData
 from dagster_census.types import CensusOutput
 
 CENSUS_API_BASE = "app.getcensus.com/api"
@@ -21,35 +23,68 @@ DEFAULT_POLL_INTERVAL = 10
 SYNC_RUN_STATUSES = {"completed", "failed", "queued", "skipped", "working"}
 
 
-class CensusResource:
-    """This class exposes methods on top of the Census REST API."""
+class CensusResource(ConfigurableResource):
+    """This resource allows users to programatically interface with the Census REST API to launch
+    syncs and monitor their progress. This currently implements only a subset of the functionality
+    exposed by the API.
 
-    def __init__(
-        self,
-        api_key: str,
-        request_max_retries: int = 3,
-        request_retry_delay: float = 0.25,
-        log: logging.Logger = get_dagster_logger(),
-    ):
-        self.api_key = api_key
+    **Examples:**
 
-        self._request_max_retries = request_max_retries
-        self._request_retry_delay = request_retry_delay
+    .. code-block:: python
 
-        self._log = log
+        import dagster as dg
+        from dagster_census import CensusResource
 
-    @property
-    def _api_key(self):
-        if self.api_key.startswith("secret-token:"):
-            return self.api_key
-        return "secret-token:" + self.api_key
+        census_resource = CensusResource(
+            api_key=dg.EnvVar("CENSUS_API_KEY")
+        )
+
+        @dg.asset
+        def census_sync_asset(census: CensusResource):
+            census.trigger_sync_and_poll(sync_id=123456)
+
+        defs = dg.Definitions(
+            assets=[census_sync_asset],
+            resources={"census": census_resource}
+        )
+    """
+
+    api_key: str = pydantic.Field(..., description="The Census API key")
+    request_max_retries: int = pydantic.Field(
+        default=3,
+        description=(
+            "The maximum number of times requests to the Census API should be retried "
+            "before failing."
+        ),
+    )
+    request_retry_delay: float = pydantic.Field(
+        default=0.25,
+        description="Time (in seconds) to wait between each request retry.",
+    )
+    request_timeout: int = pydantic.Field(
+        default=15,
+        description="Time (in seconds) after which the requests to Census are declared timed out.",
+    )
 
     @property
     def api_base_url(self) -> str:
         return f"https://{CENSUS_API_BASE}/{CENSUS_VERSION}"
 
+    @classmethod
+    def _is_dagster_maintained(cls) -> bool:
+        return True
+
+    @property
+    @cached_method
+    def _log(self) -> logging.Logger:
+        return get_dagster_logger()
+
     def make_request(
-        self, method: str, endpoint: str, data: Optional[str] = None
+        self,
+        method: str,
+        endpoint: str,
+        data: Optional[str] = None,
+        page_number: Optional[int] = None,
     ) -> Mapping[str, Any]:
         """Creates and sends a request to the desired Census API endpoint.
 
@@ -67,6 +102,11 @@ class CensusResource:
             "Content-Type": "application/json;version=2",
         }
 
+        if page_number is not None:
+            params = {"page": page_number}
+        else:
+            params = {}
+
         num_retries = 0
         while True:
             try:
@@ -74,19 +114,20 @@ class CensusResource:
                     method=method,
                     url=url,
                     headers=headers,
-                    auth=HTTPBasicAuth("bearer", self._api_key),
+                    auth=HTTPBasicAuth("bearer", self.api_key),
                     data=data,
+                    params=params,
                 )
                 response.raise_for_status()
                 return response.json()
             except RequestException as e:
                 self._log.error("Request to Census API failed: %s", e)
-                if num_retries == self._request_max_retries:
+                if num_retries == self.request_max_retries:
                     break
                 num_retries += 1
-                time.sleep(self._request_retry_delay)
+                time.sleep(self.request_retry_delay)
 
-        raise Failure(f"Max retries ({self._request_max_retries}) exceeded with url: {url}.")
+        raise Failure(f"Max retries ({self.request_max_retries}) exceeded with url: {url}.")
 
     def get_sync(self, sync_id: int) -> Mapping[str, Any]:
         """Gets details about a given sync from the Census API.
@@ -250,57 +291,32 @@ class CensusResource:
             destination=destination_details,
         )
 
+    @cached_method
+    def fetch_census_workspace_data(self) -> CensusWorkspaceData:
+        """Retrieves all Census syncs from the workspace and returns it as a CensusWorkspaceData object.
 
-@dagster_maintained_resource
-@resource(
-    config_schema={
-        "api_key": Field(
-            StringSource,
-            is_required=True,
-            description="Census API Key.",
-        ),
-        "request_max_retries": Field(
-            int,
-            default_value=3,
-            description=(
-                "The maximum number of times requests to the Census API should be retried "
-                "before failing."
-            ),
-        ),
-        "request_retry_delay": Field(
-            float,
-            default_value=0.25,
-            description="Time (in seconds) to wait between each request retry.",
-        ),
-    },
-    description="This resource helps manage Census connectors",
-)
-def census_resource(context) -> CensusResource:
-    """This resource allows users to programatically interface with the Census REST API to launch
-    syncs and monitor their progress. This currently implements only a subset of the functionality
-    exposed by the API.
-
-    **Examples:**
-
-    .. code-block:: python
-
-        from dagster import job
-        from dagster_census import census_resource
-
-        my_census_resource = census_resource.configured(
-            {
-                "api_key": {"env": "CENSUS_API_KEY"},
-            }
+        Returns:
+            CensusWorkspaceData: A snapshot of the Census workspace's syncs.
+        """
+        all_syncs = []
+        page = self.make_request(method="GET", endpoint="syncs")
+
+        last_page_number = page["pagination"]["last_page"]
+        for sync in page["data"]:
+            all_syncs.append(self._census_sync_struct_from_json(sync))
+
+        for i in range(2, last_page_number + 1):
+            page = self.make_request(method="GET", endpoint="syncs", page_number=i)
+            for sync in page["data"]:
+                all_syncs.append(self._census_sync_struct_from_json(sync))
+
+        return CensusWorkspaceData(syncs=all_syncs)
+
+    def _census_sync_struct_from_json(self, sync: dict[str, Any]) -> CensusSync:
+        return CensusSync(
+            id=sync["id"],
+            name=sync["label"] or sync["resource_identifier"],
+            source_id=sync["source_attributes"]["connection_id"],
+            destination_id=sync["destination_attributes"]["connection_id"],
+            mappings=sync["mappings"],
         )
-
-        @job(resource_defs={"census":my_census_resource})
-        def my_census_job():
-            ...
-
-    """
-    return CensusResource(
-        api_key=context.resource_config["api_key"],
-        request_max_retries=context.resource_config["request_max_retries"],
-        request_retry_delay=context.resource_config["request_retry_delay"],
-        log=context.log,
-    )

dagster_census/translator.py

@@ -0,0 +1,57 @@
+from collections.abc import Mapping
+from typing import Any
+
+import dagster as dg
+from dagster._core.definitions.metadata.metadata_set import NamespacedMetadataSet
+from dagster._record import record
+from dagster_shared.serdes import whitelist_for_serdes
+
+
+@whitelist_for_serdes
+@record
+class CensusSync:
+    """Represents a Census sync, based on data as returned from the API."""
+
+    id: int
+    name: str
+    source_id: int
+    destination_id: int
+    mappings: list[dict[str, Any]]
+
+
+@whitelist_for_serdes
+@record
+class CensusWorkspaceData:
+    """A record representing all content in a Census workspace.
+    Provided as context for the translator so that it can resolve dependencies between content.
+    """
+
+    syncs: list[CensusSync]
+
+    @property
+    def syncs_by_id(self) -> Mapping[int, CensusSync]:
+        """Returns a mapping of sync IDs to CensusSync objects."""
+        return {sync.id: sync for sync in self.syncs}
+
+
+def generate_table_schema(sync_mappings_props: list[dict[str, Any]]) -> dg.TableSchema:
+    return dg.TableSchema(
+        columns=sorted(
+            [
+                dg.TableColumn(name=mapping["to"], type=mapping.get("field_type", "unknown"))
+                for mapping in sync_mappings_props
+            ],
+            key=lambda col: col.name,
+        )
+    )
+
+
+class CensusMetadataSet(NamespacedMetadataSet):
+    sync_id: int
+    sync_name: str
+    source_id: int
+    destination_id: int
+
+    @classmethod
+    def namespace(cls) -> str:
+        return "dagster-census"

dagster_census/version.py

@@ -1 +1 @@
-__version__ = "0.28.0"
+__version__ = "0.28.2"

{dagster_census-0.28.0.dist-info → dagster_census-0.28.2.dist-info}/METADATA

@@ -1,21 +1,20 @@
 Metadata-Version: 2.4
 Name: dagster-census
-Version: 0.28.0
+Version: 0.28.2
 Summary: Package for integrating Census with Dagster.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-census
 Author: Dagster Labs
 Author-email: hello@dagsterlabs.com
 License: Apache-2.0
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.9,<3.14
+Requires-Python: >=3.10,<3.14
 License-File: LICENSE
-Requires-Dist: dagster==1.12.0
+Requires-Dist: dagster==1.12.2
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

dagster_census-0.28.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+dagster_census/__init__.py,sha256=eIDWjz13trVGLVMDjAdq0LeUkbEiWex1ho2NX8cDFOY,501
+dagster_census/ops.py,sha256=4N53-w0u3PG5kn-q1tW-7Ac-i5SCDFq9ku--g7gDig4,3659
+dagster_census/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dagster_census/resources.py,sha256=R-McSQLQTOImDVAQ3Tr-ZbvbUMQnbKpW7hkvtuwtS7I,11661
+dagster_census/translator.py,sha256=caSshQ3ZfSsl_ASNlN6p24t-15WDRN29Z1xBg7Utbp4,1510
+dagster_census/types.py,sha256=on1CIKWZMsSuw5h-HkzQPVctMQzQu2L9HgkX5gc9X54,698
+dagster_census/utils.py,sha256=fFsKjn6xlLEel1-nvAwfAaemjSTg2gYxJHY8sK2e1no,1348
+dagster_census/version.py,sha256=K-TM2fq9AmH_Dk8Cadam72wILDZ_6qftLHvY9P1Fc3I,23
+dagster_census/components/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dagster_census/components/census_component.py,sha256=vgaVfT0pFVoC8ddhJJQEGJwjfZGsjLiTlXDOeNYW4q0,7970
+dagster_census/components/census_scaffolder.py,sha256=n0YT0Rov6yJsPHAVpbtUrkxxF59z0WLCnRhYEpYtF4A,818
+dagster_census-0.28.2.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
+dagster_census-0.28.2.dist-info/METADATA,sha256=Z46H5NkOfdmNX1lcrSUc3oFTz5RtUkV3gIY8q5kjW2o,869
+dagster_census-0.28.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dagster_census-0.28.2.dist-info/top_level.txt,sha256=Lz6T6gBn89ilLGOLr8YIrBTRcWPCjeOtdNk8sdqmUnA,15
+dagster_census-0.28.2.dist-info/RECORD,,

dagster_census-0.28.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-dagster_census/__init__.py,sha256=MsOV0XOueouBCn1FUYSN36fyswIK4pfwaWD2qceNPy4,451
-dagster_census/ops.py,sha256=4N53-w0u3PG5kn-q1tW-7Ac-i5SCDFq9ku--g7gDig4,3659
-dagster_census/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dagster_census/resources.py,sha256=2kktbxyzgZOCo0VB8rUDmvK3-D9IpanYQyPTkTRlsg4,10839
-dagster_census/types.py,sha256=on1CIKWZMsSuw5h-HkzQPVctMQzQu2L9HgkX5gc9X54,698
-dagster_census/utils.py,sha256=fFsKjn6xlLEel1-nvAwfAaemjSTg2gYxJHY8sK2e1no,1348
-dagster_census/version.py,sha256=MRQGtOXBhcDKeeNOL0LiB-cllo6kfd8_KGJOvaDp0XQ,23
-dagster_census-0.28.0.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
-dagster_census-0.28.0.dist-info/METADATA,sha256=nPqqghKY9EODMfw-rCYSf3jim7E5TnT7dunnJM1Ift4,918
-dagster_census-0.28.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dagster_census-0.28.0.dist-info/top_level.txt,sha256=Lz6T6gBn89ilLGOLr8YIrBTRcWPCjeOtdNk8sdqmUnA,15
-dagster_census-0.28.0.dist-info/RECORD,,