pglearned 0.1.0__tar.gz

pglearned-0.1.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: pglearned
+Version: 0.1.0
+Summary: Python bindings for pglearned extension
+Project-URL: Homepage, https://github.com/RMTT/pglearned
+Requires-Python: >=3.7
+Requires-Dist: psycopg[binary]>=3.0.0
+Requires-Dist: grpcio>=1.50.0
+Requires-Dist: grpcio-reflection>=1.50.0
+Requires-Dist: protobuf>=4.0.0

pglearned-0.1.0/README.md

@@ -0,0 +1,77 @@
+# pglearned python
+
+Python bindings and framework for interacting with the `pglearned` PostgreSQL extension. This library allows you to collect query execution data for training and to implement custom query planning logic using Python.
+
+## Installation
+
+Navigate to the `frameworks/python` directory and install the package:
+
+```bash
+pip install .
+```
+
+Ensure you have the `pglearned` extension installed and configured in your PostgreSQL database.
+
+## Usage
+
+### 1. Data Collection
+
+Use `PglClient` to connect to your PostgreSQL database and collect query execution data (dataset) generated by the extension.
+
+```python
+from pgl import PglClient
+
+# Connect to the database
+# ensure the connection string is correct for your setup
+db_url = "postgresql://user:password@localhost:5432/dbname"
+client = PglClient(db_url)
+
+# Iterate over collected query plans from a named dataset
+# 'training_set' should match the dataset name used in your SQL queries
+for query_id, plan_data in client.qdataset_collect("training_set"):
+    print(f"Query ID: {query_id}")
+    print(f"Plan Structure: {plan_data.keys()}")
+    # plan_data is a dictionary containing the JSON plan representation
+```
+
+### 2. Custom Adapter
+
+To implement custom query processing logic (e.g., a machine learning model that selects the best plan), create a class that inherits from `PglAdapter` and serve it using `run_server`.
+
+```python
+from typing import List, Dict, Any
+from pgl import PglAdapter, run_server
+
+class MySmartAdapter(PglAdapter):
+    def choose_plan(self, plans: List[Dict[str, Any]]) -> int:
+        """
+        Choose the best plan from a list of candidates.
+        
+        Args:
+            plans: A list of candidate query plans (parsed as dictionaries).
+            
+        Returns:
+            The index of the chosen plan (0-based).
+        """
+        best_score = -1
+        best_index = 0
+        
+        for i, plan in enumerate(plans):
+            # Implement your scoring logic here
+            # For demonstration, we just print the plan cost if available
+            print(f"Plan {i} cost: {plan.get('Total Cost', 'N/A')}")
+            
+            # Simple heuristic: select plan based on some criteria
+            # ...
+            
+        return best_index
+
+if __name__ == "__main__":
+    # Initialize your adapter
+    adapter = MySmartAdapter()
+    
+    # Start the gRPC server
+    # The pglearned extension will connect to this server to request plan selections
+    print("Starting PglAdapter server on port 50051...")
+    run_server(adapter, host="0.0.0.0", port=50051)
+```

pglearned-0.1.0/pgl/__init__.py

@@ -0,0 +1,5 @@
+from .client import PglClient
+from .adapter import PglAdapter
+from .server import run_server
+
+__all__ = ["PglClient", "PglAdapter", "run_server"]

pglearned-0.1.0/pgl/adapter.py

@@ -0,0 +1,21 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any
+
+
+class PglAdapter(ABC):
+    """
+    Abstract base class for implementing a pglearned adapter.
+    """
+
+    @abstractmethod
+    def choose_plan(self, plans: List[Dict[str, Any]]) -> int:
+        """
+        Choose the best query plan from a list of candidates.
+
+        Args:
+            plans: A list of query plans (parsed as dictionaries).
+
+        Returns:
+            The index of the chosen plan (0-based).
+        """
+        pass

pglearned-0.1.0/pgl/client.py

@@ -0,0 +1,52 @@
+import psycopg
+from typing import Iterator, Tuple, Dict, Any
+
+
+class PglClient:
+    def __init__(self, dburl: str):
+        """
+        Initialize the PglClient with a database connection string.
+
+        Args:
+            dburl: libpq connection string (e.g., "postgresql://user:password@host:port/dbname")
+        """
+        self.dburl = dburl
+
+    def qdataset_collect(
+        self,
+        dataset_name: str,
+        offset: int = -1,
+        limit: int = 10,
+        method: str = "default",
+        arm: int = -1,
+    ) -> Iterator[Tuple[int, Dict[str, Any]]]:
+        """
+        Call the pgl_qdataset_collect UDF to fetch query plans.
+
+        Args:
+            dataset_name: Name of the dataset to run.
+            offset: Offset in the dataset (default -1 uses current_pos).
+            limit: Number of queries to run (default 10).
+            method: Planner method to use (default 'default').
+            arm: Planner arm to use (default -1 for all/iterator).
+
+        Yields:
+            Tuple of (id, plan), where plan is a dictionary (JSON).
+        """
+        # Connect to the database
+        with psycopg.connect(self.dburl) as conn:
+            conn.autocommit = True
+            with conn.cursor() as cur:
+                # Check and install extension if needed
+                cur.execute("SELECT 1 FROM pg_extension WHERE extname = 'pgl'")
+                if not cur.fetchone():
+                    cur.execute("CREATE EXTENSION IF NOT EXISTS pglearned")
+
+                cur.execute(
+                    "SELECT id, plan FROM pgl_qdataset_collect(%s, %s, %s, %s, %s)",
+                    (dataset_name, offset, limit, method, arm),
+                )
+
+                # Iterate over results
+                for row in cur:
+                    yield row[0], row[1]

pglearned-0.1.0/pgl/proto/pgl_rpc_pb2.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: pgl_rpc.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'pgl_rpc.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rpgl_rpc.proto\x12\x07pgl_rpc\"\"\n\x11\x43hoosePlanRequest\x12\r\n\x05plans\x18\x01 \x03(\t\"/\n\x12\x43hoosePlanResponse\x12\x19\n\x11\x63hosen_plan_index\x18\x01 \x01(\x05\x32R\n\tPglRemote\x12\x45\n\nChoosePlan\x12\x1a.pgl_rpc.ChoosePlanRequest\x1a\x1b.pgl_rpc.ChoosePlanResponseb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'pgl_rpc_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_CHOOSEPLANREQUEST']._serialized_start=26
+  _globals['_CHOOSEPLANREQUEST']._serialized_end=60
+  _globals['_CHOOSEPLANRESPONSE']._serialized_start=62
+  _globals['_CHOOSEPLANRESPONSE']._serialized_end=109
+  _globals['_PGLREMOTE']._serialized_start=111
+  _globals['_PGLREMOTE']._serialized_end=193
+# @@protoc_insertion_point(module_scope)

pglearned-0.1.0/pgl/proto/pgl_rpc_pb2_grpc.py

@@ -0,0 +1,97 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+from . import pgl_rpc_pb2 as pgl__rpc__pb2
+
+GRPC_GENERATED_VERSION = '1.78.1'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + ' but the generated code in pgl_rpc_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class PglRemoteStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.ChoosePlan = channel.unary_unary(
+                '/pgl_rpc.PglRemote/ChoosePlan',
+                request_serializer=pgl__rpc__pb2.ChoosePlanRequest.SerializeToString,
+                response_deserializer=pgl__rpc__pb2.ChoosePlanResponse.FromString,
+                _registered_method=True)
+
+
+class PglRemoteServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def ChoosePlan(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_PglRemoteServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'ChoosePlan': grpc.unary_unary_rpc_method_handler(
+                    servicer.ChoosePlan,
+                    request_deserializer=pgl__rpc__pb2.ChoosePlanRequest.FromString,
+                    response_serializer=pgl__rpc__pb2.ChoosePlanResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'pgl_rpc.PglRemote', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('pgl_rpc.PglRemote', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class PglRemote(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def ChoosePlan(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/pgl_rpc.PglRemote/ChoosePlan',
+            pgl__rpc__pb2.ChoosePlanRequest.SerializeToString,
+            pgl__rpc__pb2.ChoosePlanResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

pglearned-0.1.0/pgl/server.py

@@ -0,0 +1,90 @@
+import json
+import logging
+from concurrent import futures
+
+import grpc
+from grpc_reflection.v1alpha import reflection
+
+# Note: These imports assume the protobuf code has been generated.
+# Run: python -m grpc_tools.protoc -I../../proto --python_out=./pgl/proto --grpc_python_out=./pgl/proto ../../proto/pgl_rpc.proto
+from .proto import pgl_rpc_pb2
+from .proto import pgl_rpc_pb2_grpc
+
+from .adapter import PglAdapter
+
+logger = logging.getLogger(__name__)
+
+
+class PglRemoteAdapter(pgl_rpc_pb2_grpc.PglRemoteServicer):
+    """
+    Internal gRPC servicer that adapts the PglAdapter to the PglRemote service.
+    """
+
+    def __init__(self, adapter: PglAdapter):
+        self.adapter = adapter
+
+    def ChoosePlan(self, request, context):
+        if not request.plans:
+            return pgl_rpc_pb2.ChoosePlanResponse(chosen_plan_index=0)
+
+        # Parse JSON strings to Dicts
+        parsed_plans = []
+        for i, plan_str in enumerate(request.plans):
+            try:
+                parsed_plans.append(json.loads(plan_str))
+            except json.JSONDecodeError as e:
+                logger.error(f"Failed to parse plan at index {i}: {e}")
+                parsed_plans.append({})
+
+        try:
+            chosen_index = self.adapter.choose_plan(parsed_plans)
+        except Exception as e:
+            logger.exception("Error in user adapter logic")
+            context.abort(grpc.StatusCode.INTERNAL, f"Adapter error: {str(e)}")
+            return
+
+        # Validate return value
+        if not (0 <= chosen_index < len(request.plans)):
+            logger.error(
+                f"Invalid plan index returned: {chosen_index}. Max: {len(request.plans) - 1}"
+            )
+            chosen_index = 0
+
+        return pgl_rpc_pb2.ChoosePlanResponse(chosen_plan_index=chosen_index)
+
+
+def run_server(
+    adapter: PglAdapter, host: str = "0.0.0.0", port: int = 50051, max_workers: int = 10
+):
+    """
+    Start the gRPC server hosting the given adapter.
+
+    Args:
+        adapter: An instance of a PglAdapter subclass.
+        host: Host to bind to.
+        port: Port to bind to.
+        max_workers: Number of worker threads for the gRPC server.
+    """
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=max_workers))
+    servicer = PglRemoteAdapter(adapter)
+
+    # This registration function is generated by grpc_tools
+    pgl_rpc_pb2_grpc.add_PglRemoteServicer_to_server(servicer, server)
+    service_full_name = pgl_rpc_pb2.DESCRIPTOR.services_by_name["PglRemote"].full_name
+
+    # Enable reflection
+    SERVICE_NAMES = (
+        service_full_name,
+        reflection.SERVICE_NAME,
+    )
+    reflection.enable_server_reflection(SERVICE_NAMES, server)
+
+    address = f"{host}:{port}"
+    server.add_insecure_port(address)
+
+    print(f"PglServer starting on {address}...")
+    server.start()
+    try:
+        server.wait_for_termination()
+    except KeyboardInterrupt:
+        server.stop(0)

pglearned-0.1.0/pglearned.egg-info/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: pglearned
+Version: 0.1.0
+Summary: Python bindings for pglearned extension
+Project-URL: Homepage, https://github.com/RMTT/pglearned
+Requires-Python: >=3.7
+Requires-Dist: psycopg[binary]>=3.0.0
+Requires-Dist: grpcio>=1.50.0
+Requires-Dist: grpcio-reflection>=1.50.0
+Requires-Dist: protobuf>=4.0.0

pglearned-0.1.0/pglearned.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+pyproject.toml
+pgl/__init__.py
+pgl/adapter.py
+pgl/client.py
+pgl/server.py
+pgl/proto/__init__.py
+pgl/proto/pgl_rpc_pb2.py
+pgl/proto/pgl_rpc_pb2_grpc.py
+pglearned.egg-info/PKG-INFO
+pglearned.egg-info/SOURCES.txt
+pglearned.egg-info/dependency_links.txt
+pglearned.egg-info/requires.txt
+pglearned.egg-info/top_level.txt

pglearned-0.1.0/pglearned.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

pglearned-0.1.0/pglearned.egg-info/requires.txt

@@ -0,0 +1,4 @@
+psycopg[binary]>=3.0.0
+grpcio>=1.50.0
+grpcio-reflection>=1.50.0
+protobuf>=4.0.0

pglearned-0.1.0/pglearned.egg-info/top_level.txt

@@ -0,0 +1 @@
+pgl

pglearned-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "pglearned"
+version = "0.1.0"
+description = "Python bindings for pglearned extension"
+requires-python = ">=3.7"
+dependencies = [
+    "psycopg[binary]>=3.0.0",
+    "grpcio>=1.50.0",
+    "grpcio-reflection>=1.50.0",
+    "protobuf>=4.0.0",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/RMTT/pglearned"

pglearned-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+