catopt-query-category-theor.../catopt_query/adapters.py

83 lines
2.6 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, Any
from .protocol import LocalProblem, SharedVariables, CanonicalPlan
@dataclass
class VendorPlan:
# Minimal stand-in for a vendor-specific plan representation
shard_id: str
projection: list
predicates: list
price: float
class Adapter:
"""Abstract adapter: maps a vendor-specific plan into a canonical plan."""
def to_canonical(self, vendor_plan: VendorPlan) -> CanonicalPlan:
# Simple, deterministic mapping; can be overridden by concrete adapters
return CanonicalPlan(
projection=vendor_plan.projection,
predicates=vendor_plan.predicates,
estimated_cost=float(vendor_plan.price),
)
@dataclass
class PostgresVendorPlan:
shard_id: str
table: str
projection: list
predicates: list
price: float
class PostgresAdapter(Adapter):
"""Concrete adapter mapping a PostgresVendorPlan into a CanonicalPlan."""
def to_canonical(self, vendor_plan: VendorPlan) -> CanonicalPlan:
# Accept either a PostgresVendorPlan or a generic VendorPlan-like object
if isinstance(vendor_plan, PostgresVendorPlan):
vp = vendor_plan
return CanonicalPlan(
projection=vp.projection,
predicates=vp.predicates,
estimated_cost=float(vp.price),
)
proj = getattr(vendor_plan, "projection", [])
preds = getattr(vendor_plan, "predicates", [])
price = float(getattr(vendor_plan, "price", 0.0))
return CanonicalPlan(projection=proj, predicates=preds, estimated_cost=price)
@dataclass
class MongoVendorPlan:
shard_id: str
collection: str
projection: list
predicates: list
price: float
class MongoAdapter(Adapter):
"""Concrete adapter mapping a MongoVendorPlan into a CanonicalPlan."""
def to_canonical(self, vendor_plan: VendorPlan) -> CanonicalPlan:
# Accept either a MongoVendorPlan or a generic VendorPlan-like object
if isinstance(vendor_plan, MongoVendorPlan):
vp = vendor_plan
return CanonicalPlan(
projection=vp.projection,
predicates=vp.predicates,
estimated_cost=float(vp.price),
)
# Fallback: try to treat as a generic VendorPlan
proj = getattr(vendor_plan, "projection", [])
preds = getattr(vendor_plan, "predicates", [])
price = float(getattr(vendor_plan, "price", 0.0))
return CanonicalPlan(projection=proj, predicates=preds, estimated_cost=price)