Repository: googleapis/genai-toolbox
File: internal/sources/bigquery/bigquery.go (Go)

Code before:
		client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount)
		if err != nil {
			return nil, fmt.Errorf("error creating client from ADC: %w", err)
		}
	}

	allowedDatasets := make(map[string]struct{})
	// Get full id of allowed datasets and verify they exist.
	if len(r.AllowedDatasets) > 0 {
		for _, allowed := range r.AllowedDatasets {
			var projectID, datasetID, allowedFullID string
			if strings.Contains(allowed, ".") {
				parts := strings.Split(allowed, ".")
				if len(parts) != 2 {
					return nil, fmt.Errorf("invalid allowedDataset format: %q, expected 'project.dataset' or 'dataset'", allowed)
				}
				projectID = parts[0]
				datasetID = parts[1]
				allowedFullID = allowed
			} else {
				projectID = client.Project()
				datasetID = allowed
				allowedFullID = fmt.Sprintf("%s.%s", projectID, datasetID)
			}

			dataset := client.DatasetInProject(projectID, datasetID)
			_, err := dataset.Metadata(ctx)
			if err != nil {
				if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == http.StatusNotFound {
					return nil, fmt.Errorf("allowedDataset '%s' not found in project '%s'", datasetID, projectID)
				}
				return nil, fmt.Errorf("failed to verify allowedDataset '%s' in project '%s': %w", datasetID, projectID, err)
			}
			allowedDatasets[allowedFullID] = struct{}{}
		}
	}

	s := &Source{
		Name:                      r.Name,
		Kind:                      SourceKind,
		Project:                   r.Project,
		Location:                  r.Location,
		Client:                    client,
		RestService:               restService,
		TokenSource:               tokenSource,
		MaxQueryResultRows:        50,
		ClientCreator:             clientCreator,
		AllowedDatasets:           allowedDatasets,
		UseClientOAuth:            r.UseClientOAuth,
		ImpersonateServiceAccount: r.ImpersonateServiceAccount,
	}

Review comment (security):
Maybe add a code block to check if useClientOAuth and ImpersonateServiceAccount are provided by user at the same time. If so, we can return an error.

Code after:
	if r.WriteMode == "" {
		r.WriteMode = WriteModeAllowed
	}

	if r.WriteMode == WriteModeProtected && r.UseClientOAuth {
		return nil, fmt.Errorf("writeMode 'protected' cannot be used with useClientOAuth 'true'")
	}

	if r.UseClientOAuth && r.ImpersonateServiceAccount != "" {
		return nil, fmt.Errorf("useClientOAuth cannot be used with impersonateServiceAccount")
	}

	var client *bigqueryapi.Client
	var restService *bigqueryrestapi.Service
	var tokenSource oauth2.TokenSource
	var clientCreator BigqueryClientCreator
	var err error

	if r.UseClientOAuth {
		clientCreator, err = newBigQueryClientCreator(ctx, tracer, r.Project, r.Location, r.Name)
		if err != nil {
			return nil, fmt.Errorf("error constructing client creator: %w", err)
		}
	} else {
		// Initializes a BigQuery Google SQL source
		client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount)
		if err != nil {
			return nil, fmt.Errorf("error creating client from ADC: %w", err)
		}
	}

	allowedDatasets := make(map[string]struct{})
	// Get full id of allowed datasets and verify they exist.
	if len(r.AllowedDatasets) > 0 {
		for _, allowed := range r.AllowedDatasets {
			var projectID, datasetID, allowedFullID string
			if strings.Contains(allowed, ".") {
				parts := strings.Split(allowed, ".")
				if len(parts) != 2 {
					return nil, fmt.Errorf("invalid allowedDataset format: %q, expected 'project.dataset' or 'dataset'", allowed)
				}
				projectID = parts[0]
				datasetID = parts[1]
				allowedFullID = allowed
			} else {
				projectID = r.Project
				datasetID = allowed
				allowedFullID = fmt.Sprintf("%s.%s", projectID, datasetID)
			}

			if client != nil {