*Jf{u fdZddlZddlZddlZddlZddlZddlZddlmZm Z ddl m Z ddl m Z mZmZmZmZmZmZmZddl mZmZmZmZmZmZddlmZdd lmZdd lmZdd lm Z m!Z!m"Z"m#Z#dd l$m%Z%m&Z&dd l'm(Z(ddl)m*Z*ddl+m,Z,ddl-m.Z.m/Z/m0Z0erddl1m2Z2m3Z3ddl4m5Z5ddl6m7Z7ddl8m9Z9ej:j;Zej?j;Z;ej?j=Z=ej@jAZAej@jBZBejCjDZDdZEdeEzZFdZGejHdkZIGddee!ZJGddeJdZKGddeZLGd d!eLZMGd"d#eMZNGd$d%eLZOGd&d'eJd(ZPGd)d*eZQGd+d,eQZRGd-d.eRZSGd/d0eZTGd1d2eTZUGd3d4eTZV dSd5ee d6e#d7eeWd8eWfd9ZX dSd5ee d6e#d7eeWd8eWfd:ZYd;e d6e#d8ee dZ[ej\d?j]Z^ej\d@j]Z_ej\d?j`Zaej\d@j`ZbdAZcdBZddCZeejfdDZgdEdFdGdHdIdJdKdLZhehfdMeeidNeeieifd8eifdOZjejfdPZkdQehlDZmemfdMeeidNeeieifd8eifdRZne*re*joZoe*jpZpe*jqZqe*jrZrdSeXZoeYZpeZZqe[ZrdS)Tz psycopg copy support N)ABCabstractmethod) TracebackType)Any AsyncIteratorDictGenericIteratorListMatchIO)OptionalSequenceTupleTypeUnion TYPE_CHECKING)pq)adapt)errors)BufferConnectionTypePQGen Transformer) create_taskSelf)connection_summary)_psycopg)pgconn_encoding) copy_fromcopy_tocopy_end) BaseCursorCursor) AsyncCursor) Connection)AsyncConnectioniidarwinceZdZUdZded<dddddeefd Zd efd Z dd Z d e e e efd dfdZd eefdZd eeeedffdZdeed edfdZdS)BaseCopya Base implementation for the copy user interface. Two subclasses expose real methods with the sync/async differences. The difference between the text and binary format is managed by two different `Formatter` subclasses. Writing (the I/O part) is implemented in the subclasses by a `Writer` or `AsyncWriter` instance. Normally writing implies sending copy data to a database, but a different writer might be chosen, e.g. to stream data into a file for later use. Formatter formatterNbinarycursorzBaseCursor[ConnectionType, Any]r0c@||_|j|_|jj|_|j}|ra|j|_|jtkrD|jtkr4tj dtj |jj dn t|_|t|o|j}t!|ddpt#j|}|rt'||_n(t+|t-|j|_d|_dS)NzAthe cursor should have performed a COPY operation; its status is z instead_tx)encodingF)r1 connectionpgconn_pgconnpgresultstatus _directionCOPY_INCOPY_OUTeProgrammingErrorr ExecStatusnamebool binary_tuplesgetattrrrBinaryFormatterr. TextFormatterr _finished)selfr1r0resulttxs O/home/alex/cs2snipeproduction/venv/lib/python3.11/site-packages/psycopg/copy.py__init__zBaseCopy.__init__Us   +-   &$mDO'))do.I.I(T&(mDO&D&D&ITTT &DO >&9V%9::F!&%66S%:KF:S:S  W,R00DNN*2 8U8UVVVDNreturnc|jjd|jj}t|j}d|d|dt |ddS)N.< z at 0xx>) __class__ __module__ __qualname__rr7id)rGclsinfos rJ__repr__zBaseCopy.__repr__usX*JJT^-HJJ!$,//232222RXX22222rLc2|jrtddS)Nz!copy blocks can be used only once)rF TypeErrorrGs rJ_enterzBaseCopy._enterzs( > A?@@ @ A ArLtypesc|jjjfd|D}|jtkr,|jj||jjdS|jj ||jjdS)a Set the types expected in a COPY operation. The types must be specified as a sequence of oid or PostgreSQL type names (e.g. ``int4``, ``timestamptz[]``). This operation overcomes the lack of metadata returned by PostgreSQL when a COPY operation begins: - On :sql:`COPY TO`, `!set_types()` allows to specify what types the operation returns. If `!set_types()` is not used, the data will be returned as unparsed strings or bytes instead of Python objects. - On :sql:`COPY FROM`, `!set_types()` allows to choose what type the database expects. This is especially useful in binary copy, because PostgreSQL will apply no cast rule. chg|].}t|tr|n|/S) isinstanceintget_oid).0tregistrys rJ z&BaseCopy.set_types..s:PPPQZ3''@X-=-=a-@-@PPPrLN) r1adaptersr_r:r;r. transformerset_dumper_typesformatset_loader_types)rGr_oidsrhs @rJ set_typeszBaseCopy.set_types~s&;'-PPPP%PPP ?g % % N & 7 7dn>S T T T T T N & 7 7dn>S T T T T TrLc#K|jrtdSt|jEd{V}t |tr|Sd|_|j}||nd|j_tdS)NrLT)rF memoryviewr!r7rccommand_tuplesr1 _rowcount)rGresnrowss rJ _read_genzBaseCopy._read_gens > #c?? ""4<00000000 c: & & J ").): #rL.c#K|Ed{V}|sdS|j|}|#|Ed{Vd|_dS|SNT)rxr. parse_rowrF)rGdatarows rJ _read_row_genzBaseCopy._read_row_gens..******** 4n&&t,, ;~~'' ' ' ' ' ' ' '!DN4 rLexcc#K|sdS|jjtkrdS|j |Ed{Vr |Ed{VdSdS#t j$rYdSwxYwN)r7transaction_statusACTIVEr5cancelrxr= QueryCanceledrGrs rJ_end_copy_out_genzBaseCopy._end_copy_out_gens  F < *f 4 4 F     "nn........ #nn........         DD s5A11BBrMN)__name__rUrV__doc____annotations__rrArKstrrZr^rrrdrprrrxrrr~ BaseExceptionrrbrLrJr,r,DsG   "& 1 @3#3333 AAAAUxc3h8UTUUUU:5=& uXeCHo%>?    Xm%<trLr,ceZdZUdZdZded<dddddd eededffd Zd efd Z d ee e dee dee d dfdZ d eefdZd efdZd eeedffdZd eeedffdZdeeefd dfdZdeed dfdZdee d dfdZxZS)Copya]Manage a :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`. psycopgWriterwriterNr0rr1 Cursor[Any]r0ct|||st|}||_|j|_dSNr/)superrK LibpqWriterrwrite_writerGr1r0rrTs rJrKz Copy.__init__sK /// ) ((F l rLrMc.||Srr^r]s rJ __enter__zCopy.__enter__s  rLexc_typeexc_valexc_tbc0||dSrfinishrGrrrs rJ__exit__z Copy.__exit__s GrLc#BK |}|sdS|V)z5Implement block-by-block iteration on :sql:`COPY TO`.TNreadrGr|s rJ__iter__z Copy.__iter__s3 99;;D JJJ  rLcZ|j|S)z Read an unparsed row after a :sql:`COPY TO` operation. Return an empty string when the data is finished. r5waitrxr]s rJrz Copy.read s$ ##DNN$4$4555rL.c#BK |}|dS|V)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()`. TNread_rowrGrecords rJrowsz Copy.rowss0 ]]__F~LLL  rLcZ|j|S)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()`. r5rr~r]s rJrz Copy.read_rows&##D$6$6$8$8999rLbuffercl|j|}|r||dSdS)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`. Nr.rrrGrr|s rJrz Copy.write)sC~##F++   KK       rLr}cl|j|}|r||dSdS)z=Write a record to a table after a :sql:`COPY FROM` operation.Nr. write_rowrrGr}r|s rJrzCopy.write_row4sA~'',,   KK       rLrc&|jtkrS|j}|r|||j|d|_dS|j | |dS)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. TN r:r;r.endrrrrFr5rrrGrr|s rJrz Copy.finish:s ?g % %>%%''D " D!!! K  s # # #!DNNN O !7!7!(=1>d>>>>>>>>rLrzConnection[Any]cPeZdZdZededdfdZddeeddfdZ dS) rz/ A class to write copy data somewhere. r|rMNcdS)z1 Write some data to destination. Nrbrs rJrz Writer.writePs rLrcdS)z Called when write operations are finished. If operations finished with an error, it will be passed to ``exc``. Nrbrs rJrz Writer.finishWs rLr rrUrVrrrrrrrrbrLrJrrKsx & T   ^   (=1 T      rLrcHeZdZdZd dZdeddfdZd d eeddfd Z dS) rz? A `Writer` to write copy data to a Postgres database. r1rcN||_|j|_|jj|_dSrr1r5r6r7rGr1s rJrKzLibpqWriter.__init__e#  +- rLr|rMNc vt|tkr6|jt |j|t dStdt|tD]F}|jt |j|||tzt GdSNflushrlenMAX_BUFFER_SIZEr5rr"r7 PREFER_FLUSHrangerGr|is rJrzLibpqWriter.writejs t99 ' ' O tHH3HHC::odl;;YGGDDD )/&&x d'C'CDDC%(5DK         s -BBBr1rr rrUrVrrKrrrrrrbrLrJrr`s.... & T    ))(=1)T))))))rLrc^eZdZdZd fd Zd dZdeddfd Zdd ee ddffd Z xZ S)QueuedLibpqWritera A writer using a buffer to queue data to write to a Postgres database. `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. r1rct|tjt|_d|_d|_dSN)maxsize)rrKqueueQueue QUEUE_SIZE_queue_worker _worker_errorrGr1rTs rJrKzQueuedLibpqWriter.__init__sD    +0;z+J+J+J 37 6:rLrMNc |jdd}|sdS|jt |j|t U#t$r}||_Yd}~dSd}~wwxYw)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 thread. TiQ)blocktimeoutrN) rgetr5rr"r7rrr)rGr|exs rJworkerzQueuedLibpqWriter.workers $ V{T<HHE$$WT\4|%T%T%TUUU  V  $ $ $!#D        $sA5A A4"A//A4r|c|jsDtj|j|_d|j_|j|jr|jt|tkr|j |dStdt|tD],}|j |||tz-dS)N)targetTr) r threadingThreadrdaemonstartrrrrputrrs rJrzQueuedLibpqWriter.writes| !$+4;???DL"&DL  L      %$ $ t99 ' ' KOOD ! ! ! ! !1c$ii99 ? ? Q_)<%< =>>>> ? ?rLrc|jd|jr |jd|_|jr|jt |dSNrL)rrrjoinrrrrGrrTs rJrzQueuedLibpqWriter.finishsn  < L     DL   %$ $ srLrrr rrUrVrrKrrrrrrrrs@rJrrs;;;;;;$$$$$?&?T????*  (=1 T          rLrc:eZdZdZdeefdZdeddfdZdS) FileWriterz A `Writer` to write copy data to a file-like object. :param file: the file where to write copy data. It must be open for writing in binary mode. filec||_dSr)r)rGrs rJrKzFileWriter.__init__s  rLr|rMNc:|j|dSr)rrrs rJrzFileWriter.writes rL) rrUrVrr bytesrKrrrbrLrJrrs_RY&TrLrceZdZUdZdZded<dddddd eededffd Zd efd Z d ee e dee dee d dfdZ d eefdZd efdZd eeedffdZd eeedffdZdeeefd dfdZdeed dfdZdee d dfdZxZS) AsyncCopyz-Manage an asynchronous :sql:`COPY` operation.r AsyncWriterrNrr1AsyncCursor[Any]r0ct|||st|}||_|j|_dSr)rrKAsyncLibpqWriterrrrrs rJrKzAsyncCopy.__init__sK /// .%f--F l rLrMc2K||Srrr]s rJ __aenter__zAsyncCopy.__aenter__s  rLrrrc@K||d{VdSrrrs rJ __aexit__zAsyncCopy.__aexit__s2 kk'"""""""""""rLcPK |d{V}|sdS|WV$rrrs rJ __aiter__zAsyncCopy.__aiter__sG $$$$$$D JJJJ  rLcjK|j|d{VSrrr]s rJrzAsyncCopy.read s8_))$..*:*:;;;;;;;;;rL.cPK |d{V}|dS|WV$rrrs rJrzAsyncCopy.rows sD ==??******F~LLLL  rLcjK|j|d{VSrrr]s rJrzAsyncCopy.read_rows:_))$*<*<*>*>?????????rLrc|K|j|}|r||d{VdSdSrrrs rJrzAsyncCopy.writesW~##F++  $++d## # # # # # # # # # $ $rLr}c|K|j|}|r||d{VdSdSrrrs rJrzAsyncCopy.write_rowsW~'',,  $++d## # # # # # # # # # $ $rLrcNK|jtkr_|j}|r||d{V|j|d{Vd|_dS|j | |d{VdSrzrrs rJrzAsyncCopy.finish!s ?g % %>%%''D (kk$'''''''''+$$S)) ) ) ) ) ) ) )!DNNN/&&t'='=c'B'BCC C C C C C C C C CrL)rrUrVrrrrArKrrrrrrrrrrrrrrrrrrrrrrs@rJr r s77J  "&*. # # #" # # ' # # # # # #$#4 ./#-(#' #  ####v!6HH3HHC::odl;;YGGDDD ),,XdlD-I-IJJJJJJJJC%(5DK         s 3B B#"B#r1r rrrbrLrJrr8s....  4    )) 7)4))))))rLrc^eZdZdZd fd Zd dZdeddfd Zdd ee ddffd Z xZ S)AsyncQueuedLibpqWriterz An `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. r1r ct|tjt|_d|_dSr)rrKasynciorrrrrs rJrKzAsyncQueuedLibpqWriter.__init__ns:    -4]:-N-N-N 7; rLrMNcK |jd{V}|sdS|jt |j|d{VW)zPush data to the server when available from the copy queue. Terminate reading when the queue receives a false-y value. The function is designed to be run in a separate task. TN)rrr5rr"r7rs rJrzAsyncQueuedLibpqWriter.workertsv D********D /&&wt|T'B'BCC C C C C C C C  DrLr|cK|js&t||_t|tkr"|j|d{VdStdt|tD]2}|j|||tzd{V3dSNr)rrrrrrrrrs rJrzAsyncQueuedLibpqWriter.writes| 6&t{{}}55DL t99 ' '+//$'' ' ' ' ' ' ' ' ' '1c$ii99 E Ekood1q?/B+B&CDDDDDDDDDD E ErLrcK|jdd{V|jr&tj|jd{Vd|_t |d{VdSr)rrrr&gatherrrrs rJrzAsyncQueuedLibpqWriter.finishskooc""""""""" < ... . . . . . . .DLggnnS!!!!!!!!!!!rLr"rrrrs@rJr$r$es<<<<<< D D D D E E4 E E E E"" 7"4""""""""""rLr$ceZdZUdZejed<defdZe de de e e dffdZe d ee efde fd Ze d ee de fd Ze de fd ZdS)r-z@ A class which understand a copy format (text, binary). rmrkcH||_t|_d|_dSNF)rk bytearray _write_buffer _row_mode)rGrks rJrKzFormatter.__init__s!&&[[rLr|rM.cdSrrbrs rJr{zFormatter.parse_rowsDGCrLrcdSrrbrGrs rJrzFormatter.writes;>3rLr}cdSrrb)rGr}s rJrzFormatter.write_rows7:srLcdSrrbr]s rJrz Formatter.ends SrLN)rrUrVrrFormatrrrKrrrrrr{rrrrrrrbrLrJr-r-s IK GfG%S/)BGGG^G>E&#+.>6>>>^>:Xc]:v:::^: V   ^   rLr-ceZdZeZddedeffd Zdede e e dffdZ d e eefdefd Zd ee defd Zdefd Zde eefdefdZxZS)rEutf-8rkr4cXt|||_dSr)rrK _encoding)rGrkr4rTs rJrKzTextFormatter.__init__s& %%%!rLr|rM.c4|rt||jSdSr)parse_row_textrkrs rJr{zTextFormatter.parse_rows"  !$(899 94rLrc>||}d|_|Srz _ensure_bytes_signature_sentrs rJrzTextFormatter.write"!!&))# rLr}cd|_t||j|jt |jt kr|jt c}|_|SdSNTrL)r1format_row_textrkr0r BUFFER_SIZEr/rGr}rs rJrzTextFormatter.write_rowsYT-t/ABBB t! " "[ 0 0)-);Y[[ &FD&M3rLc<|jtc}|_|Sr)r0r/r4s rJrzTextFormatter.ends%)%7"" rLcdt|tr||jS|Sr)rcrrr;rs rJr@zTextFormatter._ensure_bytess/ dC  ;;t~.. . KrL)r9)rrUrVTEXTrmrrrKrrrrr{rrrrrr@rrs@rJrErEs F""K"3""""""f%S/)B E&#+.6 Xc] v    V% "4rLrEceZdZeZdeffd Zdedee e dffdZ de ee fdefdZd ee defd Zdefd Zde ee fdefd ZxZS) rDrkcXt|d|_dSr.)rrKrA)rGrkrTs rJrKzBinaryFormatter.__init__s) %%%$rLr|rM.c|js]|dtttkrtjdd|_|ttd}n |t krdSt ||jS)Nz5binary copy doesn't start with the expected signatureT)rAr_binary_signaturer= DataError_binary_trailerparse_row_binaryrkrs rJr{zBinaryFormatter.parse_rows# ,c+,,,-1BBBkK$(D -..001DD _ $ $4d&6777rLrc>||}d|_|Srzr?rs rJrzBinaryFormatter.writerBrLr}cd|_|js|xjtz c_d|_t ||j|jt |jtkr|jtc}|_|SdSrD) r1rAr0rNformat_row_binaryrkrrFr/rGs rJrzBinaryFormatter.write_rows# (   "3 3  #'D #t/1CDDD t! " "[ 0 0)-);Y[[ &FD&M3rLc|js+|xjtz c_|xjtz c_n|jr|xjtz c_|jt c}|_|Sr)rAr0rNrPr1r/r4s rJrzBinaryFormatter.endsv# 2   "3 3     / 1    ^ 2   / 1  %)%7"" rLcNt|trtd|S)Nz6cannot copy str data in binary mode: use bytes instead)rcrr\rs rJr@zBinaryFormatter._ensure_bytess* dC  TUU U KrL)rrUrVBINARYrmrrKrrrrr{rrrrrrr@rrs@rJrDrDs F%K%%%%%% 8f 8%S/)B 8 8 8 8E&#+.6 Xc]v V&% "4rLrDr}rIoutrMc|t}|s|dz }|S|D]b}|T||t}||}|tt |z }n|dz }|dz }cd|dd<|S)z6Convert a row of objects to the data to send for copy.N \N rr)r/ get_dumperPY_TEXTdump_dump_resub _dump_sub)r}rIrXitemdumperbs rJ_format_row_textrf%s {kk  u    ]]411F D!!A 8<< 1-- -CC 6MC u CH JrLc(|t}|tt|z }||tgt|z}|D]3}|%|t t|z }||z })|t z }4|S)z=Convert a row of objects to the data to send for binary copy.)r/ _pack_int2r dump_sequence PY_BINARY _pack_int4 _binary_null)r}rIrXadaptedres rJ_format_row_binaryrn=s {kk:c#hh  CsYK#c(($:;;G    = :c!ff%% %C 1HCC < CC JrLr|.ct|tst|}|d}|ddd|d<d|D}||S)Nr\rrc^g|]*}|dkrdntt|+S)r[N)_load_rera _load_sub)rffs rJriz#_parse_row_text..Us2 O O O11;;44HLLA$>$> O O OrL)rcrsplit load_sequence)r|rIfieldsr}s rJ_parse_row_textrwPsl dE " "T{{ ZZ  FCRCF2J O O O O OC  C  rLc<g}t|dd}d}t|D]^}t||d}|dz }|dkr&|||||z||z }I|d_||S)Nrr)) _unpack_int2r _unpack_int4appendru)r|rIr}nfieldsposrlengths rJ_parse_row_binaryrYs"$C4##A&G C 7^^dC((+ q Q;; JJtC#,./ 0 0 0 6MCC JJt      C  rLz!hz!isPGCOPY sss [ \\]s\bs\ts\ns\vs\fs\rs\\)r\rZ   \m__mapc8||dSr)grouprrs rJrbrb  rLs \\[btnvfr\\]ci|]\}}|| Srbrb)rfkvs rJ rs 2 2 2tq!a 2 2 2rLc8||dSr)rrs rJrrrrrrLr)srresysrstructr&rabcrrr_rtypingrrrr r r r r rrrrrrrrrr=rrrr_compatrrpq.miscr_cmoduler _encodingsr generatorsr!r"r#r1r$r% cursor_asyncr&r5r'connection_asyncr(PyFormatrJr^rWrjr7r?r;r<TransactionStatusrrFrrplatformrr,rrrrrr r rr$r-rErDr/rfrnrwrStructpackrhrk unpack_fromrzr{rNrPrlcompiler` _dump_replrrbrqitems _load_replrrrErTr=rQrbrLrJrsj  ########OOOOOOOOOOOOOOOOOOOOHHHHHHHHHHHHHHHH;;;;;;;;;;;;&&&&&&&&''''''''''''44444444442********))))))&&&&&&111111 .  N !  y~   -  = !  $  k/ |x' JJJJJw~&JJJZw>w>w>w>w>8% &w>w>w>t     S   **)*)*)*)*)&*)*)*)ZBBBBB BBBJ        HDHDHDHDHD/0HDHDHDV      #    *)*)*)*)*){*)*)*)Z1"1"1"1"1"-1"1"1"h!!!!!!!!2)))))I)))XDDDDDiDDDPEI #'.6y.A2EI #'.6y.A&!&!k!eCHo!!!! !F ! !c3h ! ! ! ! V]4 % V]4 % v}T"". v}T"".  " 2:+ , ,