pass organization id to launch browser (#1000)

This commit is contained in:
LawyZheng
2024-10-18 15:02:11 +08:00
committed by GitHub
parent f69016088b
commit 533c814558
2 changed files with 37 additions and 7 deletions

View File

@@ -208,6 +208,7 @@ class BrowserState:
proxy_location: ProxyLocation | None = None, proxy_location: ProxyLocation | None = None,
task_id: str | None = None, task_id: str | None = None,
workflow_run_id: str | None = None, workflow_run_id: str | None = None,
organization_id: str | None = None,
) -> None: ) -> None:
if self.pw is None: if self.pw is None:
LOG.info("Starting playwright") LOG.info("Starting playwright")
@@ -225,6 +226,7 @@ class BrowserState:
proxy_location=proxy_location, proxy_location=proxy_location,
task_id=task_id, task_id=task_id,
workflow_run_id=workflow_run_id, workflow_run_id=workflow_run_id,
organization_id=organization_id,
) )
self.browser_context = browser_context self.browser_context = browser_context
self.browser_artifacts = browser_artifacts self.browser_artifacts = browser_artifacts
@@ -312,6 +314,7 @@ class BrowserState:
proxy_location: ProxyLocation | None = None, proxy_location: ProxyLocation | None = None,
task_id: str | None = None, task_id: str | None = None,
workflow_run_id: str | None = None, workflow_run_id: str | None = None,
organization_id: str | None = None,
) -> Page: ) -> Page:
page = await self.get_working_page() page = await self.get_working_page()
if page is not None: if page is not None:
@@ -319,7 +322,11 @@ class BrowserState:
try: try:
await self.check_and_fix_state( await self.check_and_fix_state(
url=url, proxy_location=proxy_location, task_id=task_id, workflow_run_id=workflow_run_id url=url,
proxy_location=proxy_location,
task_id=task_id,
workflow_run_id=workflow_run_id,
organization_id=organization_id,
) )
except Exception as e: except Exception as e:
error_message = str(e) error_message = str(e)
@@ -327,14 +334,22 @@ class BrowserState:
raise e raise e
await self.close_current_open_page() await self.close_current_open_page()
await self.check_and_fix_state( await self.check_and_fix_state(
url=url, proxy_location=proxy_location, task_id=task_id, workflow_run_id=workflow_run_id url=url,
proxy_location=proxy_location,
task_id=task_id,
workflow_run_id=workflow_run_id,
organization_id=organization_id,
) )
await self.__assert_page() await self.__assert_page()
if not await BrowserContextFactory.validate_browser_context(await self.get_working_page()): if not await BrowserContextFactory.validate_browser_context(await self.get_working_page()):
await self.close_current_open_page() await self.close_current_open_page()
await self.check_and_fix_state( await self.check_and_fix_state(
url=url, proxy_location=proxy_location, task_id=task_id, workflow_run_id=workflow_run_id url=url,
proxy_location=proxy_location,
task_id=task_id,
workflow_run_id=workflow_run_id,
organization_id=organization_id,
) )
await self.__assert_page() await self.__assert_page()

View File

@@ -30,6 +30,7 @@ class BrowserManager:
url: str | None = None, url: str | None = None,
task_id: str | None = None, task_id: str | None = None,
workflow_run_id: str | None = None, workflow_run_id: str | None = None,
organization_id: str | None = None,
) -> BrowserState: ) -> BrowserState:
pw = await async_playwright().start() pw = await async_playwright().start()
( (
@@ -42,6 +43,7 @@ class BrowserManager:
url=url, url=url,
task_id=task_id, task_id=task_id,
workflow_run_id=workflow_run_id, workflow_run_id=workflow_run_id,
organization_id=organization_id,
) )
return BrowserState( return BrowserState(
pw=pw, pw=pw,
@@ -64,11 +66,18 @@ class BrowserManager:
return self.pages[task.task_id] return self.pages[task.task_id]
LOG.info("Creating browser state for task", task_id=task.task_id) LOG.info("Creating browser state for task", task_id=task.task_id)
browser_state = await self._create_browser_state(task.proxy_location, task.url, task.task_id) browser_state = await self._create_browser_state(
proxy_location=task.proxy_location,
url=task.url,
task_id=task.task_id,
organization_id=task.organization_id,
)
# The URL here is only used when creating a new page, and not when using an existing page. # The URL here is only used when creating a new page, and not when using an existing page.
# This will make sure browser_state.page is not None. # This will make sure browser_state.page is not None.
await browser_state.get_or_create_page(url=task.url, proxy_location=task.proxy_location, task_id=task.task_id) await browser_state.get_or_create_page(
url=task.url, proxy_location=task.proxy_location, task_id=task.task_id, organization_id=task.organization_id
)
self.pages[task.task_id] = browser_state self.pages[task.task_id] = browser_state
if task.workflow_run_id: if task.workflow_run_id:
@@ -83,13 +92,19 @@ class BrowserManager:
workflow_run_id=workflow_run.workflow_run_id, workflow_run_id=workflow_run.workflow_run_id,
) )
browser_state = await self._create_browser_state( browser_state = await self._create_browser_state(
workflow_run.proxy_location, url=url, workflow_run_id=workflow_run.workflow_run_id workflow_run.proxy_location,
url=url,
workflow_run_id=workflow_run.workflow_run_id,
organization_id=workflow_run.organization_id,
) )
# The URL here is only used when creating a new page, and not when using an existing page. # The URL here is only used when creating a new page, and not when using an existing page.
# This will make sure browser_state.page is not None. # This will make sure browser_state.page is not None.
await browser_state.get_or_create_page( await browser_state.get_or_create_page(
url=url, proxy_location=workflow_run.proxy_location, workflow_run_id=workflow_run.workflow_run_id url=url,
proxy_location=workflow_run.proxy_location,
workflow_run_id=workflow_run.workflow_run_id,
organization_id=workflow_run.organization_id,
) )
self.pages[workflow_run.workflow_run_id] = browser_state self.pages[workflow_run.workflow_run_id] = browser_state