diff --git a/origami/cli.py b/origami/cli.py index b962f216..1b91a758 100644 --- a/origami/cli.py +++ b/origami/cli.py @@ -12,10 +12,10 @@ async def _get_notebook(file_id: str, api_url: str = "https://app.noteable.io/gate/api"): - if not os.environ['NOTEABLE_TOKEN']: - raise RuntimeError('NOTEABLE_TOKEN environment variable not set') + if not os.environ["NOTEABLE_TOKEN"]: + raise RuntimeError("NOTEABLE_TOKEN environment variable not set") api_client = APIClient( - authorization_token=os.environ['NOTEABLE_TOKEN'], + authorization_token=os.environ["NOTEABLE_TOKEN"], api_base_url=api_url, ) rtu_client: RTUClient = await api_client.connect_realtime(file=file_id) @@ -28,12 +28,12 @@ def fetch(file_id: str, api_url: str = "https://app.noteable.io/gate/api"): async def _tail_notebook(file_id: str, api_url: str = "https://app.noteable.io/gate/api"): - if not os.environ['NOTEABLE_TOKEN']: - raise RuntimeError('NOTEABLE_TOKEN environment variable not set') + if not os.environ["NOTEABLE_TOKEN"]: + raise RuntimeError("NOTEABLE_TOKEN environment variable not set") setup_logging() - logging.getLogger('origami.clients.rtu').setLevel(logging.DEBUG) + logging.getLogger("origami.clients.rtu").setLevel(logging.DEBUG) api_client = APIClient( - authorization_token=os.environ['NOTEABLE_TOKEN'], + authorization_token=os.environ["NOTEABLE_TOKEN"], api_base_url=api_url, ) print("RTU Client starting initialization") diff --git a/origami/clients/api.py b/origami/clients/api.py index c305fea1..0ba22465 100644 --- a/origami/clients/api.py +++ b/origami/clients/api.py @@ -28,15 +28,15 @@ def __init__( headers: Optional[dict] = None, transport: Optional[httpx.AsyncHTTPTransport] = None, timeout: httpx.Timeout = httpx.Timeout(5.0), - rtu_client_type: str = 'origami', + rtu_client_type: str = "origami", ): # jwt and api_base_url saved as attributes because they're re-used when creating rtu client - self.jwt = authorization_token or os.environ.get('NOTEABLE_TOKEN') + self.jwt = authorization_token or os.environ.get("NOTEABLE_TOKEN") if not self.jwt: raise ValueError( - 'Must provide authorization_token or set NOTEABLE_TOKEN environment variable' + "Must provide authorization_token or set NOTEABLE_TOKEN environment variable" ) - self.api_base_url = os.environ.get('NOTEABLE_API_URL', api_base_url) + self.api_base_url = os.environ.get("NOTEABLE_API_URL", api_base_url) self.headers = {"Authorization": f"Bearer {self.jwt}"} if headers: self.headers.update(headers) @@ -48,8 +48,8 @@ def __init__( timeout=timeout, ) # Hack until Gate changes out rtu_client_type from enum to str - if rtu_client_type not in ['origami', 'origamist', 'planar_ally', 'geas']: - rtu_client_type = 'unknown' + if rtu_client_type not in ["origami", "origamist", "planar_ally", "geas"]: + rtu_client_type = "unknown" self.rtu_client_type = rtu_client_type # Only used when generating an RTUClient def add_tags_and_contextvars(self, **tags): @@ -108,7 +108,7 @@ async def create_project( self.add_tags_and_contextvars(space_id=str(space_id)) endpoint = "/projects" resp = await self.client.post( - endpoint, json={'space_id': str(space_id), "name": name, "description": description} + endpoint, json={"space_id": str(space_id), "name": name, "description": description} ) resp.raise_for_status() project = Project.parse_obj(resp.json()) @@ -134,7 +134,7 @@ async def delete_project(self, project_id: uuid.UUID) -> Project: async def list_project_files(self, project_id: uuid.UUID) -> List[File]: """List all Files in a Project. Files do not have presigned download urls included here.""" self.add_tags_and_contextvars(project_id=str(project_id)) - endpoint = f'/projects/{project_id}/files' + endpoint = f"/projects/{project_id}/files" resp = await self.client.get(endpoint) resp.raise_for_status() files = [File.parse_obj(file) for file in resp.json()] @@ -145,7 +145,7 @@ async def _multi_step_file_create( self, project_id: uuid.UUID, path: str, - file_type: Literal['file', 'notebook'], + file_type: Literal["file", "notebook"], content: bytes, ) -> File: # Uploading files using the /v1/files endpoint is a multi-step process. @@ -158,7 +158,7 @@ async def _multi_step_file_create( "project_id": str(project_id), "path": path, "type": file_type, - 'file_size_bytes': len(content), + "file_size_bytes": len(content), } resp = await self.client.post("/v1/files", json=body) resp.raise_for_status() @@ -173,8 +173,8 @@ async def _multi_step_file_create( # (2) Upload to pre-signed url # TODO: remove this hack if/when we get containers in Skaffold to be able to translate # localhost urls to the minio pod/container - if 'LOCAL_K8S' in os.environ and bool(os.environ['LOCAL_K8S']): - upload_url = upload_url.replace('localhost', 'minio') + if "LOCAL_K8S" in os.environ and bool(os.environ["LOCAL_K8S"]): + upload_url = upload_url.replace("localhost", "minio") async with httpx.AsyncClient() as plain_client: r = await plain_client.put(upload_url, content=content) r.raise_for_status() @@ -215,7 +215,7 @@ async def create_notebook( async def get_file(self, file_id: uuid.UUID) -> File: """Get metadata about a File, not including its content. Includes presigned download url.""" self.add_tags_and_contextvars(file_id=str(file_id)) - endpoint = f'/v1/files/{file_id}' + endpoint = f"/v1/files/{file_id}" resp = await self.client.get(endpoint) resp.raise_for_status() file = File.parse_obj(resp.json()) @@ -230,8 +230,8 @@ async def get_file_content(self, file_id: uuid.UUID) -> bytes: raise ValueError(f"File {file.id} does not have a presigned download url") # TODO: remove this hack if/when we get containers in Skaffold to be able to translate # localhost urls to the minio pod/container - if 'LOCAL_K8S' in os.environ and bool(os.environ['LOCAL_K8S']): - presigned_download_url = presigned_download_url.replace('localhost', 'minio') + if "LOCAL_K8S" in os.environ and bool(os.environ["LOCAL_K8S"]): + presigned_download_url = presigned_download_url.replace("localhost", "minio") async with httpx.AsyncClient() as plain_http_client: resp = await plain_http_client.get(presigned_download_url) resp.raise_for_status() @@ -243,7 +243,7 @@ async def get_file_versions(self, file_id: uuid.UUID) -> List[FileVersion]: of any previous version. Note when working with older versions, you do not want to establish an RTUClient to "catch up" past that version. """ - endpoint = f'/files/{file_id}/versions' + endpoint = f"/files/{file_id}/versions" resp = await self.client.get(endpoint) resp.raise_for_status() versions = [FileVersion.parse_obj(version) for version in resp.json()] @@ -251,7 +251,7 @@ async def get_file_versions(self, file_id: uuid.UUID) -> List[FileVersion]: async def delete_file(self, file_id: uuid.UUID) -> File: self.add_tags_and_contextvars(file_id=str(file_id)) - endpoint = f'/v1/files/{file_id}' + endpoint = f"/v1/files/{file_id}" resp = await self.client.delete(endpoint) resp.raise_for_status() file = File.parse_obj(resp.json()) @@ -268,14 +268,14 @@ async def get_datasources_for_notebook(self, file_id: uuid.UUID) -> List[DataSou return datasources async def launch_kernel( - self, file_id: uuid.UUID, kernel_name: str = 'python3', hardware_size: str = 'small' + self, file_id: uuid.UUID, kernel_name: str = "python3", hardware_size: str = "small" ) -> KernelSession: - endpoint = '/v1/sessions' + endpoint = "/v1/sessions" data = { - 'file_id': str(file_id), - 'kernel_config': { - 'kernel_name': kernel_name, - 'hardware_size_identifier': hardware_size, + "file_id": str(file_id), + "kernel_config": { + "kernel_name": kernel_name, + "hardware_size_identifier": hardware_size, }, } resp = await self.client.post(endpoint, json=data) @@ -289,7 +289,7 @@ async def launch_kernel( return kernel_session async def shutdown_kernel(self, kernel_session_id: uuid.UUID) -> None: - endpoint = f'/sessions/{kernel_session_id}' + endpoint = f"/sessions/{kernel_session_id}" resp = await self.client.delete(endpoint, timeout=60) resp.raise_for_status() logger.info("Shut down kernel", extra={"kernel_session_id": str(kernel_session_id)}) @@ -297,7 +297,7 @@ async def shutdown_kernel(self, kernel_session_id: uuid.UUID) -> None: async def get_output_collection( self, output_collection_id: uuid.UUID ) -> KernelOutputCollection: - endpoint = f'/outputs/collection/{output_collection_id}' + endpoint = f"/outputs/collection/{output_collection_id}" resp = await self.client.get(endpoint) resp.raise_for_status() return KernelOutputCollection.parse_obj(resp.json()) @@ -325,21 +325,21 @@ async def connect_realtime(self, file: Union[File, uuid.UUID, str]) -> RTUClient logger.info(f"Creating RTUClient for file {file_id}") file = await self.get_file(file_id) - if file.type != 'notebook': + if file.type != "notebook": raise ValueError(f"File {file_id} is not a notebook") if not file.presigned_download_url: raise ValueError(f"File {file_id} does not have a presigned download url") # TODO: remove this hack if/when we get containers in Skaffold to be able to translate # localhost urls to the minio pod/container - if 'LOCAL_K8S' in os.environ and bool(os.environ['LOCAL_K8S']): - file.presigned_download_url = file.presigned_download_url.replace('localhost', 'minio') + if "LOCAL_K8S" in os.environ and bool(os.environ["LOCAL_K8S"]): + file.presigned_download_url = file.presigned_download_url.replace("localhost", "minio") async with httpx.AsyncClient() as plain_http_client: resp = await plain_http_client.get(file.presigned_download_url) resp.raise_for_status() seed_notebook = Notebook.parse_obj(resp.json()) nb_builder = NotebookBuilder(seed_notebook=seed_notebook) - rtu_url = self.api_base_url.replace('http', 'ws') + '/v1/rtu' + rtu_url = self.api_base_url.replace("http", "ws") + "/v1/rtu" rtu_client = RTUClient( rtu_url=rtu_url, jwt=self.jwt, diff --git a/origami/clients/rtu.py b/origami/clients/rtu.py index ab0596b4..9309aba7 100644 --- a/origami/clients/rtu.py +++ b/origami/clients/rtu.py @@ -79,7 +79,7 @@ async def inbound_message_hook(self, contents: str) -> RTUResponse: # then a second parse to go through the discriminators to a specific event (or fall back # to error or BaseRTUResponse) data: dict = orjson.loads(contents) - data['channel_prefix'] = data.get('channel', '').split('/')[0] + data["channel_prefix"] = data.get("channel", "").split("/")[0] rtu_event = parse_obj_as(RTUResponse, data) # Debug Logging @@ -162,12 +162,12 @@ class DeltaRequestCallbackManager: # Delta is guarenteed to be in rtu_client.builder at this point """ - def __init__(self, client: 'RTUClient', delta: FileDelta): + def __init__(self, client: "RTUClient", delta: FileDelta): self.result = asyncio.Future() self.client = client self.delta = delta # keep a ref to use in self.delta_cb_ref req = NewDeltaRequest( - channel=f'files/{self.client.file_id}', data=NewDeltaRequestData(delta=delta) + channel=f"files/{self.client.file_id}", data=NewDeltaRequestData(delta=delta) ) # Register one cb by RTU request transaction id in order to catch errors and set Future self.rtu_cb_ref = client.register_transaction_id_callback( @@ -206,7 +206,7 @@ async def rtu_cb(self, msg: RTUResponse): async def delta_cb(self, delta: FileDelta): if delta.id == self.delta.id: - logger.debug("Delta squashed", extra={'delta': delta}) + logger.debug("Delta squashed", extra={"delta": delta}) if not self.result.done(): self.result.set_result(delta) self.deregister_callbacks() @@ -225,7 +225,7 @@ def __init__( file_id: uuid.UUID, file_version_id: uuid.UUID, builder: NotebookBuilder, - rtu_client_type: str = 'origami', + rtu_client_type: str = "origami", file_subscribe_timeout: int = 10, ): """ @@ -290,7 +290,7 @@ def __init__( self.register_rtu_event_callback(rtu_event=NewDeltaEvent, fn=self._on_delta_recv) # Kernel and cell state handling - self.kernel_state: str = 'not_started' # value used when there's no Kernel for a Notebook + self.kernel_state: str = "not_started" # value used when there's no Kernel for a Notebook self.cell_states: Dict[str, str] = {} self.register_rtu_event_callback( @@ -319,7 +319,7 @@ def cell_ids(self): @property def kernel_pod_name(self) -> str: """Transform the file_id into the Pod name used to build the kernels/ RTU channel""" - return f'kernels/notebook-kernel-{self.file_id.hex[:20]}' + return f"kernels/notebook-kernel-{self.file_id.hex[:20]}" def send(self, msg: RTURequest): """ @@ -423,7 +423,7 @@ async def auth_hook(self, *args, **kwargs): until we've observed an `authenticate_reply` event """ auth_request = AuthenticateRequest( - data={'token': self.jwt, 'rtu_client_type': self.rtu_client_type} + data={"token": self.jwt, "rtu_client_type": self.rtu_client_type} ) # auth_hook is the special situation that shouldn't use manager.send(), @@ -478,21 +478,21 @@ async def send_file_subscribe(self): # Second note, subscribing by delta id all-0's throws an error in Gate. if self.builder.last_applied_delta_id and self.builder.last_applied_delta_id != uuid.UUID(int=0): # type: ignore # noqa: E501 req = FileSubscribeRequest( - channel=f'files/{self.file_id}', - data={'from_delta_id': self.builder.last_applied_delta_id}, + channel=f"files/{self.file_id}", + data={"from_delta_id": self.builder.last_applied_delta_id}, ) logger.info( "Sending File subscribe request by last applied delta id", - extra={'from_delta_id': str(req.data.from_delta_id)}, + extra={"from_delta_id": str(req.data.from_delta_id)}, ) else: req = FileSubscribeRequest( - channel=f'files/{self.file_id}', - data={'from_version_id': self.file_version_id}, + channel=f"files/{self.file_id}", + data={"from_version_id": self.file_version_id}, ) logger.info( "Sending File subscribe request by version id", - extra={'from_version_id': str(req.data.from_version_id)}, + extra={"from_version_id": str(req.data.from_version_id)}, ) self.file_subscribe_timeout_task = asyncio.create_task(self.on_file_subscribe_timeout()) self.manager.send(req) @@ -651,10 +651,10 @@ async def apply_delta(self, delta: FileDelta): "Error trying to run callback while applying delta", exc_info="".join(traceback.format_tb(result.__traceback__)), extra={ - 'callback': callback, - 'delta': delta, - 'ename': repr(result), - 'traceback': "".join(traceback.format_tb(result.__traceback__)), + "callback": callback, + "delta": delta, + "ename": repr(result), + "traceback": "".join(traceback.format_tb(result.__traceback__)), }, ) @@ -675,7 +675,7 @@ async def replay_unapplied_deltas(self): if delta.parent_delta_id == self.builder.last_applied_delta_id: logger.debug( "Applying previously queued out of order delta", - extra={'delta_id': str(delta.id)}, + extra={"delta_id": str(delta.id)}, ) await self.apply_delta(delta=delta) self.unapplied_deltas.remove(delta) @@ -694,12 +694,12 @@ async def on_bulk_cell_state_update(self, msg: BulkCellStateUpdateResponse): if item.cell_id in self._execute_cell_events: # When we see that a cell we're monitoring has finished, resolve the Future to # be the cell - if item.state in ['finished_with_error', 'finished_with_no_error']: + if item.state in ["finished_with_error", "finished_with_no_error"]: logger.debug( "Cell execution for monitored cell finished", extra={ - 'cell_id': item.cell_id, - 'state': item.state, + "cell_id": item.cell_id, + "state": item.state, }, ) fut = self._execute_cell_events[item.cell_id] @@ -712,18 +712,18 @@ async def on_bulk_cell_state_update(self, msg: BulkCellStateUpdateResponse): logger.warning( "Cell execution finished for cell that doesn't exist in Notebook", extra={ - 'cell_id': item.cell_id, - 'state': item.state, + "cell_id": item.cell_id, + "state": item.state, }, ) fut.set_exception(CellNotFound(item.cell_id)) self.cell_states[item.cell_id] = item.state - logger.debug("Updated cell states", extra={'cell_states': self.cell_states}) + logger.debug("Updated cell states", extra={"cell_states": self.cell_states}) async def wait_for_kernel_idle(self): """Wait for the kernel to be idle""" logger.debug("Waiting for Kernel to be idle") - while self.kernel_state != 'idle': + while self.kernel_state != "idle": await asyncio.sleep(0.05) logger.debug("Kernel is idle") @@ -738,7 +738,7 @@ async def new_delta_request(self, delta=FileDelta) -> FileDelta: async def add_cell( self, - source: str = '', + source: str = "", cell: Optional[NotebookCell] = None, before_id: Optional[str] = None, after_id: Optional[str] = None, @@ -762,15 +762,15 @@ async def add_cell( return cell async def delete_cell(self, cell_id: str) -> NBCellsDelete: - delta = NBCellsDelete(file_id=self.file_id, properties={'id': cell_id}) + delta = NBCellsDelete(file_id=self.file_id, properties={"id": cell_id}) return await self.new_delta_request(delta) async def change_cell_type( self, cell_id: str, - cell_type: Literal['code', 'markdown', 'sql'], - code_language: str = 'python', - db_connection: str = '@noteable', + cell_type: Literal["code", "markdown", "sql"], + code_language: str = "python", + db_connection: str = "@noteable", assign_results_to: Optional[str] = None, ) -> NotebookCell: """ @@ -779,40 +779,40 @@ async def change_cell_type( - db_connection and assign_results_to only relevant when switching to SQL cell """ self.builder.get_cell(cell_id) # Raise CellNotFound if it doesn't exist - if cell_type == 'code': + if cell_type == "code": delta = CellMetadataReplace( file_id=self.file_id, resource_id=cell_id, - properties={'language': code_language, 'type': 'code'}, + properties={"language": code_language, "type": "code"}, ) await self.new_delta_request(delta) - elif cell_type == 'markdown': + elif cell_type == "markdown": delta = CellMetadataReplace( file_id=self.file_id, resource_id=cell_id, - properties={'language': 'markdown', 'type': 'markdown'}, + properties={"language": "markdown", "type": "markdown"}, ) await self.new_delta_request(delta) - elif cell_type == 'sql': + elif cell_type == "sql": delta = CellMetadataReplace( file_id=self.file_id, resource_id=cell_id, - properties={'language': 'sql', 'type': 'code'}, + properties={"language": "sql", "type": "code"}, ) await self.new_delta_request(delta) if not assign_results_to: name_suffix = "".join(random.choices(string.ascii_lowercase, k=4)) - assign_results_to = 'df_' + name_suffix + assign_results_to = "df_" + name_suffix delta = CellMetadataUpdate( file_id=self.file_id, resource_id=cell_id, properties={ - 'path': ['metadata', 'noteable'], - 'value': { - 'cell_type': 'sql', - 'db_connection': db_connection, - 'assign_results_to': assign_results_to, + "path": ["metadata", "noteable"], + "value": { + "cell_type": "sql", + "db_connection": db_connection, + "assign_results_to": assign_results_to, }, }, ) @@ -828,7 +828,7 @@ async def update_cell_content(self, cell_id: str, patch: str) -> NotebookCell: Update cell content with a diff-match-patch patch string """ delta = CellContentsUpdate( - file_id=self.file_id, resource_id=cell_id, properties={'patch': patch} + file_id=self.file_id, resource_id=cell_id, properties={"patch": patch} ) await self.new_delta_request(delta) # Grab updated cell post-squashing @@ -840,7 +840,7 @@ async def replace_cell_content(self, cell_id: str, source: str) -> NotebookCell: Replace cell content with a string """ delta = CellContentsReplace( - file_id=self.file_id, resource_id=cell_id, properties={'source': source} + file_id=self.file_id, resource_id=cell_id, properties={"source": source} ) await self.new_delta_request(delta) # Grab updated cell post-squashing @@ -896,7 +896,7 @@ async def queue_execution( # will never get executed by PA/Kernel, so we'd never see cell status and resolve future future = asyncio.Future() idx, cell = self.builder.get_cell(cell_id) - if cell.cell_type == 'code' and cell.source.strip(): + if cell.cell_type == "code" and cell.source.strip(): self._execute_cell_events[cell_id] = future futures[future] = cell_id await self.new_delta_request(delta) diff --git a/origami/models/api/files.py b/origami/models/api/files.py index 03b92ac6..8714bc36 100644 --- a/origami/models/api/files.py +++ b/origami/models/api/files.py @@ -15,7 +15,7 @@ class File(ResourceBase): space_id: uuid.UUID size: Optional[int] = None mimetype: Optional[str] = None - type: Literal['file', 'notebook'] + type: Literal["file", "notebook"] current_version_id: Optional[uuid.UUID] = None # presigned_download_url is None when listing Files in a Project, need to hit /api/v1/files/{id} # to get it. Use presigned download url to get File content including Notebooks @@ -24,7 +24,7 @@ class File(ResourceBase): @validator("url", always=True) def construct_url(cls, v, values): - noteable_url = os.environ.get('PUBLIC_NOTEABLE_URL', 'https://app.noteable.io') + noteable_url = os.environ.get("PUBLIC_NOTEABLE_URL", "https://app.noteable.io") return f"{noteable_url}/f/{values['id']}/{values['path']}" diff --git a/origami/models/api/projects.py b/origami/models/api/projects.py index 06ba8a1a..38788bf9 100644 --- a/origami/models/api/projects.py +++ b/origami/models/api/projects.py @@ -15,5 +15,5 @@ class Project(ResourceBase): @validator("url", always=True) def construct_url(cls, v, values): - noteable_url = os.environ.get('PUBLIC_NOTEABLE_URL', 'https://app.noteable.io') + noteable_url = os.environ.get("PUBLIC_NOTEABLE_URL", "https://app.noteable.io") return f"{noteable_url}/p/{values['id']}" diff --git a/origami/models/api/spaces.py b/origami/models/api/spaces.py index 504baf12..c7e3b07d 100644 --- a/origami/models/api/spaces.py +++ b/origami/models/api/spaces.py @@ -13,5 +13,5 @@ class Space(ResourceBase): @validator("url", always=True) def construct_url(cls, v, values): - noteable_url = os.environ.get('PUBLIC_NOTEABLE_URL', 'https://app.noteable.io') + noteable_url = os.environ.get("PUBLIC_NOTEABLE_URL", "https://app.noteable.io") return f"{noteable_url}/s/{values['id']}" diff --git a/origami/models/api/users.py b/origami/models/api/users.py index 9e505bcb..56e99c05 100644 --- a/origami/models/api/users.py +++ b/origami/models/api/users.py @@ -19,5 +19,5 @@ class User(ResourceBase): @validator("auth_type", always=True) def construct_auth_type(cls, v, values): - if values.get('principal_sub'): + if values.get("principal_sub"): return values["principal_sub"].split("|")[0] diff --git a/origami/models/deltas/delta_types/cell_contents.py b/origami/models/deltas/delta_types/cell_contents.py index fbd9d195..9d0785fc 100644 --- a/origami/models/deltas/delta_types/cell_contents.py +++ b/origami/models/deltas/delta_types/cell_contents.py @@ -15,7 +15,7 @@ class CellContentsUpdateProperties(BaseModel): class CellContentsUpdate(CellContentsDelta): # resource_id should be cell id to update - delta_action: Literal['update'] = 'update' + delta_action: Literal["update"] = "update" properties: CellContentsUpdateProperties @@ -25,7 +25,7 @@ class CellContentsReplaceProperties(BaseModel): class CellContentsReplace(CellContentsDelta): # resource_id should be cell id to replace - delta_action: Literal['replace'] = 'replace' + delta_action: Literal["replace"] = "replace" properties: CellContentsReplaceProperties diff --git a/origami/models/deltas/delta_types/cell_execute.py b/origami/models/deltas/delta_types/cell_execute.py index 47ed3cf1..bf4be78f 100644 --- a/origami/models/deltas/delta_types/cell_execute.py +++ b/origami/models/deltas/delta_types/cell_execute.py @@ -12,24 +12,24 @@ class CellExecuteDelta(FileDeltaBase): class CellExecute(CellExecuteDelta): # execute single cel # resource_id should be cell id to run - delta_action: Literal['execute'] = 'execute' + delta_action: Literal["execute"] = "execute" class CellExecuteAfter(CellExecuteDelta): # execute specific cell id and all cells after it # resource_id should be cell id to run - delta_action: Literal['execute_after'] = 'execute_after' + delta_action: Literal["execute_after"] = "execute_after" class CellExecuteBefore(CellExecuteDelta): # execute all cells up to specific cell, inclusive of that cell id # resource_id should be cell id to run - delta_action: Literal['execute_before'] = 'execute_before' + delta_action: Literal["execute_before"] = "execute_before" class CellExecuteAll(CellExecuteDelta): # execute all cells - delta_action: Literal['execute_all'] = 'execute_all' + delta_action: Literal["execute_all"] = "execute_all" CellExecuteDeltas = Annotated[ diff --git a/origami/models/deltas/delta_types/cell_metadata.py b/origami/models/deltas/delta_types/cell_metadata.py index c8d09a77..ee17fab3 100644 --- a/origami/models/deltas/delta_types/cell_metadata.py +++ b/origami/models/deltas/delta_types/cell_metadata.py @@ -20,7 +20,7 @@ class CellMetadataUpdateProperties(BaseModel): class CellMetadataUpdate(CellMetadataDelta): # resource_id should be cell id to update - delta_action: Literal['update'] = 'update' + delta_action: Literal["update"] = "update" properties: CellMetadataUpdateProperties @@ -32,7 +32,7 @@ class CellMetadataReplaceProperties(BaseModel): class CellMetadataReplace(CellMetadataDelta): # resource_id should be cell id to replace - delta_action: Literal['replace'] = 'replace' + delta_action: Literal["replace"] = "replace" properties: CellMetadataReplaceProperties diff --git a/origami/models/deltas/delta_types/cell_output_collection.py b/origami/models/deltas/delta_types/cell_output_collection.py index 861425e8..749c7d08 100644 --- a/origami/models/deltas/delta_types/cell_output_collection.py +++ b/origami/models/deltas/delta_types/cell_output_collection.py @@ -16,7 +16,7 @@ class CellOutputCollectionReplaceData(BaseModel): class CellOutputCollectionReplace(CellOutputCollectionDelta): # resource_id should be cell id to replace with new output ocllection id - delta_action: Literal['replace'] = 'replace' + delta_action: Literal["replace"] = "replace" properties: CellOutputCollectionReplaceData diff --git a/origami/models/deltas/delta_types/nb_cells.py b/origami/models/deltas/delta_types/nb_cells.py index 1e9b87ec..9f77b495 100644 --- a/origami/models/deltas/delta_types/nb_cells.py +++ b/origami/models/deltas/delta_types/nb_cells.py @@ -17,7 +17,7 @@ class NBCellsAddProperties(BaseModel): class NBCellsAdd(NBCellsDelta): - delta_action: Literal['add'] = 'add' + delta_action: Literal["add"] = "add" properties: NBCellsAddProperties @@ -26,7 +26,7 @@ class NBCellsDeleteProperties(BaseModel): class NBCellsDelete(NBCellsDelta): - delta_action: Literal['delete'] = 'delete' + delta_action: Literal["delete"] = "delete" properties: NBCellsDeleteProperties @@ -36,7 +36,7 @@ class NBCellsMoveProperties(BaseModel): class NBCellsMove(NBCellsDelta): - delta_action: Literal['move'] = 'move' + delta_action: Literal["move"] = "move" properties: NBCellsMoveProperties diff --git a/origami/models/deltas/delta_types/nb_metadata.py b/origami/models/deltas/delta_types/nb_metadata.py index 18853e9a..716d28dd 100644 --- a/origami/models/deltas/delta_types/nb_metadata.py +++ b/origami/models/deltas/delta_types/nb_metadata.py @@ -16,7 +16,7 @@ class NBMetadataProperties(BaseModel): class NBMetadataUpdate(NBMetadataDelta): - delta_action: Literal['update'] = 'update' + delta_action: Literal["update"] = "update" properties: NBMetadataProperties diff --git a/origami/models/notebook.py b/origami/models/notebook.py index fbb9cdfb..90f2b9b3 100644 --- a/origami/models/notebook.py +++ b/origami/models/notebook.py @@ -101,27 +101,27 @@ def output_collection_id(self) -> Optional[Union[str, uuid.UUID]]: def make_sql_cell( cell_id: Optional[str] = None, - source: str = '', - db_connection: str = '@noteable', + source: str = "", + db_connection: str = "@noteable", assign_results_to: Optional[str] = None, ) -> CodeCell: cell_id = cell_id or str(uuid.uuid4()) # Remove first line of source if it starts with %%sql. That is the right syntax for regular # code cells with sql magic support, but Noteable SQL cells should have just the sql source - if source.startswith('%%sql'): + if source.startswith("%%sql"): lines = source.splitlines() - source = '\n'.join(lines[1:]) + source = "\n".join(lines[1:]) if not assign_results_to: name_suffix = "".join(random.choices(string.ascii_lowercase, k=4)) - assign_results_to = 'df_' + name_suffix + assign_results_to = "df_" + name_suffix metadata = { - 'language': 'sql', - 'type': 'code', - 'noteable': { - 'cell_type': 'sql', - 'db_connection': db_connection, - 'assign_results_to': assign_results_to, + "language": "sql", + "type": "code", + "noteable": { + "cell_type": "sql", + "db_connection": db_connection, + "assign_results_to": assign_results_to, }, } return CodeCell(cell_id=cell_id, source=source, metadata=metadata) diff --git a/origami/models/rtu/base.py b/origami/models/rtu/base.py index 40d608cb..eb378079 100644 --- a/origami/models/rtu/base.py +++ b/origami/models/rtu/base.py @@ -19,11 +19,11 @@ class BaseRTURequest(BaseModel): class Config: # do not include channel_prefix when serializing to dict / json - fields = {'channel_prefix': {'exclude': True}} + fields = {"channel_prefix": {"exclude": True}} @root_validator def set_channel_prefix(cls, values): - values['channel_prefix'] = values['channel'].split('/')[0] + values["channel_prefix"] = values["channel"].split("/")[0] return values @@ -38,9 +38,9 @@ class BaseRTUResponse(BaseModel): class Config: # do not include channel_prefix when serializing to dict / json - fields = {'channel_prefix': {'exclude': True}} + fields = {"channel_prefix": {"exclude": True}} @root_validator def set_channel_prefix(cls, values): - values['channel_prefix'] = values['channel'].split('/')[0] + values["channel_prefix"] = values["channel"].split("/")[0] return values diff --git a/origami/models/rtu/channels/files.py b/origami/models/rtu/channels/files.py index dbd7ff07..a6c9e735 100644 --- a/origami/models/rtu/channels/files.py +++ b/origami/models/rtu/channels/files.py @@ -24,11 +24,11 @@ class FilesRequest(BaseRTURequest): - channel_prefix: Literal['files'] = 'files' + channel_prefix: Literal["files"] = "files" class FilesResponse(BaseRTUResponse): - channel_prefix: Literal['files'] = 'files' + channel_prefix: Literal["files"] = "files" # When an RTU Client wants to get document model updates from a Notebook, it subscribes to the files @@ -51,11 +51,11 @@ def exactly_one_field(cls, values): return values # If not, raise a validation error - raise ValidationError('Exactly one field must be set') + raise ValidationError("Exactly one field must be set") class FileSubscribeRequest(FilesRequest): - event: Literal['subscribe_request'] = 'subscribe_request' + event: Literal["subscribe_request"] = "subscribe_request" data: FileSubscribeRequestData @@ -70,17 +70,17 @@ class FileSubscribeReplyData(BaseModel): class FileSubscribeReply(FilesResponse): - event: Literal['subscribe_reply'] = 'subscribe_reply' + event: Literal["subscribe_reply"] = "subscribe_reply" data: FileSubscribeReplyData # Clients typically do not need to unsubscribe, they can just close the websocket connection class FileUnsubscribeRequest(FilesRequest): - event: Literal['unsubscribe_request'] = 'unsubscribe_request' + event: Literal["unsubscribe_request"] = "unsubscribe_request" class FileUnsubscribeReply(FilesResponse): - event: Literal['unsubscribe_reply'] = 'unsubscribe_reply' + event: Literal["unsubscribe_reply"] = "unsubscribe_reply" data: BooleanReplyData @@ -94,17 +94,17 @@ class NewDeltaRequestData(BaseModel): class NewDeltaRequest(FilesRequest): - event: Literal['new_delta_request'] = 'new_delta_request' + event: Literal["new_delta_request"] = "new_delta_request" data: NewDeltaRequestData class NewDeltaReply(FilesResponse): - event: Literal['new_delta_reply'] = 'new_delta_reply' + event: Literal["new_delta_reply"] = "new_delta_reply" data: BooleanReplyData class NewDeltaEvent(FilesResponse): - event: Literal['new_delta_event'] = 'new_delta_event' + event: Literal["new_delta_event"] = "new_delta_event" data: FileDelta @@ -116,14 +116,14 @@ class UpdateOutputCollectionEventData(BaseModel): class UpdateOutputCollectionEvent(FilesResponse): - event: Literal['update_output_collection_event'] = 'update_output_collection_event' + event: Literal["update_output_collection_event"] = "update_output_collection_event" data: UpdateOutputCollectionEventData # If Cells are streaming multiple outputs like a pip install or for loop and print, then we'll get # append to output events class AppendOutputEvent(FilesResponse): - event: Literal['append_output_event'] = 'append_output_event' + event: Literal["append_output_event"] = "append_output_event" data: KernelOutput @@ -134,12 +134,12 @@ class UpdateUserCellSelectionRequestData(BaseModel): class UpdateUserCellSelectionRequest(FilesRequest): - event: Literal['update_user_cell_selection_request'] = 'update_user_cell_selection_request' + event: Literal["update_user_cell_selection_request"] = "update_user_cell_selection_request" data: UpdateUserCellSelectionRequestData class UpdateUserCellSelectionReply(FilesResponse): - event: Literal['update_user_cell_selection_reply'] = 'update_user_cell_selection_reply' + event: Literal["update_user_cell_selection_reply"] = "update_user_cell_selection_reply" data: BooleanReplyData @@ -152,7 +152,7 @@ class UpdateUserFileSubscriptionEventData(BaseModel): class UpdateUserFileSubscriptionEvent(FilesResponse): - event: Literal['update_user_file_subscription_event'] = 'update_user_file_subscription_event' + event: Literal["update_user_file_subscription_event"] = "update_user_file_subscription_event" data: UpdateUserFileSubscriptionEventData @@ -161,7 +161,7 @@ class RemoveUserFileSubscriptionEventData(BaseModel): class RemoveUserFileSubscriptionEvent(FilesResponse): - event: Literal['remove_user_file_subscription_event'] = 'remove_user_file_subscription_event' + event: Literal["remove_user_file_subscription_event"] = "remove_user_file_subscription_event" data: RemoveUserFileSubscriptionEventData @@ -172,7 +172,7 @@ class UsageMetricsEventData(BaseModel): class UsageMetricsEvent(FilesResponse): - event: Literal['usage_metrics_event'] = 'usage_metrics_event' + event: Literal["usage_metrics_event"] = "usage_metrics_event" data: UsageMetricsEventData @@ -184,30 +184,30 @@ class TransformViewToCodeRequestData(BaseModel): filters: Any ignore_index: bool = True overrides: dict = Field(default_factory=dict) - target_cell_type: str = 'code' - variable_name: str = 'df' + target_cell_type: str = "code" + variable_name: str = "df" class TransformViewToCodeRequest(FilesRequest): - event: Literal['transform_view_to_code_request'] = 'transform_view_to_code_request' + event: Literal["transform_view_to_code_request"] = "transform_view_to_code_request" data: TransformViewToCodeRequestData class TransformViewToCodeReply(FilesResponse): - event: Literal['transform_view_to_code_reply'] = 'transform_view_to_code_reply' + event: Literal["transform_view_to_code_reply"] = "transform_view_to_code_reply" data: BooleanReplyData # Widgets, ugh. Not attempting to model the payload, no current plan on doing anything with them # on the Origami side. class V0CreateWidgetModelEvent(FilesResponse): - event: Literal['v0_create_widget_model_event'] = 'v0_create_widget_model_event' + event: Literal["v0_create_widget_model_event"] = "v0_create_widget_model_event" data: Any # When the API squashes Deltas, it will emit a new file versions changed event class FileVersionsChangedEvent(FilesResponse): - event: Literal['v0_file_versions_changed_event'] = 'v0_file_versions_changed_event' + event: Literal["v0_file_versions_changed_event"] = "v0_file_versions_changed_event" data: Optional[dict] @@ -219,7 +219,7 @@ class FileVersionsChangedEvent(FilesResponse): UpdateUserCellSelectionRequest, TransformViewToCodeRequest, ], - Field(discriminator='event'), + Field(discriminator="event"), ] FileResponses = Annotated[ @@ -238,5 +238,5 @@ class FileVersionsChangedEvent(FilesResponse): AppendOutputEvent, UsageMetricsEvent, ], - Field(discriminator='event'), + Field(discriminator="event"), ] diff --git a/origami/models/rtu/channels/kernels.py b/origami/models/rtu/channels/kernels.py index 41b49d9a..10761f96 100644 --- a/origami/models/rtu/channels/kernels.py +++ b/origami/models/rtu/channels/kernels.py @@ -13,11 +13,11 @@ class KernelsRequest(BaseRTURequest): - channel_prefix: Literal['kernels'] = 'kernels' + channel_prefix: Literal["kernels"] = "kernels" class KernelsResponse(BaseRTUResponse): - channel_prefix: Literal['kernels'] = 'kernels' + channel_prefix: Literal["kernels"] = "kernels" class KernelSubscribeRequestData(BaseModel): @@ -25,7 +25,7 @@ class KernelSubscribeRequestData(BaseModel): class KernelSubscribeRequest(KernelsRequest): - event: Literal['subscribe_request'] = 'subscribe_request' + event: Literal["subscribe_request"] = "subscribe_request" data: KernelSubscribeRequestData @@ -36,12 +36,12 @@ class KernelSubscribeReplyData(BaseModel): class KernelSubscribeReply(KernelsResponse): - event: Literal['subscribe_reply'] = 'subscribe_reply' + event: Literal["subscribe_reply"] = "subscribe_reply" data: KernelSubscribeReplyData class KernelStatusUpdateResponse(KernelsResponse): - event: Literal['kernel_status_update_event'] = 'kernel_status_update_event' + event: Literal["kernel_status_update_event"] = "kernel_status_update_event" data: KernelStatusUpdate @@ -51,7 +51,7 @@ class BulkCellStateUpdateData(BaseModel): class BulkCellStateUpdateResponse(KernelsResponse): - event: Literal['bulk_cell_state_update_event'] = 'bulk_cell_state_update_event' + event: Literal["bulk_cell_state_update_event"] = "bulk_cell_state_update_event" data: BulkCellStateUpdateData @@ -59,19 +59,19 @@ class BulkCellStateUpdateResponse(KernelsResponse): # On connect to a new Kernel, Clients can send a request to trigger an event. Otherwise events occur # after cell execution automatically. class VariableExplorerUpdateRequest(KernelsRequest): - event: Literal['variable_explorer_update_request'] = 'variable_explorer_update_request' + event: Literal["variable_explorer_update_request"] = "variable_explorer_update_request" # It is confusing but variable_explorer_update_request can either be an RTU client to Gate server # (RTURequest) or also be propogated out by Gate from another client, meaning it comes in as a # server-to-client (RTUResponse) so we need to model it just to avoid warning about unmodeled msgs class VariableExplorerUpdateRequestPropogated(KernelsResponse): - event: Literal['variable_explorer_update_request'] = 'variable_explorer_update_request' + event: Literal["variable_explorer_update_request"] = "variable_explorer_update_request" data: dict = Field(default_factory=dict) class VariableExplorerResponse(KernelsResponse): - event: Literal['variable_explorer_event'] = 'variable_explorer_event' + event: Literal["variable_explorer_event"] = "variable_explorer_event" class IntegratedAIRequestData(BaseModel): @@ -85,17 +85,17 @@ class IntegratedAIRequestData(BaseModel): class IntegratedAIRequest(KernelsRequest): - event: Literal['integrated_ai_request'] = 'integrated_ai_request' + event: Literal["integrated_ai_request"] = "integrated_ai_request" data: IntegratedAIRequestData class IntegratedAIReply(KernelsResponse): - event: Literal['integrated_ai_reply'] = 'integrated_ai_reply' + event: Literal["integrated_ai_reply"] = "integrated_ai_reply" data: BooleanReplyData class IntegratedAIEvent(KernelsResponse): - event: Literal['integrated_ai_event'] = 'integrated_ai_event' + event: Literal["integrated_ai_event"] = "integrated_ai_event" # same data as the IntegratedAIRequest, just echoed back out data: IntegratedAIRequestData @@ -110,17 +110,17 @@ class IntegratedAIResultData(BaseModel): # this is sidecar to gate as a result of calling the OpenAIHandler method (OpenAI response, # error, etc); after that, Gate propogates the data out as an IntegratedAIEvent class IntegratedAIResult(KernelsRequest): - event: Literal['integrated_ai_result'] = 'integrated_ai_result' + event: Literal["integrated_ai_result"] = "integrated_ai_result" data: IntegratedAIResultData class IntegratedAIResultReply(KernelsResponse): - event: Literal['integrated_ai_result_reply'] = 'integrated_ai_result_reply' + event: Literal["integrated_ai_result_reply"] = "integrated_ai_result_reply" data: BooleanReplyData class IntegratedAIResultEvent(KernelsResponse): - event: Literal['integrated_ai_result_event'] = 'integrated_ai_result_event' + event: Literal["integrated_ai_result_event"] = "integrated_ai_result_event" data: IntegratedAIResultData diff --git a/origami/models/rtu/channels/system.py b/origami/models/rtu/channels/system.py index 9341ba4f..6705e17b 100644 --- a/origami/models/rtu/channels/system.py +++ b/origami/models/rtu/channels/system.py @@ -19,24 +19,24 @@ class SystemRequest(BaseRTURequest): - channel: str = 'system' - channel_prefix: Literal['system'] = 'system' + channel: str = "system" + channel_prefix: Literal["system"] = "system" class SystemResponse(BaseRTUResponse): - channel: str = 'system' - channel_prefix: Literal['system'] = 'system' + channel: str = "system" + channel_prefix: Literal["system"] = "system" # The first thing RTU Clients should do after websocket connection is authenticate with a JWT, # same access token as what is included in Authorization bearer headers for API requests class AuthenticateRequestData(BaseModel): token: str - rtu_client_type: str = 'origami' + rtu_client_type: str = "origami" class AuthenticateRequest(SystemRequest): - event: Literal['authenticate_request'] = 'authenticate_request' + event: Literal["authenticate_request"] = "authenticate_request" data: AuthenticateRequestData @@ -46,21 +46,21 @@ class AuthenticateReplyData(BaseModel): class AuthenticateReply(SystemResponse): - event: Literal['authenticate_reply'] = 'authenticate_reply' + event: Literal["authenticate_reply"] = "authenticate_reply" data: AuthenticateReplyData # Below is all mainly used for debug, App devs don't need to do anything with these usually class PingRequest(SystemRequest): - event: Literal['ping_request'] = 'ping_request' + event: Literal["ping_request"] = "ping_request" class PingResponse(SystemResponse): - event: Literal['ping_response'] = 'ping_response' + event: Literal["ping_response"] = "ping_response" class WhoAmIRequest(SystemRequest): - event: Literal['whoami_request'] = 'whoami_request' + event: Literal["whoami_request"] = "whoami_request" class WhoAmIResponseData(BaseModel): @@ -68,7 +68,7 @@ class WhoAmIResponseData(BaseModel): class WhoAmIResponse(SystemResponse): - event: Literal['whoami_response'] = 'whoami_response' + event: Literal["whoami_response"] = "whoami_response" data: WhoAmIResponseData @@ -78,7 +78,7 @@ class WhoAmIResponse(SystemResponse): PingRequest, WhoAmIRequest, ], - Field(discriminator='event'), + Field(discriminator="event"), ] SystemResponses = Annotated[ Union[ @@ -86,5 +86,5 @@ class WhoAmIResponse(SystemResponse): PingResponse, WhoAmIResponse, ], - Field(discriminator='event'), + Field(discriminator="event"), ] diff --git a/origami/models/rtu/discriminators.py b/origami/models/rtu/discriminators.py index 0bd58e80..7d1c211d 100644 --- a/origami/models/rtu/discriminators.py +++ b/origami/models/rtu/discriminators.py @@ -15,7 +15,7 @@ KernelRequests, SystemRequests, ], - Field(discriminator='channel_prefix'), + Field(discriminator="channel_prefix"), ] # Use: pydantic.pares_obj_as(RTUResponse, ) @@ -29,7 +29,7 @@ KernelResponses, SystemResponses, ], - Field(discriminator='channel_prefix'), + Field(discriminator="channel_prefix"), ], RTUError, BaseRTUResponse, diff --git a/origami/models/rtu/errors.py b/origami/models/rtu/errors.py index ff4f91d5..82e9d65c 100644 --- a/origami/models/rtu/errors.py +++ b/origami/models/rtu/errors.py @@ -11,24 +11,24 @@ class ErrorData(BaseModel): # Error when we send over a request that doesn't match any handlers class InvalidEvent(BaseRTUResponse): - event: Literal['invalid_event'] = 'invalid_event' + event: Literal["invalid_event"] = "invalid_event" data: ErrorData # Error when the payload of our request has a validation error class InvalidData(BaseRTUResponse): - event: Literal['invalid_data'] = 'invalid_data' + event: Literal["invalid_data"] = "invalid_data" data: ErrorData # Error when RTU session isn't authenticated or the request does not pass RBAC checks class PermissionDenied(BaseRTUResponse): - event: Literal['permission_denied'] = 'permission_denied' + event: Literal["permission_denied"] = "permission_denied" data: ErrorData class InconsistentStateEvent(BaseRTUResponse): - event: Literal['inconsistent_state_event'] = 'inconsistent_state_event' + event: Literal["inconsistent_state_event"] = "inconsistent_state_event" data: ErrorData diff --git a/origami/notebook/builder.py b/origami/notebook/builder.py index 34410ddc..1630cc89 100644 --- a/origami/notebook/builder.py +++ b/origami/notebook/builder.py @@ -112,7 +112,7 @@ def apply_delta(self, delta: FileDelta) -> None: handler(delta) self.last_applied_delta_id = delta.id except Exception as e: # noqa: E722 - logger.exception("Error squashing Delta into NotebookBuilder", extra={'delta': delta}) + logger.exception("Error squashing Delta into NotebookBuilder", extra={"delta": delta}) raise e def add_cell(self, delta: NBCellsAdd): @@ -199,7 +199,7 @@ def update_cell_metadata(self, delta: CellMetadataUpdate): if delta.resource_id in self.deleted_cell_ids: logger.debug( f"Skipping update_cell_metadata for deleted cell {delta.resource_id}", - extra={'delta_properties_path': delta.properties.path}, + extra={"delta_properties_path": delta.properties.path}, ) return @@ -210,7 +210,7 @@ def update_cell_metadata(self, delta: CellMetadataUpdate): # and we end up emitting a cell execution timing metadata as it gets deleted logger.warning( "Got update_cell_metadata for cell that isn't in notebook or deleted_cell_ids", # noqa: E501 - extra={'delta_properties_path': delta.properties.path}, + extra={"delta_properties_path": delta.properties.path}, ) return diff --git a/pyproject.toml b/pyproject.toml index 7027f1da..d13e3ffd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,6 @@ exclude = ''' | profiling )/ ''' -skip-string-normalization = true [tool.isort] line_length = 100