buildgrid 0.2.53__tar.gz → 0.3.1__tar.gz

{buildgrid-0.2.53 → buildgrid-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildgrid
-Version: 0.2.53
+Version: 0.3.1
 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.53 → buildgrid-0.3.1}/buildgrid/server/app/settings/parser.py

@@ -1025,7 +1025,7 @@ def load_sql_scheduler(
         sql_ro = sql_ro or sql
         sql_notifier = sql_notifier or sql
 
-        logstream_url, logstream_credentials, logstream_instance = get_logstream_connection_info(logstream)
+        logstream_url, logstream_credentials = get_logstream_connection_info(logstream)
         logstream_channel: grpc.Channel | None = None
         if logstream_url is not None:
             logstream_credentials = logstream_credentials or {}
@@ -1081,7 +1081,6 @@ def load_sql_scheduler(
             ),
             bot_session_keepalive_timeout=bot_session_keepalive_timeout,
             logstream_channel=logstream_channel,
-            logstream_instance=logstream_instance,
             asset_client=asset_client,
             queued_action_retention_hours=queued_action_retention_hours,
             completed_action_retention_hours=completed_action_retention_hours,
@@ -2150,7 +2149,6 @@ def load_metering_service_client(
 def load_asset_client(
     url: str,
     credentials: ClientCredentials | None = None,
-    instance_name: str = "",
     request_timeout: float = 5.0,
     retries: int = 3,
 ) -> AssetClient:
@@ -2162,7 +2160,6 @@ def load_asset_client(
 
             - !asset-client
               url: https://remote-asset.com
-              instance-name: dev
               credentials:
                 tls-client-cert: /path/to/cert
                 auth-token: /path/to/token
@@ -2179,7 +2176,7 @@ def load_asset_client(
         server_cert=credentials.get("tls-server-cert"),
         timeout=request_timeout,
     )
-    return AssetClient(channel=channel, instance_name=instance_name, retries=retries)
+    return AssetClient(channel=channel, retries=retries)
 
 
 @object_tag("!introspection")
@@ -2378,18 +2375,16 @@ def _validate_server_credentials(credentials: dict[str, str] | None) -> None:
         sys.exit(-1)
 
 
-def get_logstream_connection_info(logstream: Any) -> tuple[str | None, dict[str, str] | None, str | None]:
+def get_logstream_connection_info(logstream: Any) -> tuple[str | None, dict[str, str] | None]:
     logstream_url = None
     credentials = None
-    logstream_instance_name = None
     if logstream:
         logstream_url = logstream["url"]
         credentials = logstream.get("credentials")
         if not _validate_url_and_credentials(logstream_url, credentials=credentials):
             sys.exit(-1)
-        logstream_instance_name = logstream.get("instance-name", "")
 
-    return logstream_url, credentials, logstream_instance_name
+    return logstream_url, credentials
 
 
 def get_schema(strict: bool = False) -> Any:

{buildgrid-0.2.53 → buildgrid-0.3.1}/buildgrid/server/app/settings/schema.yml

@@ -472,7 +472,7 @@ definitions:
       priority-assignment-percentage:   { type: number, minimum: 0 }
       metering-service-client:          { "$ref": "#/definitions/metering-service-client" }
       metering-throttle-action:         { type: string, enum: [deprioritize, reject] }
-      logstream:                        { "$ref": "#/definitions/grpc-connection-options" }
+      logstream:                        { "$ref": "#/definitions/logstream-connection-options" }
       asset-client:                     { "$ref": "#/definitions/asset-client" }
       action-browser-url:               { type: string }
       poll-interval:                    { type: number, minimum: 0 }
@@ -645,12 +645,11 @@ definitions:
     properties:
       kind:            { type: string, enum: [ "!asset-client" ] }
       url:             { type: string }
-      instance-name:   { type: string }
       channel-options: { type: object, propertyNames: { pattern: "^[a-z0-9-]+$" } }
       request-timeout: { type: number }
       credentials:     { "$ref": "#/definitions/grpc-credentials" }
       retries:         { type: number }
-    required: [kind, url, instance-name]
+    required: [kind, url]
 
   metering-service-client:
     type: object
@@ -783,16 +782,15 @@ definitions:
       retry-limit:               { type: integer, min: 0 }
     required: [instance-name, index, high-watermark, low-watermark, batch-size]
 
-  grpc-connection-options:
+  logstream-connection-options:
     type: object
     propertyNames: { pattern: "^[A-Za-z0-9-]*$" }
     properties:
       url:              { type: string }
-      instance-name:    { type: string }
       channel-options:  { type: object, propertyNames: { pattern: "^[a-z0-9-]+$" } }
       request-timeout:  { type: number }
       credentials:      { "$ref": "#/definitions/grpc-credentials" }
-    required: [url, instance-name]
+    required: [url]
 
   grpc-credentials:
     type: object

{buildgrid-0.2.53 → buildgrid-0.3.1}/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.53 → buildgrid-0.3.1}/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.53 → buildgrid-0.3.1}/buildgrid/server/client/asset.py

@@ -41,13 +41,11 @@ class AssetClient:
     def __init__(
         self,
         channel: grpc.Channel,
-        instance_name: str,
         retries: int = 0,
         max_backoff: int = 64,
         should_backoff: bool = True,
     ) -> None:
         self._channel = channel
-        self._instance_name = instance_name
         self._push_stub = PushStub(channel)
         self._fetch_stub = FetchStub(channel)
         self._retrier = GrpcRetrier(retries=retries, max_backoff=max_backoff, should_backoff=should_backoff)
@@ -68,6 +66,7 @@ class AssetClient:
         expire_at: datetime | None = None,
         referenced_blobs: Iterable[Digest] = [],
         referenced_directories: Iterable[Digest] = [],
+        instance_name: str,
     ) -> PushBlobResponse:
         def _push_blob() -> PushBlobResponse:
             qualifiers_pb = [Qualifier(name=name, value=value) for name, value in qualifiers.items()]
@@ -77,7 +76,7 @@ class AssetClient:
                 expire_at_pb.FromDatetime(expire_at)
 
             request = PushBlobRequest(
-                instance_name=self._instance_name,
+                instance_name=instance_name,
                 uris=uris,
                 qualifiers=qualifiers_pb,
                 expire_at=expire_at_pb,
@@ -98,6 +97,7 @@ class AssetClient:
         expire_at: datetime | None = None,
         referenced_blobs: Iterable[Digest] = [],
         referenced_directories: Iterable[Digest] = [],
+        instance_name: str,
     ) -> PushDirectoryResponse:
         def _push_directory() -> PushDirectoryResponse:
             qualifiers_pb = [Qualifier(name=name, value=value) for name, value in qualifiers.items()]
@@ -107,7 +107,7 @@ class AssetClient:
                 expire_at_pb.FromDatetime(expire_at)
 
             request = PushDirectoryRequest(
-                instance_name=self._instance_name,
+                instance_name=instance_name,
                 uris=uris,
                 qualifiers=qualifiers_pb,
                 expire_at=expire_at_pb,

{buildgrid-0.2.53 → buildgrid-0.3.1}/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
@@ -218,7 +218,6 @@ class Scheduler:
         metering_throttle_action: MeteringThrottleAction | None = None,
         bot_session_keepalive_timeout: int = 600,
         logstream_channel: Channel | None = None,
-        logstream_instance: str | None = None,
         asset_client: AssetClient | None = None,
         queued_action_retention_hours: float | None = None,
         completed_action_retention_hours: float | None = None,
@@ -268,7 +267,6 @@ class Scheduler:
         self.metering_throttle_action = metering_throttle_action or MeteringThrottleAction.DEPRIORITIZE
         self.bot_session_keepalive_timeout = bot_session_keepalive_timeout
         self.logstream_channel = logstream_channel
-        self.logstream_instance = logstream_instance
         self.asset_client = asset_client
         self.queued_action_retention_hours = queued_action_retention_hours
         self.completed_action_retention_hours = completed_action_retention_hours
@@ -414,7 +412,9 @@ class Scheduler:
                 LOGGER.exception("Checking ActionCache for action failed.", tags=dict(digest=action_digest))
 
         # Extend retention for action
-        self._update_action_retention(action, action_digest, self.queued_action_retention_hours)
+        self._update_action_retention(
+            action, action_digest, self.queued_action_retention_hours, instance_name=current_instance()
+        )
 
         return self.create_operation_for_new_job(
             action=action,
@@ -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.")
 
@@ -1999,11 +1975,11 @@ class Scheduler:
         return None
 
     def _create_logstream_for_job(self, job: JobEntry, log_tags: Tags) -> None:
-        if self.logstream_channel and self.logstream_instance is not None:
+        if self.logstream_channel:
             try:
                 action_digest = string_to_digest(job.action_digest)
                 parent_base = f"{action_digest.hash}_{action_digest.size_bytes}_{int(time())}"
-                with logstream_client(self.logstream_channel, self.logstream_instance) as ls_client:
+                with logstream_client(self.logstream_channel, job.instance_name) as ls_client:
                     stdout_stream = ls_client.create(f"{parent_base}_stdout")
                     stderr_stream = ls_client.create(f"{parent_base}_stderr")
                 job.stdout_stream_name = stdout_stream.name
@@ -2527,9 +2503,12 @@ class Scheduler:
             Action.FromString(job.action),
             string_to_digest(job.action_digest),
             retention_hours=self.completed_action_retention_hours,
+            instance_name=job.instance_name,
         )
         if action_result.ByteSize() > 0:
-            self._update_action_result_retention(action_result, retention_hours=self.action_result_retention_hours)
+            self._update_action_result_retention(
+                action_result, retention_hours=self.action_result_retention_hours, instance_name=job.instance_name
+            )
 
         worker_duration = None
         if job.worker_start_timestamp is not None and job.worker_completed_timestamp is not None:
@@ -2773,7 +2752,9 @@ class Scheduler:
         except Exception as exc:
             LOGGER.exception("Cannot publish resource usage.", tags=dict(job_name=job_name), exc_info=exc)
 
-    def _update_action_retention(self, action: Action, action_digest: Digest, retention_hours: float | None) -> None:
+    def _update_action_retention(
+        self, action: Action, action_digest: Digest, retention_hours: float | None, instance_name: str
+    ) -> None:
         if not self.asset_client or not retention_hours:
             return
         uri = DIGEST_URI_TEMPLATE.format(digest_hash=action_digest.hash)
@@ -2790,6 +2771,7 @@ class Scheduler:
                 expire_at=expire_at,
                 referenced_blobs=referenced_blobs,
                 referenced_directories=referenced_directories,
+                instance_name=instance_name,
             )
             LOGGER.debug(
                 "Extended the retention of action.", tags=dict(digest=action_digest, retention_hours=retention_hours)
@@ -2798,7 +2780,9 @@ class Scheduler:
             LOGGER.exception("Failed to push action as an asset.", tags=dict(digest=action_digest))
             # Not a fatal path, don't reraise here
 
-    def _update_action_result_retention(self, action_result: ActionResult, retention_hours: float | None) -> None:
+    def _update_action_result_retention(
+        self, action_result: ActionResult, retention_hours: float | None, instance_name: str
+    ) -> None:
         if not self.asset_client or not retention_hours:
             return
         digest = None
@@ -2837,6 +2821,7 @@ class Scheduler:
                 expire_at=expire_at,
                 referenced_blobs=referenced_blobs,
                 referenced_directories=referenced_directories,
+                instance_name=instance_name,
             )
             LOGGER.debug(
                 "Extended the retention of action result.", tags=dict(digest=digest, retention_hours=retention_hours)
@@ -2853,18 +2838,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,
@@ -2969,9 +2945,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.53 → buildgrid-0.3.1}/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(