o "4g'@s0dZddlmZddlmZmZddlmZddlm Z m Z m Z m Z ddl mZddl mZdd lmZdd lmZmZmZmZdd lmZmZdd lmZmZmZmZe rhdd lm Z ddl!m"Z"ddl#m$Z$ej%j&Z&ej%j'Z'ej(j)Z)GdddedZ*GdddeZ+Gddde+Z,Gddde,Z-dS)z7 Objects to support the COPY protocol (async version). ) annotations)ABCabstractmethod) TracebackType)Any AsyncIteratorSequence TYPE_CHECKING)pq)errors)Self)BaseCopyMAX_BUFFER_SIZE QUEUE_SIZE PREFER_FLUSH)copy_tocopy_end)aspawnagatherAQueueAWorker)Buffer) AsyncCursor)AsyncConnectioncseZdZUdZdZded<dddd2fd d Zd3ddZd4ddZd5ddZd6dd Z d7d"d#Z d8d%d&Z d9d)d*Z d:d-d.Z d;d0d1ZZS)< AsyncCopyajManage an asynchronous :sql:`COPY` operation. :param cursor: the cursor where the operation is performed. :param binary: if `!True`, write binary format. :param writer: the object to write to destination. If not specified, write to the `!cursor` connection. Choosing `!binary` is not necessary if the cursor has executed a :sql:`COPY` operation, because the operation result describes the format too. The parameter is useful when a `!Copy` object is created manually and no operation is performed on the cursor, such as when using ``writer=``\ `~psycopg.copy.FileWriter`. psycopg AsyncWriterwriterN)binaryrcursorAsyncCursor[Any]r bool | NoneAsyncWriter | Nonecs.tj||d|st|}||_|j|_dS)N)r)super__init__AsyncLibpqWriterrwrite_write)selfr rr __class__\/var/www/html/authentication-server/venv/lib/python3.10/site-packages/psycopg/_copy_async.pyr%2s  zAsyncCopy.__init__returnr cs||SN)_enterr)r,r,r- __aenter__@szAsyncCopy.__aenter__exc_typetype[BaseException] | Noneexc_valBaseException | Noneexc_tbTracebackType | NoneNonecs||IdHdSr/)finish)r)r3r5r7r,r,r- __aexit__DszAsyncCopy.__aexit__AsyncIterator[Buffer]cCs" |IdH}|s dS|Vq)z5Implement block-by-block iteration on :sql:`COPY TO`.TN)readr)datar,r,r- __aiter__NszAsyncCopy.__aiter__rc|j|IdHS)z Read an unparsed row after a :sql:`COPY TO` operation. Return an empty string when the data is finished. N) connectionwait _read_genr1r,r,r-r=VszAsyncCopy.readAsyncIterator[tuple[Any, ...]]cCs& |IdH}|durdS|Vq)z Iterate on the result of a :sql:`COPY TO` operation record by record. Note that the records returned will be tuples of unparsed strings or bytes, unless data types are specified using `set_types()`. TN)read_row)r)recordr,r,r-rows^szAsyncCopy.rowstuple[Any, ...] | NonecrA)a Read a parsed row of data from a table after a :sql:`COPY TO` operation. Return `!None` when the data is finished. Note that the records returned will be tuples of unparsed strings or bytes, unless data types are specified using `set_types()`. N)rBrC _read_row_genr1r,r,r-rFks zAsyncCopy.read_rowbuffer Buffer | strc*|j|}|r||IdHdSdS)z Write a block of data to a table after a :sql:`COPY FROM` operation. If the :sql:`COPY` is in binary format `!buffer` must be `!bytes`. In text mode it can be either `!bytes` or `!str`. N) formatterr'r()r)rKr?r,r,r-r'vs  zAsyncCopy.writerow Sequence[Any]crM)z=Write a record to a table after a :sql:`COPY FROM` operation.N)rN write_rowr()r)rOr?r,r,r-rQs  zAsyncCopy.write_rowexccs|jtkr#|j}|r||IdH|j|IdHd|_dS|s'dS|jj t kr/dS|j IdH|j |IdHdS)aTerminate the copy operation and free the resources allocated. You shouldn't need to call this function yourself: it is usually called by exit. It is available if, despite what is documented, you end up using the `Copy` object outside a block. NT) _directionCOPY_INrNendr(rr: _finished_pgconntransaction_statusACTIVErB _try_cancelrC_end_copy_out_gen)r)rRr?r,r,r-r:s    zAsyncCopy.finish)r r!rr"rr#)r.r )r3r4r5r6r7r8r.r9)r.r<)r.r)r.rE)r.rI)rKrLr.r9)rOrPr.r9rRr6r.r9)__name__ __module__ __qualname____doc____annotations__r%r2r;r@r=rHrFr'rQr: __classcell__r,r,r*r-rs      rzAsyncConnection[Any]c@s*eZdZdZed ddZddd d ZdS)rzG A class to write copy data somewhere (for async connections). r?rr.r9cdS)zWrite some data to destination.Nr,r>r,r,r-r'szAsyncWriter.writeNrRr6crc)z Called when write operations are finished. If operations finished with an error, it will be passed to ``exc``. Nr,r)rRr,r,r-r:szAsyncWriter.finishr?rr.r9r/r\)r]r^r_r`rr'r:r,r,r,r-rs  rc@s4eZdZdZdZdddZdd d ZddddZd S)r&zE An `AsyncWriter` to write copy data to a Postgres database. psycopg.copyr r!cCs||_|j|_|jj|_dSr/)r rBpgconnrWr)r r,r,r-r%szAsyncLibpqWriter.__init__r?rr.r9c stt|tkr|jt|j|tdIdHdStdt|tD]}|jt|j|||ttdIdHq dS)Nflushr)lenrrBrCrrWrranger)r?ir,r,r-r's " zAsyncLibpqWriter.writeNrRr6cs||rdt|jd|}||jjd}nd}z|jt|j|IdH}Wntj y6|s3YdSw|g|j _ dS)Nzerror from Python: z - replace) typer_encoderW _encodingrBrCre QueryCanceledr _results)r)rRmsgbmsgresr,r,r-r:szAsyncLibpqWriter.finishr r!rer/r\)r]r^r_r`r%r'r:r,r,r,r-r&s   r&csJeZdZdZdZdfdd Zdd d Zdd dZddfdd ZZS)AsyncQueuedLibpqWriterz `AsyncWriter` using a buffer to queue data to write. `write()` returns immediately, so that the main thread can be CPU-bound formatting messages, while a worker thread can be IO-bound waiting to write on the connection. rfr r!cs(t|ttd|_d|_d|_dS)N)maxsize)r$r%rr_queue_worker _worker_errorrhr*r,r-r%s   zAsyncQueuedLibpqWriter.__init__r.r9c slz |jIdH}|sWdS|jt|j|tdIdHqty5}z ||_WYd}~dSd}~ww)zPush data to the server when available from the copy queue. Terminate reading when the queue receives a false-y value, or in case of error. The function is designed to be run in a separate task. TNri) r|getrBrCrrWr BaseExceptionr~)r)r?exr,r,r-workers zAsyncQueuedLibpqWriter.workerr?rcsz|js t|j|_|jr|jt|tkr!|j|IdHdStdt|tD]}|j|||tIdHq)dS)Nr) r}rrr~rkrr|putrlrmr,r,r-r' s   zAsyncQueuedLibpqWriter.writeNrRr6csR|jdIdH|jrt|jIdHd|_|jr|jt|IdHdS)N)r|rr}rr~r$r:rdr*r,r-r: szAsyncQueuedLibpqWriter.finishry)r.r9rer/r\) r]r^r_r`r%rr'r:rbr,r,r*r-rzs  rzN).r` __future__rabcrrtypesrtypingrrrr r r rs_compatr _copy_baserrrr generatorsrr_acompatrrrrr cursor_asyncrconnection_asyncr ExecStatusrTCOPY_OUTTransactionStatusrYrrr&rzr,r,r,r-s.        /