Merge branch 'feature/use-personal-branch' of github.com:dbanalyticsc…

…o/spectacles into feature/use-personal-branch

spectacles/cli.py

@@ -17,7 +17,11 @@
 
 import spectacles.printer as printer
 import spectacles.tracking as tracking
-from spectacles.client import DEFAULT_API_VERSION, LookerClient
+from spectacles.client import (
+    DEFAULT_API_VERSION,
+    LOOKML_VALIDATION_TIMEOUT,
+    LookerClient,
+)
 from spectacles.exceptions import (
     GenericValidationError,
     LookerApiError,
@@ -27,6 +31,7 @@
 from spectacles.logger import set_file_handler
 from spectacles.runner import Runner
 from spectacles.utils import log_duration
+from spectacles.validators.data_test import DATA_TEST_CONCURRENCY
 
 __version__ = importlib.metadata.version("spectacles")
 
@@ -343,6 +348,7 @@ def main() -> None:
                 remote_reset=args.remote_reset,
                 pin_imports=pin_imports,
                 use_personal_branch=args.use_personal_branch,
+                concurrency=args.concurrency,
             )
         )
     elif args.command == "content":
@@ -379,6 +385,7 @@ def main() -> None:
                 severity=args.severity,
                 pin_imports=pin_imports,
                 use_personal_branch=args.use_personal_branch,
+                timeout=args.timeout,
             )
         )
 
@@ -612,6 +619,12 @@ def _build_lookml_subparser(
             "validator to fail. The default is 'warning'."
         ),
     )
+    subparser.add_argument(
+        "--timeout",
+        type=int,
+        default=LOOKML_VALIDATION_TIMEOUT,
+        help="Specify the timeout for the LookML validation in seconds.",
+    )
     _build_validator_subparser(subparser_action, subparser)
 
 
@@ -719,6 +732,16 @@ def _build_assert_subparser(
     _build_validator_subparser(subparser_action, subparser)
     _build_select_subparser(subparser_action, subparser)
 
+    subparser.add_argument(
+        "--concurrency",
+        type=int,
+        default=DATA_TEST_CONCURRENCY,
+        help=(
+            "Specify the number of concurrent queries you want to have running "
+            f"against your data warehouse. The default is {DATA_TEST_CONCURRENCY}."
+        ),
+    )
+
 
 def _build_content_subparser(
     subparser_action: argparse._SubParsersAction,  # type: ignore[type-arg]
@@ -790,6 +813,7 @@ async def run_lookml(
     severity: str,
     pin_imports: Dict[str, str],
     use_personal_branch: bool,
+    timeout: int,
 ) -> None:
     # Don't trust env to ignore .netrc credentials
     async_client = httpx.AsyncClient(trust_env=False)
@@ -799,7 +823,7 @@ async def run_lookml(
         )
         runner = Runner(client, project, remote_reset, pin_imports, use_personal_branch)
 
-        results = await runner.validate_lookml(ref, severity)
+        results = await runner.validate_lookml(ref, severity, timeout)
     finally:
         await async_client.aclose()
 
@@ -909,6 +933,7 @@ async def run_assert(
     remote_reset: bool,
     pin_imports: Dict[str, str],
     use_personal_branch: bool,
+    concurrency: int,
 ) -> None:
     # Don't trust env to ignore .netrc credentials
     async_client = httpx.AsyncClient(trust_env=False)
@@ -918,7 +943,7 @@ async def run_assert(
         )
         runner = Runner(client, project, remote_reset, pin_imports, use_personal_branch)
 
-        results = await runner.validate_data_tests(ref, filters)
+        results = await runner.validate_data_tests(ref, filters, concurrency)
     finally:
         await async_client.aclose()
 

spectacles/client.py

@@ -15,6 +15,7 @@
 
 DEFAULT_API_VERSION = 4.0
 TIMEOUT_SEC = 300
+LOOKML_VALIDATION_TIMEOUT = 7200
 MAX_ASYNC_CONNECTIONS = 200
 DEFAULT_MAX_TRIES = 3
 BACKOFF_EXCEPTIONS = (
@@ -862,10 +863,14 @@ async def content_validation(self) -> JsonDict:
         return result  # type: ignore[no-any-return]
 
     @backoff.on_exception(backoff.expo, BACKOFF_EXCEPTIONS, max_tries=DEFAULT_MAX_TRIES)
-    async def lookml_validation(self, project: str) -> JsonDict:
-        logger.debug(f"Validating LookML for project '{project}'")
+    async def lookml_validation(
+        self, project: str, timeout: int = LOOKML_VALIDATION_TIMEOUT
+    ) -> JsonDict:
+        logger.debug(
+            f"Validating LookML for project '{project}' with timeout {timeout} seconds."
+        )
         url = utils.compose_url(self.api_url, path=["projects", project, "validate"])
-        response = await self.post(url=url, timeout=7200)
+        response = await self.post(url=url, timeout=timeout)
 
         try:
             response.raise_for_status()

spectacles/runner.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Set, Tuple
 
-from spectacles.client import LookerClient
+from spectacles.client import LOOKML_VALIDATION_TIMEOUT, LookerClient
 from spectacles.exceptions import LookerApiError, SpectaclesException, SqlError
 from spectacles.logger import GLOBAL_LOGGER as logger
 from spectacles.lookml import CompiledSql, Explore, build_project
@@ -17,6 +17,7 @@
     LookMLValidator,
     SqlValidator,
 )
+from spectacles.validators.data_test import DATA_TEST_CONCURRENCY
 from spectacles.validators.sql import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_QUERY_CONCURRENCY,
@@ -486,6 +487,7 @@ async def validate_data_tests(
         self,
         ref: Optional[str] = None,
         filters: Optional[List[str]] = None,
+        concurrency: int = DATA_TEST_CONCURRENCY,
     ) -> JsonDict:
         if filters is None:
             filters = ["*/*"]
@@ -504,16 +506,21 @@ async def validate_data_tests(
                 f"{'explore' if explore_count == 1 else 'explores'}"
             )
             tests = await validator.get_tests(project)
-            await validator.validate(tests)
+            await validator.validate(tests, concurrency)
 
         results = project.get_results(validator="data_test")
         return results
 
-    async def validate_lookml(self, ref: Optional[str], severity: str) -> JsonDict:
+    async def validate_lookml(
+        self,
+        ref: Optional[str],
+        severity: str,
+        timeout: int = LOOKML_VALIDATION_TIMEOUT,
+    ) -> JsonDict:
         async with self.branch_manager(ref=ref):
             validator = LookMLValidator(self.client)
             print_header(f"Validating LookML in project {self.project} [{severity}]")
-            results = await validator.validate(self.project, severity)
+            results = await validator.validate(self.project, severity, timeout)
         return results
 
     async def validate_content(

spectacles/validators/data_test.py

@@ -6,7 +6,9 @@
 from spectacles.exceptions import DataTestError, SpectaclesException
 from spectacles.lookml import Explore, Project
 
-QUERY_SLOT_LIMIT = 15  # This is the per-user query limit in Looker for most instances
+DATA_TEST_CONCURRENCY = (
+    15  # This is the per-user query limit in Looker for most instances
+)
 
 
 @dataclass
@@ -94,9 +96,11 @@ async def get_tests(self, project: Project) -> List[DataTest]:
 
         return selected_tests
 
-    async def validate(self, tests: List[DataTest]) -> List[DataTestError]:
+    async def validate(
+        self, tests: List[DataTest], concurrency: int = DATA_TEST_CONCURRENCY
+    ) -> List[DataTestError]:
         data_test_errors: List[DataTestError] = []
-        query_slot = asyncio.Semaphore(QUERY_SLOT_LIMIT)
+        query_slot = asyncio.Semaphore(concurrency)
 
         async def run_test(test: DataTest, query_slot: asyncio.Semaphore) -> None:
             async with query_slot:

spectacles/validators/lookml.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional
 
-from spectacles.client import LookerClient
+from spectacles.client import LOOKML_VALIDATION_TIMEOUT, LookerClient
 from spectacles.exceptions import LookMLError
 
 # Define constants for severity levels
@@ -31,11 +31,16 @@ class LookMLValidator:
     def __init__(self, client: LookerClient):
         self.client = client
 
-    async def validate(self, project: str, severity: str = "warning") -> Dict[str, Any]:
+    async def validate(
+        self,
+        project: str,
+        severity: str = "warning",
+        timeout: int = LOOKML_VALIDATION_TIMEOUT,
+    ) -> Dict[str, Any]:
         severity_level: int = NAME_TO_LEVEL[severity]
         validation_results = await self.client.cached_lookml_validation(project)
         if not validation_results or validation_results.get("stale"):
-            validation_results = await self.client.lookml_validation(project)
+            validation_results = await self.client.lookml_validation(project, timeout)
         errors = []
         lookml_url: Optional[str] = None
         for error in validation_results["errors"]: