feat(server): allow partial source startup checks

[Deeven-Seru]

Summary

• add checkAtStartup per-source config to skip startup connectivity checks
• add --allow-partial-sources flag to skip all source checks at startup
• make source initialization respect startup-check flag via context
• update docs + add regression test

Testing

• go test ./cmd/internal -run TestParseToolFileCheckAtStartup
• go test ./internal/server
• go test ./internal/sources/...

Fixes #2341

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the server's resilience and flexibility by allowing partial source availability at startup. It introduces mechanisms to bypass initial connectivity checks for data sources, either globally via a command-line flag or on a per-source basis through configuration. This prevents the server from failing to start due to transient network issues or unavailable dependencies, making it more robust in dynamic environments. Additionally, it includes a security improvement for HTTP tool URL handling.

Highlights

• Configurable Startup Connectivity Checks: Introduced a new CLI flag --allow-partial-sources to globally skip startup connectivity checks for all sources, and a checkAtStartup: false configuration option per source to skip checks individually.
• Context-aware Source Initialization: Modified source initialization logic to respect the new startup check flags, allowing the server to start even if some sources are temporarily unavailable.
• Enhanced URL Parsing for HTTP Tools: Refactored HTTP tool URL parsing to prevent host overrides from relative paths, improving security and predictability.
• Documentation and Testing: Updated CLI and source configuration documentation to reflect the new features, and added a regression test for the checkAtStartup functionality.

Changelog

• cmd/internal/persistent_flags.go
  • Added a new boolean flag --allow-partial-sources to skip all source startup connectivity checks.
• cmd/internal/tools_file_test.go
  • Imported the internal/sources package.
  • Added TestParseToolFileCheckAtStartup to verify the parsing and behavior of the checkAtStartup configuration for sources.
• cmd/root.go
  • Modified handleDynamicReload to accept and pass the allowPartialSources flag.
  • Updated validateReloadEdits to receive and utilize the allowPartialSources flag during configuration initialization.
  • Passed the AllowPartialSources option to the watchChanges function.
• docs/en/reference/cli.md
  • Documented the new --allow-partial-sources CLI flag in the reference documentation.
• docs/en/resources/sources/_index.md
  • Added documentation for the checkAtStartup configuration option for individual sources.
• internal/server/config.go
  • Added AllowPartialSources field to the ServerConfig struct.
  • Modified UnmarshalYAMLSourceConfig to parse the checkAtStartup field from source configurations and wrap the source config accordingly.
• internal/server/server.go
  • Updated InitializeConfigs to check the AllowPartialSources flag and conditionally skip startup checks for all sources.
  • Integrated sources.WithStartupCheck into source initialization to pass the checkAtStartup context.
• internal/sources/alloydbpg/alloydb_pg.go
  • Wrapped the pool.Ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/clickhouse/clickhouse.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/cloudsqlmssql/cloud_sql_mssql.go
  • Wrapped the db.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/cloudsqlmysql/cloud_sql_mysql.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/cloudsqlpg/cloud_sql_pg.go
  • Wrapped the pool.Ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/cockroachdb/cockroachdb.go
  • Wrapped the pool.Ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/firebird/firebird.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/mindsdb/mindsdb.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/mongodb/mongodb.go
  • Wrapped the client.Ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/mssql/mssql.go
  • Wrapped the db.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/mysql/mysql.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/oceanbase/oceanbase.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/oracle/oracle.go
  • Wrapped the db.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/postgres/postgres.go
  • Wrapped the pool.Ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/redis/redis.go
  • Removed an unused err variable.
  • Wrapped Redis client ping calls (cluster and standalone) with sources.CheckConnectivity to conditionally skip startup checks.
• internal/sources/singlestore/singlestore.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/snowflake/snowflake.go
  • Wrapped the db.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/sources.go
  • Introduced StartupCheckConfig interface and startupCheckWrapper struct for managing startup checks.
  • Added WrapSourceConfigWithStartupCheck to apply startup check overrides to source configurations.
  • Implemented CheckAtStartup to determine if a source config requires startup checks.
  • Defined startupCheckKey and WithStartupCheck to pass startup check preferences via context.
  • Provided ShouldCheckAtStartup to retrieve startup check preference from context.
  • Created CheckConnectivity to conditionally execute a connectivity check function based on context.
• internal/sources/sqlite/sqlite.go
  • Wrapped the db.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/tidb/tidb.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/trino/trino.go
  • Wrapped the pool.PingContext call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/valkey/valkey.go
  • Wrapped the client ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/sources/yugabytedb/yugabytedb.go
  • Wrapped the pool.Ping call with sources.CheckConnectivity to conditionally skip the startup check.
• internal/tools/http/http.go
  • Refactored getURL to robustly parse base and relative URLs, preventing host overrides.
  • Added validation to ensure the base URL includes a scheme and host, and that the path is relative.
• internal/tools/http/http_url_test.go
  • Added a new test file http_url_test.go.
  • Implemented TestGetURLHostOverride to verify the new URL parsing logic correctly handles and rejects host overrides in paths.

Activity

• The pull request addresses and fixes issue Skip connectivity checks on start up #2341.
• The author provided specific testing instructions for cmd/internal, internal/server, and internal/sources/... to validate the changes.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a valuable feature to allow the server to start even if some data sources are unavailable, utilizing a new --allow-partial-sources flag and a per-source checkAtStartup configuration. While it includes a security hardening measure in the HTTP tool to prevent host and scheme overrides, the new URL construction logic remains vulnerable to path traversal attacks within the same host. This should be mitigated by validating the resolved path or properly escaping path parameters. Additionally, a suggestion is made to improve error handling in the valkey source to prevent fatal errors on connection failure, aligning it with the PR's goal.

[Deeven-Seru]

Update: rebased this PR onto main to keep it focused on startup checks only. The HTTP tool host-override changes live in #2655, so the path-traversal comment is now out of scope here. Also updated valkey init to return an error instead of log.Fatalf, so allow-partial-sources can skip it cleanly.

[gemini-code-assist]

Code Review

This pull request introduces a valuable feature for allowing partial source availability during server startup, controlled by both a global flag and per-source configurations. The implementation is clean and robust, utilizing context to propagate settings and a centralized helper function for connectivity checks. The changes are consistently applied across all data sources, and the documentation has been updated accordingly. I appreciate that you also took the opportunity to fix some related issues, such as replacing log.Fatalf with proper error returns and using the correct context for ping operations. I have a couple of minor suggestions in internal/sources/redis/redis.go to improve error wrapping for consistency with other parts of the codebase.

[Deeven-Seru]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a feature to allow partial source startup by skipping connectivity checks. This is controlled by a global --allow-partial-sources flag and a per-source checkAtStartup configuration option. The implementation correctly uses context to propagate the check setting and refactors all data sources to use a centralized CheckConnectivity helper, which is a great improvement. The documentation and tests have also been updated. My feedback includes a suggestion to enhance test coverage for the new configuration option and a potential refactoring to simplify how the static configuration is handled during dynamic reloads.

[Deeven-Seru]

Updates pushed:

• Expanded TestParseToolFileCheckAtStartup to cover explicit true + invalid type cases.
• Kept allowPartialSources as a parameter for now to avoid a broader Server struct refactor; happy to follow up if you’d prefer the config to live on server.Server.

[Deeven-Seru]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a feature to allow the server to start even if some data sources are unavailable. This is achieved by adding a global --allow-partial-sources flag and a per-source checkAtStartup configuration option. The core logic is implemented via a new CheckConnectivity helper in the internal/sources package, which is now used by all data source initialization methods to conditionally perform connectivity checks. The changes also include corresponding documentation updates and a new test for parsing the per-source configuration. The implementation is robust and consistently applied across all relevant parts of the codebase, and it also includes beneficial refactoring, such as replacing a log.Fatalf call with proper error handling in the Valkey source.

_{Note: Security Review did not run due to the size of the PR.}

[averikitsch]

Thank you for the PR. We will review it shortly.

[Deeven-Seru]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a valuable feature to allow the server to start even if some data sources are unavailable, controlled by a new --allow-partial-sources flag and a per-source checkAtStartup option. The implementation is well-designed, centralizing the logic in internal/sources/sources.go and consistently applying it across all source initializations by refactoring them to use a new sources.CheckConnectivity helper. The changes are accompanied by documentation updates and a regression test. The overall implementation is solid and I have no comments for improvement.