Skip to content

Infrastructure Layer

The infrastructure layer implements the repository interfaces defined in the domain layer. It is the only layer that touches DuckDB, pickle files, or external HTTP calls.

DuckDB Adapter

src.infrastructure.db.duckdb_adapter

DuckDB adapter – infrastructure layer database connection.

All SQL lives here. Domain layer remains pure Python.

Functions

get_connection 
get_connection(read_only: bool = True) -> Generator[duckdb.DuckDBPyConnection, None, None]

Context manager for a DuckDB connection.

Parameters:

Name Type Description Default
read_only bool

If True (default), opens a read-only connection safe for concurrent API workers. Set False for write operations.

 True

Yields:

Type Description
DuckDBPyConnection

An active DuckDB connection that is closed on exit.
Source code in src/infrastructure/db/duckdb_adapter.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
@contextmanager
def get_connection(read_only: bool = True) -> Generator[duckdb.DuckDBPyConnection, None, None]:
    """Context manager for a DuckDB connection.

    Args:
        read_only: If True (default), opens a read-only connection safe for
                   concurrent API workers. Set False for write operations.

    Yields:
        An active DuckDB connection that is closed on exit.
    """
    conn = duckdb.connect(database=_DB_PATH, read_only=read_only)
    logger.debug("duckdb.connection.opened", path=_DB_PATH, read_only=read_only)
    try:
        yield conn
    finally:
        conn.close()
        logger.debug("duckdb.connection.closed")

Customer Repository (DuckDB)

src.infrastructure.repositories.customer_repository

DuckDB implementation of CustomerRepository.

Classes

DuckDBCustomerRepository

Bases: CustomerRepository

Reads Customer entities from the DuckDB warehouse.

Source code in src/infrastructure/repositories/customer_repository.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class DuckDBCustomerRepository(CustomerRepository):
    """Reads Customer entities from the DuckDB warehouse."""

    def get_by_id(self, customer_id: str) -> Customer | None:
        """Fetch a single customer by ID."""
        with get_connection() as conn:
            row = conn.execute(
                """
                SELECT customer_id, industry, plan_tier, signup_date, mrr, churn_date
                FROM raw.customers
                WHERE customer_id = ?
                """,
                [customer_id],
            ).fetchone()

        if row is None:
            return None
        return self._row_to_entity(row)

    def get_all_active(self) -> Sequence[Customer]:
        """Return all customers without a churn_date."""
        with get_connection() as conn:
            rows = conn.execute(
                """
                SELECT customer_id, industry, plan_tier, signup_date, mrr, churn_date
                FROM raw.customers
                WHERE churn_date IS NULL
                ORDER BY mrr DESC
                """
            ).fetchall()
        return [self._row_to_entity(row) for row in rows]

    def get_sample(self, n: int) -> Sequence[Customer]:
        """Return a random sample of n customers using DuckDB reservoir sampling.

        DuckDB's SAMPLE clause does not support parameterized queries — only
        literal integer constants are accepted. n is safe to interpolate because
        callers must clamp it to ≤ 100 before passing it in.
        """
        with get_connection() as conn:
            rows = conn.execute(
                f"""
                SELECT customer_id, industry, plan_tier, signup_date, mrr, churn_date
                FROM raw.customers
                USING SAMPLE reservoir({n} ROWS) REPEATABLE(42)
                """
            ).fetchall()
        return [self._row_to_entity(row) for row in rows]

    def save(self, customer: Customer) -> None:
        """Upsert a customer record."""
        with get_connection(read_only=False) as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO raw.customers
                    (customer_id, industry, plan_tier, signup_date, mrr, churn_date)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                [
                    customer.customer_id,
                    customer.industry.value,
                    customer.plan_tier.value,
                    customer.signup_date,
                    float(customer.mrr.amount),
                    customer.churn_date,
                ],
            )

    @staticmethod
    def _row_to_entity(row: tuple) -> Customer:  # type: ignore[type-arg]
        customer_id, industry, plan_tier, signup_date, mrr, churn_date = row
        return Customer(
            customer_id=str(customer_id),
            industry=Industry(industry),
            plan_tier=PlanTier(plan_tier),
            signup_date=date.fromisoformat(str(signup_date)),
            mrr=MRR(amount=Decimal(str(mrr))),
            churn_date=date.fromisoformat(str(churn_date)) if churn_date else None,
        )
Functions
get_by_id 
get_by_id(customer_id: str) -> Customer | None

Fetch a single customer by ID.

Source code in src/infrastructure/repositories/customer_repository.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def get_by_id(self, customer_id: str) -> Customer | None:
    """Fetch a single customer by ID."""
    with get_connection() as conn:
        row = conn.execute(
            """
            SELECT customer_id, industry, plan_tier, signup_date, mrr, churn_date
            FROM raw.customers
            WHERE customer_id = ?
            """,
            [customer_id],
        ).fetchone()

    if row is None:
        return None
    return self._row_to_entity(row)
get_all_active 
get_all_active() -> Sequence[Customer]

Return all customers without a churn_date.

Source code in src/infrastructure/repositories/customer_repository.py 
34
35
36
37
38
39
40
41
42
43
44
45
def get_all_active(self) -> Sequence[Customer]:
    """Return all customers without a churn_date."""
    with get_connection() as conn:
        rows = conn.execute(
            """
            SELECT customer_id, industry, plan_tier, signup_date, mrr, churn_date
            FROM raw.customers
            WHERE churn_date IS NULL
            ORDER BY mrr DESC
            """
        ).fetchall()
    return [self._row_to_entity(row) for row in rows]
get_sample 
get_sample(n: int) -> Sequence[Customer]

Return a random sample of n customers using DuckDB reservoir sampling.

DuckDB's SAMPLE clause does not support parameterized queries — only literal integer constants are accepted. n is safe to interpolate because callers must clamp it to ≤ 100 before passing it in.

Source code in src/infrastructure/repositories/customer_repository.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def get_sample(self, n: int) -> Sequence[Customer]:
    """Return a random sample of n customers using DuckDB reservoir sampling.

    DuckDB's SAMPLE clause does not support parameterized queries — only
    literal integer constants are accepted. n is safe to interpolate because
    callers must clamp it to ≤ 100 before passing it in.
    """
    with get_connection() as conn:
        rows = conn.execute(
            f"""
            SELECT customer_id, industry, plan_tier, signup_date, mrr, churn_date
            FROM raw.customers
            USING SAMPLE reservoir({n} ROWS) REPEATABLE(42)
            """
        ).fetchall()
    return [self._row_to_entity(row) for row in rows]
save 
save(customer: Customer) -> None

Upsert a customer record.

Source code in src/infrastructure/repositories/customer_repository.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def save(self, customer: Customer) -> None:
    """Upsert a customer record."""
    with get_connection(read_only=False) as conn:
        conn.execute(
            """
            INSERT OR REPLACE INTO raw.customers
                (customer_id, industry, plan_tier, signup_date, mrr, churn_date)
            VALUES (?, ?, ?, ?, ?, ?)
            """,
            [
                customer.customer_id,
                customer.industry.value,
                customer.plan_tier.value,
                customer.signup_date,
                float(customer.mrr.amount),
                customer.churn_date,
            ],
        )

Functions

Usage Repository (DuckDB)

src.infrastructure.repositories.usage_repository

DuckDB implementation of UsageRepository.

Classes

DuckDBUsageRepository

Bases: UsageRepository

Reads UsageEvent entities from the DuckDB warehouse.

Source code in src/infrastructure/repositories/usage_repository.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class DuckDBUsageRepository(UsageRepository):
    """Reads UsageEvent entities from the DuckDB warehouse."""

    def get_events_for_customer(
        self,
        customer_id: str,
        since: datetime | None = None,
    ) -> Sequence[UsageEvent]:
        since_clause = "AND timestamp >= ?" if since else ""
        params = [customer_id] + ([since] if since else [])

        with get_connection() as conn:
            rows = conn.execute(
                f"""
                SELECT event_id, customer_id, timestamp, event_type, feature_adoption_score
                FROM raw.usage_events
                WHERE customer_id = ? {since_clause}
                ORDER BY timestamp DESC
                """,
                params,
            ).fetchall()

        return [self._row_to_entity(row) for row in rows]

    def get_event_count_last_n_days(self, customer_id: str, days: int) -> int:
        with get_connection() as conn:
            result = conn.execute(
                """
                SELECT COUNT(*) FROM raw.usage_events
                WHERE customer_id = ?
                  AND timestamp >= CURRENT_TIMESTAMP - INTERVAL (?) DAY
                """,
                [customer_id, days],
            ).fetchone()
        return int(result[0]) if result else 0

    @staticmethod
    def _row_to_entity(row: tuple) -> UsageEvent:  # type: ignore[type-arg]
        event_id, customer_id, timestamp, event_type, adoption_score = row
        return UsageEvent(
            event_id=str(event_id),
            customer_id=str(customer_id),
            timestamp=timestamp if isinstance(timestamp, datetime) else datetime.fromisoformat(str(timestamp)),
            event_type=EventType(event_type),
            feature_adoption_score=FeatureAdoptionScore(value=float(adoption_score)),
        )

Functions

Model Registry

src.infrastructure.ml.model_registry

Model registry — loads model artifacts from the ZenML Model Control Plane.

Functions

load_model cached 
load_model(name: str) -> Any

Load the production model artifact from the ZenML Model Control Plane.

Parameters:

Name Type Description Default
name str

Model name, e.g. "churn_model". Must match the Model(name=...) in the training pipeline.

 required

Returns:

Type Description
Any

The deserialized model object.

Raises:

Type Description
Exception

If the ZenML server is unreachable or no production artifact exists.
Source code in src/infrastructure/ml/model_registry.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
@lru_cache(maxsize=4)
def load_model(name: str) -> Any:  # noqa: ANN401 — generic model loader, return type depends on artifact
    """Load the production model artifact from the ZenML Model Control Plane.

    Args:
        name: Model name, e.g. "churn_model". Must match the Model(name=...) in
              the training pipeline.

    Returns:
        The deserialized model object.

    Raises:
        Exception: If the ZenML server is unreachable or no production artifact exists.
    """
    from zenml.client import Client
    from zenml.enums import ModelStages

    client = Client()
    model_version = client.get_model_version(
        model_name_or_id=name,
        model_version_name_or_number_or_id=ModelStages.PRODUCTION,
    )
    artifact = model_version.get_artifact("calibrated_model")
    model = artifact.load()
    logger.info(
        "model.loaded_from_zenml_mcp",
        name=name,
        version=str(model_version.version),
    )
    return model
get_model_metadata 
get_model_metadata(name: str) -> dict[str, Any]

Return metadata (version, training date, metrics) for a model artifact.

Source code in src/infrastructure/ml/model_registry.py 
50
51
52
53
54
55
56
def get_model_metadata(name: str) -> dict[str, Any]:
    """Return metadata (version, training date, metrics) for a model artifact."""
    meta_path = _MODELS_DIR / f"{name}_metadata.json"
    if not meta_path.exists():
        return {"version": "unknown"}
    with open(meta_path) as f:
        return json.load(f)  # type: ignore[no-any-return]