feat: support datasource in count

cmd/flags/common.go

@@ -77,3 +77,7 @@ func AddDirectoryPrefixFlag(cmd *cobra.Command) {
 func AddRefangFlag(cmd *cobra.Command) {
 	cmd.Flags().Bool("refang", false, "Refang an input (convert '[.]' back to '.' and so on)")
 }
+
+func AddDatasourceFlag(cmd *cobra.Command) {
+	cmd.Flags().StringP("datasource", "D", "scans", "Datasources to search: scans (urlscan.io), hostnames, incidents, notifications, certificates (urlscan Pro)")
+}

cmd/search/count.go

@@ -6,6 +6,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/urlscan/urlscan-cli/api"
+	"github.com/urlscan/urlscan-cli/cmd/flags"
 	"github.com/urlscan/urlscan-cli/pkg/utils"
 )
 
@@ -24,6 +25,8 @@ var countCmd = &cobra.Command{
 			return cmd.Usage()
 		}
 
+		datasource, _ := cmd.Flags().GetString("datasource")
+
 		reader := utils.StringReaderFromCmdArgs(args)
 		q, err := reader.ReadString()
 		if err != nil {
@@ -35,7 +38,7 @@ var countCmd = &cobra.Command{
 			return err
 		}
 
-		it, err := client.Search(q, api.IteratorSize(0))
+		it, err := client.Search(q, api.IteratorSize(0), api.IteratorDatasource(datasource))
 		if err != nil {
 			return err
 		}
@@ -53,5 +56,6 @@ var countCmd = &cobra.Command{
 }
 
 func init() {
+	flags.AddDatasourceFlag(countCmd)
 	RootCmd.AddCommand(countCmd)
 }

cmd/search/root.go

@@ -95,7 +95,7 @@ func init() {
 	flags.AddSizeFlag(RootCmd, 100) // non-pro user's max size is 100
 	flags.AddLimitFlag(RootCmd)
 	flags.AddAllFlag(RootCmd)
+	flags.AddDatasourceFlag(RootCmd)
 	RootCmd.Flags().String("search-after", "", "For retrieving the next batch of results, value of the sort attribute of the last (oldest) result you received (comma-separated)")
-	RootCmd.Flags().StringP("datasource", "D", "scans", "Datasources to search: scans (urlscan.io), hostnames, incidents, notifications, certificates (urlscan Pro)")
 	RootCmd.Flags().StringP("collapse", "c", "", "Field to collapse results on")
 }

docs/urlscan_search_count.md

@@ -16,7 +16,8 @@ urlscan search count <query> [flags]
 ### Options
 
 ```
-  -h, --help   help for count
+  -D, --datasource string   Datasources to search: scans (urlscan.io), hostnames, incidents, notifications, certificates (urlscan Pro) (default "scans")
+  -h, --help                help for count
 ```
 
 ### SEE ALSO