buildgrid 0.2.52__tar.gz → 0.3.0__tar.gz

{buildgrid-0.2.52 → buildgrid-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildgrid
-Version: 0.2.52
+Version: 0.3.0
 Summary: A remote execution service
 License: Apache License, Version 2.0
 Project-URL: Homepage, https://buildgrid.build
@@ -42,7 +42,7 @@ Requires-Dist: fakeredis>=2.10.1; extra == "redis"
 Requires-Dist: redis>=4.5.1; extra == "redis"
 Requires-Dist: hiredis; extra == "redis"
 Provides-Extra: docs
-Requires-Dist: Sphinx; extra == "docs"
+Requires-Dist: Sphinx<=8; extra == "docs"
 Requires-Dist: sphinx-click; extra == "docs"
 Requires-Dist: sphinx-rtd-theme; extra == "docs"
 Requires-Dist: sphinxcontrib-apidoc; extra == "docs"
@@ -77,6 +77,7 @@ Requires-Dist: pytest-xdist; extra == "dev"
 Requires-Dist: memray; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: grpcio-tools; extra == "dev"
+Requires-Dist: bump4version; extra == "dev"
 Provides-Extra: mypy
 Requires-Dist: mypy; extra == "mypy"
 Requires-Dist: SQLAlchemy[mypy]; extra == "mypy"

{buildgrid-0.2.52 → buildgrid-0.3.0}/buildgrid/server/cas/storage/index/sql.py

@@ -17,8 +17,7 @@
 SQLIndex
 ==================
 
-A SQL index implementation. This can be pointed to either a remote SQL server
-or a local SQLite database.
+A SQL index implementation. This must be pointed to a remote SQL server.
 
 """
 
@@ -35,6 +34,8 @@ from sqlalchemy.orm import InstrumentedAttribute, Session, load_only
 from sqlalchemy.orm.exc import StaleDataError
 from sqlalchemy.orm.query import Query
 from sqlalchemy.orm.session import Session as SessionType
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.sql.functions import coalesce
 
 from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest
 from buildgrid._protos.google.rpc import code_pb2
@@ -49,10 +50,8 @@ from buildgrid.server.sql.provider import SqlProvider
 
 from ..storage_abc import StorageABC
 from .index_abc import IndexABC
-from .sql_dialect_delegates import PostgreSQLDelegate, SQLiteDelegate
 
 LOGGER = buildgrid_logger(__name__)
-DIALECT_DELEGATES = {"postgresql": PostgreSQLDelegate, "sqlite": SQLiteDelegate}
 
 INLINE_BLOB_SIZE_HARD_MAXIMUM = 1000000000
 
@@ -151,9 +150,6 @@ class SQLIndex(IndexABC):
             unknown_args = kwargs_keys - available_options
             raise TypeError(f"Unknown keyword arguments: [{unknown_args}]")
 
-        # Dialect-specific initialization
-        self._dialect_delegate = DIALECT_DELEGATES.get(self._sql.dialect)
-
         if inclause_limit > 0:
             if inclause_limit > window_size:
                 LOGGER.warning(
@@ -377,14 +373,32 @@ class SQLIndex(IndexABC):
 
         digest_blob_pairs = sorted(digest_blob_pairs, key=lambda pair: (pair[0].hash, pair[0].size_bytes))
 
-        if self._dialect_delegate:
-            try:
-                self._dialect_delegate._save_digests_to_index(  # type: ignore
-                    digest_blob_pairs, session, self._max_inline_blob_size
-                )
-                return
-            except AttributeError:
-                pass
+        # See discussion of __table__ typing in https://github.com/sqlalchemy/sqlalchemy/issues/9130
+        index_table = cast(Table, IndexEntry.__table__)
+        update_time = datetime.utcnow()
+        new_rows = [
+            {
+                "digest_hash": digest.hash,
+                "digest_size_bytes": digest.size_bytes,
+                "accessed_timestamp": update_time,
+                "inline_blob": (blob if digest.size_bytes <= self._max_inline_blob_size else None),
+                "deleted": False,
+            }
+            for (digest, blob) in digest_blob_pairs
+        ]
+
+        base_insert_stmt = insert(index_table).values(new_rows)
+
+        update_stmt = base_insert_stmt.on_conflict_do_update(
+            index_elements=["digest_hash"],
+            set_={
+                "accessed_timestamp": update_time,
+                "inline_blob": coalesce(base_insert_stmt.excluded.inline_blob, index_table.c.inline_blob),
+                "deleted": False,
+            },
+        )
+
+        session.execute(update_stmt)
 
         update_time = datetime.utcnow()
         # Figure out which digests we can just update

{buildgrid-0.2.52 → buildgrid-0.3.0}/buildgrid/server/cas/storage/sql.py

@@ -26,8 +26,7 @@ from io import BytesIO
 from typing import IO, Any, Iterator, Sequence, TypedDict, cast
 
 from sqlalchemy import CursorResult, delete, func, select
-from sqlalchemy.dialects.postgresql import insert as postgresql_insert
-from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm.exc import StaleDataError
 
@@ -59,7 +58,7 @@ class SQLStorage(StorageABC):
         self._sql_ro = sql_ro_provider or sql_provider
         self._inclause_limit = self._sql.default_inlimit
 
-        supported_dialects = ["postgresql", "sqlite"]
+        supported_dialects = ["postgresql"]
 
         if self._sql.dialect not in supported_dialects:
             raise RuntimeError(
@@ -71,14 +70,6 @@ class SQLStorage(StorageABC):
         with self._sql.session() as session:
             session.query(BlobEntry).first()
 
-    def _sqlite_bulk_insert(self, new_rows: list[DigestRow]) -> None:
-        with self._sql.session() as session:
-            session.execute(sqlite_insert(BlobEntry).values(new_rows).on_conflict_do_nothing())
-
-    def _postgresql_bulk_insert(self, new_rows: list[DigestRow]) -> None:
-        with self._sql.session() as session:
-            session.execute(postgresql_insert(BlobEntry).values(new_rows).on_conflict_do_nothing())
-
     def _bulk_insert(self, digests: list[tuple[Digest, bytes]]) -> None:
         # Sort digests by hash to ensure consistent order to minimize deadlocks
         # when BatchUpdateBlobs requests have overlapping blobs
@@ -87,12 +78,8 @@ class SQLStorage(StorageABC):
             for (digest, blob) in sorted(digests, key=lambda x: x[0].hash)
         ]
 
-        if self._sql.dialect == "sqlite":
-            self._sqlite_bulk_insert(new_rows)
-        elif self._sql.dialect == "postgresql":
-            self._postgresql_bulk_insert(new_rows)
-        else:
-            raise RuntimeError(f"Unsupported dialect {self._sql.dialect} for bulk_insert")
+        with self._sql.session() as session:
+            session.execute(insert(BlobEntry).values(new_rows).on_conflict_do_nothing())
 
     @timed(METRIC.STORAGE.STAT_DURATION, type=TYPE)
     def has_blob(self, digest: Digest) -> bool:

{buildgrid-0.2.52 → buildgrid-0.3.0}/buildgrid/server/scheduler/impl.py

@@ -31,7 +31,7 @@ from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
 from google.protobuf.timestamp_pb2 import Timestamp
 from grpc import Channel
 from sqlalchemy import ColumnExpressionArgument, CursorResult, and_, delete, func, insert, or_, select, text, update
-from sqlalchemy.dialects import postgresql, sqlite
+from sqlalchemy.dialects import postgresql
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session, joinedload
 from sqlalchemy.sql.expression import Insert, Select
@@ -760,11 +760,10 @@ class Scheduler:
             )
 
     def _notify_job_updated(self, job_names: str | list[str], session: Session) -> None:
-        if self._sql.dialect == "postgresql":
-            if isinstance(job_names, str):
-                job_names = [job_names]
-            for job_name in job_names:
-                session.execute(text(f"NOTIFY {NotificationChannel.JOB_UPDATED.value}, '{job_name}';"))
+        if isinstance(job_names, str):
+            job_names = [job_names]
+        for job_name in job_names:
+            session.execute(text(f"NOTIFY {NotificationChannel.JOB_UPDATED.value}, '{job_name}';"))
 
     def _get_operation(self, operation_name: str, session: Session) -> OperationEntry | None:
         statement = (
@@ -776,7 +775,7 @@ class Scheduler:
 
     def _batch_timeout_jobs(self, job_select_stmt: Select[Any], status_code: int, message: str) -> int:
         """Timeout all jobs selected by a query"""
-        with self._sql.session(sqlite_lock_immediately=True, exceptions_to_not_raise_on=[Exception]) as session:
+        with self._sql.session(exceptions_to_not_raise_on=[Exception]) as session:
             # Get the full list of jobs to timeout
             job_entries = session.execute(job_select_stmt).scalars().all()
             jobs = []
@@ -838,7 +837,7 @@ class Scheduler:
             .limit(1)
             .with_for_update(skip_locked=True)
         )
-        with self._sql.session(sqlite_lock_immediately=True, exceptions_to_not_raise_on=[Exception]) as session:
+        with self._sql.session(exceptions_to_not_raise_on=[Exception]) as session:
             job = session.execute(stale_job_statement).scalar_one_or_none()
             if not job:
                 return False
@@ -916,22 +915,14 @@ class Scheduler:
         job.cancelled = True
 
         # If the job was assigned to a bot, we need to update the quota / capacity
-        if self._sql.dialect == "postgresql":
-            update_query = (
-                update(BotEntry)
-                .where(BotEntry.bot_id == job.worker_name)
-                .values(capacity=BotEntry.capacity + 1)
-                .returning(BotEntry.cohort)
-            )
-            if cohort := session.execute(update_query).scalar_one_or_none():
-                self._update_instance_quota_usage(session, cohort, job.instance_name, -1, guard=None)
-        else:
-            if bot := session.execute(
-                select(BotEntry).where(BotEntry.bot_id == job.worker_name).with_for_update()
-            ).scalar_one_or_none():
-                bot.capacity += 1
-                if bot.cohort:
-                    self._update_instance_quota_usage(session, bot.cohort, job.instance_name, -1, guard=None)
+        update_query = (
+            update(BotEntry)
+            .where(BotEntry.bot_id == job.worker_name)
+            .values(capacity=BotEntry.capacity + 1)
+            .returning(BotEntry.cohort)
+        )
+        if cohort := session.execute(update_query).scalar_one_or_none():
+            self._update_instance_quota_usage(session, cohort, job.instance_name, -1, guard=None)
 
         session.add(
             JobHistoryEntry(
@@ -1157,8 +1148,7 @@ class Scheduler:
         self._notify_job_updated(job.name, session)
 
         LOGGER.debug("Assigned job to bot", tags=log_tags)
-        if self._sql.dialect == "postgresql":
-            session.execute(text(f"NOTIFY {NotificationChannel.JOB_ASSIGNED.value}, '{bot.name}';"))
+        session.execute(text(f"NOTIFY {NotificationChannel.JOB_ASSIGNED.value}, '{bot.name}';"))
 
     def _match_bot_by_sampling(
         self, session: Session, query: Select[tuple[BotEntry]], sampling: SamplingConfig
@@ -1523,7 +1513,7 @@ class Scheduler:
             )
 
         updated = False
-        with self._sql.session(sqlite_lock_immediately=True) as session:
+        with self._sql.session() as session:
             job = session.execute(job_statement).scalar_one_or_none()
             if job is not None:
                 self._match_job_to_bot(session, job, failure_backoff, bot_assignment_fn, assigner_name)
@@ -1558,7 +1548,7 @@ class Scheduler:
             )
 
         updated = False
-        with self._sql.session(sqlite_lock_immediately=True) as session:
+        with self._sql.session() as session:
             job = session.execute(job_statement).scalar_one_or_none()
             if job is not None:
                 self._match_job_to_bot(session, job, failure_backoff, bot_assignment_fn, assigner_name)
@@ -1702,22 +1692,8 @@ class Scheduler:
         return num_rows_deleted
 
     def _insert_on_conflict_do_nothing(self, model: type[OrmBase]) -> Insert:
-        # `Insert.on_conflict_do_nothing` is a SQLAlchemy "generative method", it
-        # returns a modified copy of the statement it is called on. For
-        # some reason mypy can't understand this, so the errors are ignored here.
-        if self._sql.dialect == "sqlite":
-            sqlite_insert: sqlite.Insert = sqlite.insert(model)
-            return sqlite_insert.on_conflict_do_nothing()
-
-        elif self._sql.dialect == "postgresql":
-            insertion: postgresql.Insert = postgresql.insert(model)
-            return insertion.on_conflict_do_nothing()
-
-        else:
-            # Fall back to the non-specific insert implementation. This doesn't
-            # support `ON CONFLICT DO NOTHING`, so callers need to be careful to
-            # still catch IntegrityErrors if other database backends are possible.
-            return insert(model)
+        insertion: postgresql.Insert = postgresql.insert(model)
+        return insertion.on_conflict_do_nothing()
 
     def get_or_create_client_identity_in_store(
         self, session: Session, client_id: ClientIdentityEntry
@@ -1744,7 +1720,7 @@ class Scheduler:
         try:
             session.execute(insertion)
 
-        # Handle unique constraint violation when using an unsupported database (ie. not PostgreSQL or SQLite)
+        # Handle unique constraint violation when using an unsupported database (ie. not PostgreSQL)
         except IntegrityError:
             LOGGER.debug("Handled IntegrityError when inserting client identity.")
 
@@ -1777,7 +1753,7 @@ class Scheduler:
         try:
             session.execute(insertion)
 
-        # Handle unique constraint violation when using an unsupported database (ie. not PostgreSQL or SQLite)
+        # Handle unique constraint violation when using an unsupported database (ie. not PostgreSQL)
         except IntegrityError:
             LOGGER.debug("Handled IntegrityError when inserting request metadata.")
 
@@ -1949,7 +1925,9 @@ class Scheduler:
             )
 
     @timed(METRIC.SCHEDULER.ASSIGNMENT_DURATION)
-    def _fetch_job_for_bot(self, session: Session, bot: BotEntry, log_tags: Tags) -> JobEntry | None:
+    def _fetch_job_for_bot(
+        self, session: Session, bot: BotEntry, usage_diffs: InstanceQuotaUsageDiffs, log_tags: Tags
+    ) -> JobEntry | None:
         # Attempt to fetch a new job for a bot to work on.
         # This can help if there are usually more jobs available than bots.
 
@@ -1959,15 +1937,22 @@ class Scheduler:
             if bot.instance_name != "*":
                 job_statement = job_statement.where(self._job_in_instance_pool())
 
-            # Prioritize instance where usage <= max_quota - bot.capacity
-            # `- bot.capacity` to avoid over-assigning jobs to bots when nearing quota limits
-            instances_query = select(InstanceQuota.instance_name).where(
-                InstanceQuota.bot_cohort == bot.cohort,
-                InstanceQuota.current_usage <= InstanceQuota.max_quota - bot.capacity,
-            )
-            instances = session.execute(instances_query).scalars().all()
-            if instances:
-                job_statement = job_statement.where(JobEntry.instance_name.in_(instances))
+            if bot.cohort:
+                # Prioritize instance where usage <= max_quota - bot.capacity
+                # `- bot.capacity` to avoid over-assigning jobs to bots when nearing quota limits
+                instances_query = select(InstanceQuota.instance_name).where(
+                    InstanceQuota.bot_cohort == bot.cohort,
+                    InstanceQuota.current_usage <= InstanceQuota.max_quota - bot.capacity,
+                )
+                instances: set[str] = set()
+                instances.update(session.execute(instances_query).scalars().all())
+                # Always allow scheduling more jobs of an instance if we're returning usage
+                instances.update(
+                    {instance for (cohort, instance), diff in usage_diffs.items() if cohort == bot.cohort and diff < 0}
+                )
+
+                if instances:
+                    job_statement = job_statement.where(JobEntry.instance_name.in_(instances))
 
             if next_job := session.execute(job_statement).scalar_one_or_none():
                 log_tags["db.next_job_name"] = next_job.name
@@ -2321,7 +2306,7 @@ class Scheduler:
 
         for _ in range(fetch_limit):
             # Try to fill up the newly free capacity with new jobs.
-            if new_job := self._fetch_job_for_bot(session, bot, log_tags):
+            if new_job := self._fetch_job_for_bot(session, bot, usage_diffs, log_tags):
                 if bot.cohort:
                     usage_diffs[(bot.cohort, new_job.instance_name)] += 1
                 synchronized_leases.append(new_job.to_lease_proto())
@@ -2844,18 +2829,9 @@ class Scheduler:
         if self.bot_locality_hint_limit == 0:
             return
 
-        # Insert new hint with dialect-specific seq handling
-        if self._sql.dialect == "postgresql":
-            # For PostgreSQL, use the sequence to get the next seq number
-            next_seq = None
-        else:
-            # For SQLite, manually find the max seq number
-            max_seq = session.execute(
-                select(func.coalesce(func.max(BotLocalityHintEntry.sequence_number), 0)).where(
-                    BotLocalityHintEntry.bot_name == bot_name
-                )
-            ).scalar_one()
-            next_seq = max_seq + 1
+        # Insert new hint with seq handling
+        # For PostgreSQL, use the sequence to get the next seq number
+        next_seq = None
 
         new_hint = BotLocalityHintEntry(
             bot_name=bot_name,
@@ -2960,9 +2936,6 @@ class Scheduler:
         # `greatest(0,_)` is needed if this feature is released when there are already running jobs
         # TODO: remove the safe-guard after the next minor version bump
         new_usage: Any = func.greatest(0, InstanceQuota.current_usage + delta)
-        if self._sql.dialect == "sqlite":
-            # SQLite does not support `greatest`, so we use a simpler update for it.
-            new_usage = InstanceQuota.current_usage + delta
 
         update_usage_query = (
             update(InstanceQuota)

{buildgrid-0.2.52 → buildgrid-0.3.0}/buildgrid/server/scheduler/notifier.py

@@ -25,7 +25,7 @@ from sqlalchemy import select as sql_select
 from sqlalchemy.orm import Session
 
 from buildgrid.server.logging import buildgrid_logger
-from buildgrid.server.sql.models import BotEntry, JobEntry
+from buildgrid.server.sql.models import BotEntry
 from buildgrid.server.sql.provider import SqlProvider
 from buildgrid.server.threading import ContextWorker
 
@@ -79,10 +79,7 @@ class Notifier(Generic[T]):
         while not shutdown_requested.is_set():
             try:
                 with self._sql.session() as session:
-                    if self._sql.dialect == "postgresql":
-                        self._listen_for_updates(shutdown_requested, session)
-                    else:
-                        self._poll_for_updates(shutdown_requested, session)
+                    self._listen_for_updates(shutdown_requested, session)
             except Exception as e:
                 LOGGER.warning(
                     f"OperationsNotifier encountered exception: {e}.",
@@ -122,23 +119,6 @@ class Notifier(Generic[T]):
                 notify = dbapi_connection.notifies.pop()
                 self.notify(notify.payload)
 
-    def _poll(self, names: list[str], session: Session) -> dict[str, T]:
-        raise NotImplementedError()
-
-    def _poll_for_updates(self, shutdown_requested: Event, session: Session) -> None:
-        prev_data: dict[str, T] = {}
-        while not shutdown_requested.is_set():
-            with self._lock:
-                names = list(self._listeners)
-
-            next_data: dict[str, T] = self._poll(names, session)
-            for name in next_data:
-                if name not in prev_data or prev_data[name] != next_data[name]:
-                    self.notify(name)
-
-            prev_data = next_data
-            shutdown_requested.wait(timeout=self.poll_interval)
-
     def notify(self, listener_name: str) -> None:
         with self._lock:
             if listener_name in self._listeners:
@@ -179,27 +159,11 @@ class OperationsNotifier(Notifier[tuple[bool, int]]):
         """
         super().__init__(sql_provider, NotificationChannel.JOB_UPDATED, "OperationsNotifier", poll_interval)
 
-    def _poll(self, names: list[str], session: Session) -> dict[str, tuple[bool, int]]:
-        # Only query for the minimal amount of data required.
-        # The subscribers can choose how they want to act (e.g. by querying the full job data).
-        statement = sql_select(JobEntry.name, JobEntry.cancelled, JobEntry.stage).where(JobEntry.name.in_(names))
-        next_data: dict[str, tuple[bool, int]] = {}
-        for [name, cancelled, stage] in session.execute(statement).all():
-            next_data[name] = (cancelled, stage)
-        return next_data
-
 
 class BotNotifier(Notifier[str]):
     def __init__(self, sql_provider: SqlProvider, poll_interval: float = 1.0) -> None:
         super().__init__(sql_provider, NotificationChannel.JOB_ASSIGNED, "BotNotifier", poll_interval)
 
-    def _poll(self, names: list[str], session: Session) -> dict[str, str]:
-        statement = sql_select(BotEntry.name, BotEntry.lease_id).where(BotEntry.name.in_(names))
-        next_data = {}
-        for name, lease_id in session.execute(statement).all():
-            next_data[name] = lease_id
-        return next_data
-
     def listener_count_for_instance(self, instance_name: str) -> int:
         with self._lock:
             stmt = sql_select(func.count(BotEntry.name)).where(

{buildgrid-0.2.52 → buildgrid-0.3.0}/buildgrid/server/sql/alembic/versions/0c17a7cb2bc5_initial_database_state.py

@@ -119,7 +119,6 @@ def upgrade() -> None:
         ["worker_completed_timestamp"],
         unique=False,
         postgresql_where=sa.text("worker_completed_timestamp IS NOT NULL"),
-        sqlite_where=sa.text("worker_completed_timestamp IS NOT NULL"),
     )
     op.create_index(
         "ix_worker_start_timestamp",
@@ -127,7 +126,6 @@ def upgrade() -> None:
         ["worker_start_timestamp"],
         unique=False,
         postgresql_where=sa.text("worker_start_timestamp IS NOT NULL"),
-        sqlite_where=sa.text("worker_start_timestamp IS NOT NULL"),
     )
     op.create_table(
         "platform_properties",
@@ -220,13 +218,11 @@ def downgrade() -> None:
         "ix_worker_start_timestamp",
         table_name="jobs",
         postgresql_where=sa.text("worker_start_timestamp IS NOT NULL"),
-        sqlite_where=sa.text("worker_start_timestamp IS NOT NULL"),
     )
     op.drop_index(
         "ix_worker_completed_timestamp",
         table_name="jobs",
         postgresql_where=sa.text("worker_completed_timestamp IS NOT NULL"),
-        sqlite_where=sa.text("worker_completed_timestamp IS NOT NULL"),
     )
     op.drop_index(op.f("ix_jobs_worker_name"), table_name="jobs")
     op.drop_index("ix_jobs_stage_property_label", table_name="jobs")

{buildgrid-0.2.52 → buildgrid-0.3.0}/buildgrid/server/sql/models.py

@@ -37,7 +37,8 @@ from buildgrid.server.enums import LeaseState, OperationStage
 
 bigint = Annotated[int, "bigint"]
 # This gives us something to reference in the type_annotation_map to specify the JSONB variant when
-# using postgresql. When we drop SQLite support this won't be necessary versus just using JSONB in
+# using postgresql.
+# TODO now SQLite support has been dropped this won't be necessary versus just using JSONB in
 # the model directly.
 json = Annotated[JSON, "json"]
 
@@ -133,7 +134,6 @@ class JobEntry(Base):
             "worker_completed_timestamp",
             unique=False,
             postgresql_where=worker_completed_timestamp.isnot(None),
-            sqlite_where=worker_completed_timestamp.isnot(None),
         ),
         Index(
             "ix_jobs_property_label_stage",