diff --git a/README.md b/README.md index ab7e4aa..6ce451a 100644 --- a/README.md +++ b/README.md @@ -2,12 +2,30 @@ Python common library codes to be shared by other RDA python utility programs. -## Installing and using in another RDA python repo +## Environment setup -`rda-python-common` is the foundation that every other `rda-python-*` repo -builds on. To consume it from a new or existing repo, follow these steps. +Create a Python environment first; the install command in the next section +runs inside whichever environment you activate here. + +### Option A — Python venv (DECS machines) + +```bash +python3 -m venv $ENVHOME # e.g. /glade/u/home/gdexdata/gdexmsenv +source $ENVHOME/bin/activate +``` + +### Option B — Conda (DAV/Casper) -### 1. Install the package +```bash +conda create --prefix $ENVHOME python=3.12 # e.g. /glade/work/gdexdata/conda-envs/pg-gdex +conda activate $ENVHOME +``` + +## Installing rda-python-common + +Pick whichever install mode fits your workflow. All four pull in the +transitive dependencies (`psycopg`, `rda-python-globus`, `unidecode`, +`hvac`) automatically. For local development, clone this repo alongside your project and install it in editable mode so that changes are picked up without re-installing: @@ -18,6 +36,15 @@ cd rda-python-common pip install -e . ``` +To test a specific branch (e.g. an in-progress feature or fix branch), pass +`-b/--branch` to `git clone`: + +```bash +git clone -b https://github.com/NCAR/rda-python-common.git +cd rda-python-common +pip install -e . +``` + For a regular (non-editable) install from a checkout: ```bash @@ -30,10 +57,75 @@ For a production install on a system that uses the published distribution: pip install rda_python_common ``` -The package brings in its own transitive dependencies (`psycopg2-binary`, -`rda-python-globus`, `unidecode`, `hvac`). +### PostgreSQL driver: psycopg v3 (default) and psycopg2 (fallback) + +`rda-python-common` uses **psycopg v3** by default. `pg_dbi.py` +auto-detects which driver is installed at import time and prefers psycopg v3 +when both are present; no code changes are needed to switch drivers. -### 2. Declare it as a dependency in your project +The required dependency is the base `psycopg` package, which works whether +psycopg was compiled from source or installed via a binary wheel. If psycopg +is not available on your system, install whichever driver works: + +```bash +pip install psycopg || pip install psycopg2 +``` + +To explicitly install the legacy psycopg2 driver: + +```bash +pip install "rda_python_common[psycopg2]" # build from source +pip install "rda_python_common[psycopg2-binary]" # pre-built wheel +``` + +## Configuration: COMMONUSER and ADMINUSER + +`PGLOG['COMMONUSER']` is the shared common user that setuid-wrapped programs +execute as (default `gdexdata`), and `PGLOG['ADMINUSER']` is the admin +specialist user that receives email notifications and is permitted to invoke +`pgstart_` (default `zji`). + +Both values are initialized via the `SETPGLOG(key, default)` helper, which +reads the environment variable `PG` and falls back to the supplied +default when the variable is unset: + +```python +# pg_log.py (class-based) +self.SETPGLOG("COMMONUSER", "gdexdata") # reads $PGCOMMONUSER +self.SETPGLOG("ADMINUSER", "zji") # reads $PGADMINUSER + +# PgLOG.py (module-level) exposes the same helper as a function +SETPGLOG("COMMONUSER", "gdexdata") +SETPGLOG("ADMINUSER", "zji") +``` + +To override the defaults per environment **once** so the values persist +across `pip install --upgrade`, set the environment variables: + +```bash +export PGCOMMONUSER=gdexdata # overrides PGLOG['COMMONUSER'] +export PGADMINUSER=zji # overrides PGLOG['ADMINUSER'] +``` + +Place these `export` lines in `$ENVHOME/bin/activate` (venv), or set them as +conda environment variables so they are applied whenever the environment is +activated: + +```bash +conda env config vars set PGCOMMONUSER=gdexdata PGADMINUSER=zji +conda activate $ENVHOME # reactivate to pick up the values +``` + +If the variables are unset, the built-in defaults (`gdexdata` / `zji`) are +used, preserving existing behavior. + +## Using rda-python-common in another RDA python repo + +`rda-python-common` is the foundation that every other `rda-python-*` repo +builds on. Once it is installed in the active environment, consuming it from +a new or existing repo takes three short steps. + +### 1. Declare it as a dependency in your project Add `rda_python_common` to the `dependencies` list of your project's `pyproject.toml` so that downstream installs pull it in automatically: @@ -52,9 +144,10 @@ This is the same pattern used by `rda-python-dsarch`, `rda-python-dsupdt`, `rda-python-dsrqst`, `rda-python-dscheck`, `rda-python-metrics`, and `rda-python-miscs`. -### 3. Import the modules you need +### 2. Import the modules you need -Two import styles are supported (see [Usage examples](#usage-examples) below): +Two import styles are supported (see [Usage examples](#usage-examples) below +for fuller patterns): ```python # Preferred for new code -- import the class from the lower-case module @@ -66,26 +159,26 @@ from rda_python_common import PgLOG, PgDBI PgLOG.pglog("hello", PgLOG.LOGWRN) ``` -### 4. Verify the install +### 3. Verify the install ```bash python -c "import rda_python_common; print(rda_python_common.__version__)" ``` -You should see the installed version (currently `2.1.11`). If the import +You should see the installed version (currently `3.0.0`). If the import fails, double-check that the active Python environment is the one where you ran `pip install`. ## Modules -All shared functionality lives under `src/rda_python_common/` and is organised as -a single-inheritance class hierarchy. Each module defines exactly one class; -later classes extend earlier ones, so an application that instantiates the -top-of-chain class (typically `PgOPT` or `PgCMD`) gets every helper through one -object. +All shared functionality lives under `src/rda_python_common/` and is organised +as a (mostly) single-inheritance class hierarchy. Each module defines exactly +one class; later classes extend earlier ones, so an application that +instantiates the top-of-chain class (typically `PgOPT` or `PgCMD`) gets every +helper through one object. -Inheritance tree (top-down; multi-inheritance shown as two arrows -converging on the same child): +The inheritance tree below is read top-down; the two multi-inheritance joins +are shown as two arrows converging on the same child: ``` PgLOG @@ -122,6 +215,8 @@ The tree is single inheritance everywhere except at two join points: operations (`PgDBI`) it needs to keep the shared `wfile` table and the per-dataset `wfile_` partitions in sync. +Each class lives in its own module. Walking the tree from the root: + - **`pg_log.py`** — `PgLOG`. Root of the hierarchy. Provides the central logging facility (bit-mask `logact` flags such as `MSGLOG`, `WARNLG`, `ERRLOG`, `EXITLG`), e-mail dispatch, system-command execution, process @@ -145,7 +240,8 @@ The tree is single inheritance everywhere except at two join points: long-running batch jobs coordinate cleanly. - **`pg_dbi.py`** — `PgDBI(PgLOG)`. PostgreSQL database interface built on - `psycopg2`. Wraps connection management, batch `INSERT`/`SELECT`/ + `psycopg` (v3 by default, with `psycopg2` as an opt-in fallback). Wraps + connection management, batch `INSERT`/`SELECT`/ `UPDATE`/`DELETE`, transaction control, and credential lookup from `.pgpass` or OpenBao. All RDA tools talk to the `rdadb` database through this class. @@ -179,9 +275,9 @@ The tree is single inheritance everywhere except at two join points: ## Usage examples -Each class lives in its own submodule. Import the class you need, then -either instantiate it directly or subclass it to add application-specific -state and methods. +The patterns below show the typical ways the classes above are used in +practice. Import the class you need, then either instantiate it directly or +subclass it to add application-specific state and methods. ### 1. Direct instantiation — use the helpers as-is diff --git a/pyproject.toml b/pyproject.toml index 798f2a6..201f5ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "rda_python_common" -version = "2.1.11" +version = "3.0.0" authors = [ { name="Zaihua Ji", email="zji@ucar.edu" }, ] @@ -18,13 +18,20 @@ classifiers = [ "Development Status :: 5 - Production/Stable", ] dependencies = [ - "psycopg2-binary", + "psycopg", "psutil", "rda-python-globus", "unidecode", "hvac" ] +[project.optional-dependencies] +# Allow opting in to the legacy psycopg2 driver instead of psycopg (v3). +# pg_dbi.py auto-detects which driver is installed and prefers psycopg (v3) +# when both are available. +psycopg2 = ["psycopg2"] # psycopg2 built from source +psycopg2-binary = ["psycopg2-binary"] # psycopg2 pre-built C extension + [project.urls] "Homepage" = "https://github.com/NCAR/rda-python-common" diff --git a/requirements.txt b/requirements.txt index 7f4a1fe..6be6009 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ iniconfig==2.1.0 packaging==24.2 pluggy==1.5.0 -psycopg2-binary==2.9.10 +psycopg[binary] psutil pytest==8.3.5 rda-python-globus diff --git a/src/rda_python_common/PgDBI.py b/src/rda_python_common/PgDBI.py index 0bb7bf7..e83a819 100644 --- a/src/rda_python_common/PgDBI.py +++ b/src/rda_python_common/PgDBI.py @@ -17,12 +17,53 @@ import time import hvac from datetime import datetime -import psycopg2 as PgSQL -from psycopg2.extras import execute_values -from psycopg2.extras import execute_batch from os import path as op from . import PgLOG +# Prefer psycopg (v3); fall back to psycopg2 if v3 is not installed. +try: + import psycopg as PgSQL + PG_DRIVER = 'psycopg3' + + def execute_values(cursor, sql, argslist, page_size=100): + """Compatibility shim for psycopg2.extras.execute_values on psycopg3. + + Rewrites ``VALUES %s`` placeholder to ``VALUES (%s, %s, ...)`` based on + the column count inferred from the first row, then dispatches to + psycopg3's ``executemany`` (which already batches efficiently). + """ + if not argslist: return + ncol = len(argslist[0]) + row_ph = '(' + ','.join(['%s'] * ncol) + ')' + new_sql = re.sub(r'(?i)\bVALUES\s+%s\b', 'VALUES ' + row_ph, sql, count=1) + cursor.executemany(new_sql, argslist) + + def execute_batch(cursor, sql, argslist, page_size=100): + """Compatibility shim for psycopg2.extras.execute_batch on psycopg3.""" + cursor.executemany(sql, argslist) + + def get_pgcode(pgerr): + """Return SQLSTATE for a psycopg3 error (via err.diag.sqlstate).""" + diag = getattr(pgerr, 'diag', None) + return getattr(diag, 'sqlstate', None) if diag is not None else None + + def get_pgerror(pgerr): + """Return primary error message for a psycopg3 error (via err.diag.message_primary).""" + diag = getattr(pgerr, 'diag', None) + return getattr(diag, 'message_primary', None) if diag is not None else None +except ImportError: + import psycopg2 as PgSQL + from psycopg2.extras import execute_values, execute_batch + PG_DRIVER = 'psycopg2' + + def get_pgcode(pgerr): + """Return SQLSTATE for a psycopg2 error (via err.pgcode).""" + return getattr(pgerr, 'pgcode', None) + + def get_pgerror(pgerr): + """Return primary error message for a psycopg2 error (via err.pgerror).""" + return getattr(pgerr, 'pgerror', None) + pgdb = None # reference to a connected database object curtran = 0 # 0 - no transaction, 1 - in transaction NMISSES = [] # array of mising userno @@ -439,8 +480,8 @@ def check_dberror(pgerr, pgcnt, sqlstr, ary, logact = PGDBI['ERRLOG']): ret = PgLOG.FAILURE - pgcode = pgerr.pgcode - pgerror = pgerr.pgerror + pgcode = get_pgcode(pgerr) + pgerror = get_pgerror(pgerr) dberror = "{} {}".format(pgcode, pgerror) if pgcode and pgerror else str(pgerr) if pgcnt < PgLOG.PGLOG['DBRETRY']: if not pgcode: @@ -517,7 +558,7 @@ def pgconnect(reconnect = 0, pgcnt = 0, autocommit = True): reconnect = 0 # initial connection while True: - config = {'database' : PGDBI['DBNAME'], + config = {'dbname' : PGDBI['DBNAME'], 'user' : PGDBI['LNNAME']} if PGDBI['DBSHOST'] == PgLOG.PGLOG['HOSTNAME']: config['host'] = 'localhost' @@ -526,7 +567,7 @@ def pgconnect(reconnect = 0, pgcnt = 0, autocommit = True): if not PGDBI['DBPORT']: PGDBI['DBPORT'] = get_dbport(PGDBI['DBNAME']) if PGDBI['DBPORT']: config['port'] = PGDBI['DBPORT'] config['password'] = '***' - sqlstr = "psycopg2.connect(**{})".format(config) + sqlstr = "{}.connect(**{})".format(PG_DRIVER, config) config['password'] = get_pgpass_password() if PgLOG.PGLOG['DBGLEVEL']: PgLOG.pgdbg(1000, sqlstr) try: diff --git a/src/rda_python_common/PgFile.py b/src/rda_python_common/PgFile.py index 9f58fa9..925b26b 100644 --- a/src/rda_python_common/PgFile.py +++ b/src/rda_python_common/PgFile.py @@ -788,7 +788,7 @@ def delete_backup_file(file, endpoint = None, logact = 0): return PgLOG.FAILURE # -# reset local file/directory information to make them writable for PgLOG.PGLOG['GDEXUSER'] +# reset local file/directory information to make them writable for PgLOG.PGLOG['COMMONUSER'] # file - file name (mandatory) # info - gathered file info with option 14, None means file not exists # diff --git a/src/rda_python_common/PgLOG.py b/src/rda_python_common/PgLOG.py index c10449d..c71b899 100644 --- a/src/rda_python_common/PgLOG.py +++ b/src/rda_python_common/PgLOG.py @@ -60,7 +60,7 @@ EMLSUM = (0x08000) # record as email summary EMEROL = (0x10000) # record error as email only EMLALL = (0x1D208) # all email acts -DOSUDO = (0x20000) # add 'sudo -u PGLOG['GDEXUSER']' +DOSUDO = (0x20000) # add 'sudo -u PGLOG['COMMONUSER']' NOTLOG = (0x40000) # do not log any thing OVRIDE = (0x80000) # do override existing file or record NOWAIT = (0x100000) # do not wait on globus task to finish @@ -100,9 +100,8 @@ 'BACKROOT': "/DRDATA/DECS", # backup path for desaster recovering tape on hpss 'OLDAROOT': "/FS/DSS", # old root path on hpss 'OLDBROOT': "/DRDATA/DSS", # old backup tape on hpss - 'GDEXUSER' : "gdexdata", # common gdex user name - 'GDEXEMAIL' : "zji", # specialist to receipt email intead of common gdex user name - 'SUDOGDEX' : 0, # 1 to allow sudo to PGLOG['GDEXUSER'] + # COMMONUSER and ADMINUSER are set below via SETPGLOG (env overrides PG) + 'SUDOGDEX' : 0, # 1 to allow sudo to PGLOG['COMMONUSER'] 'HOSTNAME' : '', # current host name the process in running on 'OBJCTSTR' : "object", 'BACKUPNM' : "quasar", @@ -140,10 +139,22 @@ 'EMLPORT' : 25 } -PGLOG['RDAUSER'] = PGLOG['GDEXUSER'] +def SETPGLOG(key, default): + """Set ``PGLOG[key]`` from environment variable ``PG`` or fall back + to ``default`` if the variable is unset. Used to make per-environment + overrides (e.g. PGCOMMONUSER, PGADMINUSER) survive package upgrades.""" + PGLOG[key] = os.environ.get('PG' + key, default) + +SETPGLOG("COMMONUSER", "gdexdata") +SETPGLOG("ADMINUSER", "zji") + +PGLOG['RDAUSER'] = PGLOG['COMMONUSER'] PGLOG['RDAGRP'] = PGLOG['GDEXGRP'] -PGLOG['RDAEMAIL'] = PGLOG['GDEXEMAIL'] +PGLOG['RDAEMAIL'] = PGLOG['ADMINUSER'] PGLOG['SUDORDA'] = PGLOG['SUDOGDEX'] +# backwards-compat aliases (deprecated: use COMMONUSER / ADMINUSER) +PGLOG['GDEXUSER'] = PGLOG['COMMONUSER'] +PGLOG['GDEXEMAIL'] = PGLOG['ADMINUSER'] HOSTTYPES = { 'rda' : 'dsg_mach', @@ -317,15 +328,15 @@ def send_python_email(subject = None, receiver = None, msg = None, sender = None docc = False if cc else True if not sender: sender = PGLOG['CURUID'] - if sender != PGLOG['GDEXUSER']: docc = False - if sender == PGLOG['GDEXUSER']: sender = PGLOG['GDEXEMAIL'] + if sender != PGLOG['COMMONUSER']: docc = False + if sender == PGLOG['COMMONUSER']: sender = PGLOG['ADMINUSER'] if sender.find('@') == -1: sender += "@ucar.edu" if not receiver: receiver = PGLOG['EMLADDR'] if PGLOG['EMLADDR'] else PGLOG['CURUID'] - if receiver == PGLOG['GDEXUSER']: receiver = PGLOG['GDEXEMAIL'] + if receiver == PGLOG['COMMONUSER']: receiver = PGLOG['ADMINUSER'] if receiver.find('@') == -1: receiver += "@ucar.edu" - if docc and not re.match(PGLOG['GDEXUSER'], sender): add_carbon_copy(sender, 1) + if docc and not re.match(PGLOG['COMMONUSER'], sender): add_carbon_copy(sender, 1) emlmsg = EmailMessage() emlmsg.set_content(msg) emlmsg['From'] = sender @@ -1167,6 +1178,10 @@ def get_command(cmdstr = None): if not cmdstr: cmdstr = sys.argv[0] cmdstr = op.basename(cmdstr) + if cmdstr.startswith('setuid_'): + euser = pwd.getpwuid(os.geteuid()).pw_name + if euser == PGLOG['COMMONUSER']: + cmdstr = cmdstr[len('setuid_'):] ms = re.match(r'^(.+)\.(py|pl)$', cmdstr) if ms: return ms.group(1) @@ -1182,11 +1197,11 @@ def get_local_command(cmd, asuser = None): cuser = PGLOG['SETUID'] if PGLOG['SETUID'] else PGLOG['CURUID'] if not asuser or cuser == asuser: return cmd - if cuser == PGLOG['GDEXUSER']: + if cuser == PGLOG['COMMONUSER']: wrapper = "pgstart_" + asuser if valid_command(wrapper): return "{} {}".format(wrapper, cmd) - elif PGLOG['SUDOGDEX'] and asuser == PGLOG['GDEXUSER']: - return "sudo -u {} {}".format(PGLOG['GDEXUSER'], cmd) # sudo as user gdexdata + elif PGLOG['SUDOGDEX'] and asuser == PGLOG['COMMONUSER']: + return "sudo -u {} {}".format(PGLOG['COMMONUSER'], cmd) # sudo as user gdexdata return cmd @@ -1209,12 +1224,12 @@ def get_hpss_command(cmd, asuser = None, hcmd = None): if not hcmd: hcmd = 'hsi' if asuser and cuser != asuser: - if cuser == PGLOG['GDEXUSER']: + if cuser == PGLOG['COMMONUSER']: return "{} sudo -u {} {}".format(hcmd, asuser, cmd) # setuid wrapper as user asuser - elif PGLOG['SUDOGDEX'] and asuser == PGLOG['GDEXUSER']: - return "sudo -u {} {} {}".format(PGLOG['GDEXUSER'], hcmd, cmd) # sudo as user gdexdata + elif PGLOG['SUDOGDEX'] and asuser == PGLOG['COMMONUSER']: + return "sudo -u {} {} {}".format(PGLOG['COMMONUSER'], hcmd, cmd) # sudo as user gdexdata - if cuser != PGLOG['GDEXUSER']: + if cuser != PGLOG['COMMONUSER']: if re.match(r'^ls ', cmd) and hcmd == 'hsi': return "hpss" + cmd # use 'hpssls' instead of 'hsi ls' elif re.match(r'^htar -tvf', hcmd): @@ -1231,8 +1246,8 @@ def get_sync_command(host, asuser = None): host = get_short_host(host) - if (not (PGLOG['SETUID'] and PGLOG['SETUID'] == PGLOG['GDEXUSER']) and - (not asuser or asuser == PGLOG['GDEXUSER'])): + if (not (PGLOG['SETUID'] and PGLOG['SETUID'] == PGLOG['COMMONUSER']) and + (not asuser or asuser == PGLOG['COMMONUSER'])): return "sync" + host return host + "-sync" @@ -1246,7 +1261,7 @@ def set_suid(cuid = 0): if cuid != PGLOG['EUID'] or cuid != PGLOG['RUID']: os.setreuid(cuid, cuid) PGLOG['SETUID'] = pwd.getpwuid(cuid).pw_name - if not (PGLOG['SETUID'] == PGLOG['GDEXUSER'] or cuid == PGLOG['RUID']): + if not (PGLOG['SETUID'] == PGLOG['COMMONUSER'] or cuid == PGLOG['RUID']): set_specialist_environments(PGLOG['SETUID']) PGLOG['CURUID'] == PGLOG['SETUID'] # set CURUID to a specific specialist @@ -1262,12 +1277,12 @@ def set_common_pglog(): PGLOG['EUID'] = os.geteuid() PGLOG['CURUID'] = pwd.getpwuid(PGLOG['RUID']).pw_name try: - PGLOG['RDAUID'] = PGLOG['GDEXUID'] = pwd.getpwnam(PGLOG['GDEXUSER']).pw_uid + PGLOG['RDAUID'] = PGLOG['GDEXUID'] = pwd.getpwnam(PGLOG['COMMONUSER']).pw_uid PGLOG['RDAGID'] = PGLOG['GDEXGID'] = grp.getgrnam(PGLOG['GDEXGRP']).gr_gid except: PGLOG['RDAUID'] = PGLOG['GDEXUID'] = 0 PGLOG['RDAGID'] = PGLOG['GDEXGID'] = 0 - if PGLOG['CURUID'] == PGLOG['GDEXUSER']: PGLOG['SETUID'] = PGLOG['GDEXUSER'] + if PGLOG['CURUID'] == PGLOG['COMMONUSER']: PGLOG['SETUID'] = PGLOG['COMMONUSER'] PGLOG['HOSTNAME'] = get_host() for htype in HOSTTYPES: diff --git a/src/rda_python_common/PgOPT.py b/src/rda_python_common/PgOPT.py index 5794e1b..7b7e3cf 100644 --- a/src/rda_python_common/PgOPT.py +++ b/src/rda_python_common/PgOPT.py @@ -624,7 +624,7 @@ def set_email_logact(): def validate_dsowner(aname, dsid = None, logname = None, pgds = 0, logact = 0): if not logname: logname = (params['LN'] if 'LN' in params else PgLOG.PGLOG['CURUID']) - if logname == PgLOG.PGLOG['GDEXUSER']: return 1 + if logname == PgLOG.PGLOG['COMMONUSER']: return 1 dsids = {} if dsid: @@ -1638,7 +1638,7 @@ def send_request_email_notice(pgrqst, errmsg, fcount, rstat, readyfile = None, p exclude = (einfo['SENDER'] if errmsg else einfo['RECEIVER']) if not errmsg and pgcntl and pgcntl['ccemail']: PgLOG.add_carbon_copy(pgcntl['ccemail'], 1, exclude, pgrqst['specialist']) - if PgLOG.PGLOG['CURUID'] != pgrqst['specialist'] and PgLOG.PGLOG['CURUID'] != PgLOG.PGLOG['GDEXUSER']: + if PgLOG.PGLOG['CURUID'] != pgrqst['specialist'] and PgLOG.PGLOG['CURUID'] != PgLOG.PGLOG['COMMONUSER']: PgLOG.add_carbon_copy(PgLOG.PGLOG['CURUID'], 1, exclude) if 'CC' in params: PgLOG.add_carbon_copy(params['CC'], 0, exclude) einfo['CCD'] = PgLOG.PGLOG['CCDADDR'] diff --git a/src/rda_python_common/PgSIG.py b/src/rda_python_common/PgSIG.py index 4efe376..28ac6f0 100644 --- a/src/rda_python_common/PgSIG.py +++ b/src/rda_python_common/PgSIG.py @@ -1089,7 +1089,7 @@ def record_background(bcmd, logact = PgLOG.LOGWRN): aname = bcmd mp = r"^\s*(\S+)\s+(\d+)\s+1\s+.*{}(.*)$".format(aname) - pc = "ps -u {},{} -f | grep ' 1 ' | grep {}".format(PgLOG.PGLOG['CURUID'], PgLOG.PGLOG['GDEXUSER'], aname) + pc = "ps -u {},{} -f | grep ' 1 ' | grep {}".format(PgLOG.PGLOG['CURUID'], PgLOG.PGLOG['COMMONUSER'], aname) for i in range(2): buf = PgLOG.pgsystem(pc, logact, 20+1024) if buf: @@ -1100,7 +1100,7 @@ def record_background(bcmd, logact = PgLOG.LOGWRN): (uid, sbid, acmd) = ms.groups() bid = int(sbid) if bid in CBIDS: return -1 - if uid == PgLOG.PGLOG['GDEXUSER']: + if uid == PgLOG.PGLOG['COMMONUSER']: acmd = re.sub(r'^\.(pl|py)\s+', '', acmd, 1) if re.match(r'^{}{}'.format(aname, acmd), bcmd): continue CBIDS[bid] = bcmd diff --git a/src/rda_python_common/__init__.py b/src/rda_python_common/__init__.py index 8c0a536..361d2c9 100644 --- a/src/rda_python_common/__init__.py +++ b/src/rda_python_common/__init__.py @@ -22,7 +22,7 @@ from . import PgLOG, PgUtil, PgDBI, PgFile, PgLock, PgCMD, PgSIG, PgOPT, PgSplit -__version__ = "2.1.11" +__version__ = "3.0.0" __all__ = [ "PgLOG", diff --git a/src/rda_python_common/pg_dbi.py b/src/rda_python_common/pg_dbi.py index 096e612..4858e8d 100644 --- a/src/rda_python_common/pg_dbi.py +++ b/src/rda_python_common/pg_dbi.py @@ -13,24 +13,75 @@ import time import hvac from datetime import datetime -import psycopg2 as PgSQL -from psycopg2.extras import execute_values -from psycopg2.extras import execute_batch from os import path as op from .pg_log import PgLOG +# Driver selection: prefer psycopg (v3) as the default driver; fall back to +# psycopg2 when psycopg is not installed. Both drivers share enough surface +# (connect(**config), Error/OperationalError, cursor.execute/executemany/ +# fetchone/description, connection.commit/rollback/close/autocommit) for this +# module to use either transparently. +try: + import psycopg as PgSQL + PG_DRIVER = 'psycopg3' + + # psycopg3 removed extras.execute_values / extras.execute_batch. + # cursor.executemany() in psycopg3 is efficient by default (it uses prepared + # statements internally), so the shims below provide matching signatures. + + def execute_values(cursor, sql, argslist, page_size=100): + """Driver-neutral shim providing psycopg2.extras.execute_values() on psycopg3. + + Replaces the single ``VALUES %s`` placeholder in ``sql`` (as expected by + psycopg2's execute_values) with a per-row ``VALUES (%s, %s, ...)`` tuple + built from the first row, then calls ``cursor.executemany()``. Lets the + same call site work under either psycopg (v3) or psycopg2. + """ + if not argslist: return + ncol = len(argslist[0]) + row_ph = '(' + ','.join(['%s']*ncol) + ')' + new_sql = re.sub(r'(?i)\bVALUES\s+%s\b', 'VALUES ' + row_ph, sql, count=1) + cursor.executemany(new_sql, argslist) + + def execute_batch(cursor, sql, argslist, page_size=100): + """Driver-neutral shim providing psycopg2.extras.execute_batch() on psycopg3.""" + cursor.executemany(sql, argslist) + + def get_pgcode(pgerr): + """Return the 5-char SQLSTATE code from a psycopg3 error, or None.""" + diag = getattr(pgerr, 'diag', None) + return getattr(diag, 'sqlstate', None) if diag is not None else None + + def get_pgerror(pgerr): + """Return the primary error message from a psycopg3 error, or None.""" + diag = getattr(pgerr, 'diag', None) + return getattr(diag, 'message_primary', None) if diag is not None else None +except ImportError: + import psycopg2 as PgSQL + from psycopg2.extras import execute_values, execute_batch + PG_DRIVER = 'psycopg2' + + def get_pgcode(pgerr): + """Return the 5-char SQLSTATE code from a psycopg2 error, or None.""" + return getattr(pgerr, 'pgcode', None) + + def get_pgerror(pgerr): + """Return the server error message from a psycopg2 error, or None.""" + return getattr(pgerr, 'pgerror', None) + class PgDBI(PgLOG): """PostgreSQL Database Interface layer extending PgLOG. Provides a high-level API for connecting to and querying PostgreSQL databases - using psycopg2. Supports single and batch INSERT, SELECT, UPDATE, and DELETE - operations, transaction management, schema introspection, user lookups, usage - tracking, and credential retrieval from .pgpass or OpenBao. + using psycopg (v3) when available, falling back to psycopg2. Supports single + and batch INSERT, SELECT, UPDATE, and DELETE operations, transaction + management, schema introspection, user lookups, usage tracking, and + credential retrieval from .pgpass or OpenBao. Inherits all logging and utility helpers from PgLOG. Instance Attributes: - pgdb (connection | None): Active psycopg2 connection, or None when disconnected. + pgdb (connection | None): Active psycopg/psycopg2 connection, or None when disconnected. curtran (int): Transaction counter: 0 = idle, >0 = inside a transaction. NMISSES (list): Cached list of scientist IDs (userno) not found in the DB. LMISSES (list): Cached list of login names not found in the DB. @@ -40,7 +91,8 @@ class PgDBI(PgLOG): SYSDOWN (dict): Cache of system-down status records keyed by hostname. PGDBI (dict): Active connection and configuration parameters. PGSIGNS (list): Special comparison sign tokens recognised by get_field_condition(). - CHCODE (int): psycopg2 type code for CHAR columns (used to strip trailing spaces). + CHCODE (int): Driver type code for CHAR columns (used to strip trailing spaces); + set from psycopg/psycopg2 depending on which driver is in use. DBPORTS (dict): Mapping of database names to non-default TCP port numbers. DBPASS (dict): Credentials loaded from .pgpass, keyed by (host, port, db, user). DBBAOS (dict): Credentials loaded from OpenBao, keyed by database name. @@ -544,14 +596,15 @@ def valid_table(self, tname, pre = None, suf = None, logact = 0): return tbname def check_dberror(self, pgerr, pgcnt, sqlstr, ary, logact = None): - """Classify a psycopg2 error and decide whether to retry or abort. + """Classify a psycopg/psycopg2 error and decide whether to retry or abort. Handles connection errors (08xxx, 57xxx), lock errors (55xxx), aborted transactions (25P02), and missing-table errors (42P01 with ADDTBL flag). Retries up to PGLOG['DBRETRY'] times; exits after that threshold. Args: - pgerr (psycopg2.Error): The caught database exception. + pgerr (PgSQL.Error): The caught database exception + (psycopg.Error or psycopg2.Error). pgcnt (int): Current retry count (0-based). sqlstr (str): SQL statement that caused the error, for logging. ary: Bound values that were passed to the statement, for logging. @@ -562,8 +615,8 @@ def check_dberror(self, pgerr, pgcnt, sqlstr, ary, logact = None): """ if logact is None: logact = self.PGDBI['ERRLOG'] ret = self.FAILURE - pgcode = pgerr.pgcode - pgerror = pgerr.pgerror + pgcode = get_pgcode(pgerr) + pgerror = get_pgerror(pgerr) dberror = "{} {}".format(pgcode, pgerror) if pgcode and pgerror else str(pgerr) if pgcnt < self.PGLOG['DBRETRY']: if not pgcode: @@ -640,15 +693,15 @@ def pgconnect(self, reconnect = 0, pgcnt = 0, autocommit = True): autocommit (bool): Whether to enable autocommit on the new connection. Returns: - connection | int: psycopg2 connection on success, self.FAILURE on error. + connection | int: psycopg/psycopg2 connection on success, self.FAILURE on error. """ if self.pgdb: if reconnect and not self.pgdb.closed: return self.pgdb # no need reconnect elif reconnect: reconnect = 0 # initial connection while True: - config = {'database': self.PGDBI['DBNAME'], - 'user': self.PGDBI['LNNAME']} + config = {'dbname': self.PGDBI['DBNAME'], + 'user': self.PGDBI['LNNAME']} if self.PGDBI['DBSHOST'] == self.PGLOG['HOSTNAME']: config['host'] = 'localhost' else: @@ -656,7 +709,7 @@ def pgconnect(self, reconnect = 0, pgcnt = 0, autocommit = True): if not self.PGDBI['DBPORT']: self.PGDBI['DBPORT'] = self.get_dbport(self.PGDBI['DBNAME']) if self.PGDBI['DBPORT']: config['port'] = self.PGDBI['DBPORT'] config['password'] = '***' - sqlstr = "psycopg2.connect(**{})".format(config) + sqlstr = "{}.connect(**{})".format(PG_DRIVER, config) config['password'] = self.get_pgpass_password() if self.PGLOG['DBGLEVEL']: self.pgdbg(1000, sqlstr) try: @@ -675,7 +728,7 @@ def pgcursor(self): errors. The search path includes PGDBI['SCPATH'] when it differs from SCNAME. Returns: - cursor | int: psycopg2 cursor on success, self.FAILURE on error. + cursor | int: psycopg/psycopg2 cursor on success, self.FAILURE on error. """ pgcur = None if not self.pgdb: @@ -924,7 +977,8 @@ def pgmadd(self, tablename, records, logact = None, getid = None): """Insert multiple records into a database table efficiently. When getid is set, executes individual inserts to capture each returned ID. - Otherwise uses psycopg2 execute_values() for a single bulk INSERT. + Otherwise uses execute_values() (psycopg2's bulk helper, or the + executemany()-based shim on psycopg v3) for a single bulk INSERT. Args: tablename (str): Target table name. @@ -1365,8 +1419,9 @@ def pghupdt(self, tablename, record, cnddict, logact = None): def pgmupdt(self, tablename, records, cnddicts, logact = None): """Update multiple rows using parallel value and condition dicts. - Uses psycopg2 execute_batch() for efficient bulk updates. The number of - values in records and cnddicts must match. + Uses execute_batch() (psycopg2's bulk helper, or the executemany()-based + shim on psycopg v3) for efficient bulk updates. The number of values in + records and cnddicts must match. Args: tablename (str): Target table name. @@ -1508,7 +1563,8 @@ def pghdel(self, tablename, cnddict, logact = None): def pgmdel(self, tablename, cnddicts, logact = None): """Delete multiple rows using a multi-value condition dict. - Uses psycopg2 execute_batch() for efficient bulk deletes. + Uses execute_batch() (psycopg2's bulk helper, or the executemany()-based + shim on psycopg v3) for efficient bulk deletes. Args: tablename (str): Target table name. diff --git a/src/rda_python_common/pg_file.py b/src/rda_python_common/pg_file.py index 068d8ef..3e8ec5b 100644 --- a/src/rda_python_common/pg_file.py +++ b/src/rda_python_common/pg_file.py @@ -987,7 +987,7 @@ def delete_backup_file(self, file, endpoint = None, logact = 0): return self.FINISH return self.FAILURE - # reset local file/directory information to make them writable for self.PGLOG['GDEXUSER'] + # reset local file/directory information to make them writable for self.PGLOG['COMMONUSER'] # file - file name (mandatory) # info - gathered file info with option 14, None means file not exists def reset_local_info(self, file, info = None, logact = 0): diff --git a/src/rda_python_common/pg_log.py b/src/rda_python_common/pg_log.py index ffc916f..bb23ea2 100644 --- a/src/rda_python_common/pg_log.py +++ b/src/rda_python_common/pg_log.py @@ -81,7 +81,7 @@ class PgLOG: EMLSUM = (0x08000) # record as email summary EMEROL = (0x10000) # record error as email only EMLALL = (0x1D208) # all email acts - DOSUDO = (0x20000) # add 'sudo -u self.PGLOG['GDEXUSER']' + DOSUDO = (0x20000) # add 'sudo -u self.PGLOG['COMMONUSER']' NOTLOG = (0x40000) # do not log any thing OVRIDE = (0x80000) # do override existing file or record NOWAIT = (0x100000) # do not wait on globus task to finish @@ -124,9 +124,8 @@ def __init__(self): 'SETUID': '', # the login name for suid if it is different to the CURUID 'FILEMODE': 0o664, # default 8-base file mode 'EXECMODE': 0o775, # default 8-base executable file mode or directory mode - 'GDEXUSER': "gdexdata", # common gdex user name - 'GDEXEMAIL': "zji", # specialist to receipt email intead of common gdex user name - 'SUDOGDEX': 0, # 1 to allow sudo to self.PGLOG['GDEXUSER'] + # COMMONUSER and ADMINUSER are set below via SETPGLOG (env overrides PG) + 'SUDOGDEX': 0, # 1 to allow sudo to self.PGLOG['COMMONUSER'] 'HOSTNAME': '', # current host name the process in running on 'OBJCTSTR': "object", 'BACKUPNM': "quasar", @@ -161,10 +160,6 @@ def __init__(self): 'EMLSRVR': "ndir.ucar.edu", # UCAR email server and port 'EMLPORT': 25 } - self.PGLOG['RDAUSER'] = self.PGLOG['GDEXUSER'] - self.PGLOG['RDAGRP'] = self.PGLOG['GDEXGRP'] - self.PGLOG['RDAEMAIL'] = self.PGLOG['GDEXEMAIL'] - self.PGLOG['SUDORDA'] = self.PGLOG['SUDOGDEX'] self.HOSTTYPES = { 'rda': 'dsg_mach', 'crlogin': 'dav', @@ -188,6 +183,13 @@ def __init__(self): self.set_common_pglog() self.OUTPUT = None + def SETPGLOG(self, key, default): + """Set ``self.PGLOG[key]`` from environment variable ``PG`` or + fall back to ``default`` if the variable is unset. Used to make + per-environment overrides (e.g. PGCOMMONUSER, PGADMINUSER) survive + package upgrades.""" + self.PGLOG[key] = os.environ.get('PG' + key, default) + def open_output(self, outfile=None): """Open the result output destination. @@ -392,14 +394,14 @@ def send_python_email(self, subject=None, receiver=None, msg=None, sender=None, docc = False if cc else True if not sender: sender = self.PGLOG['CURUID'] - if sender != self.PGLOG['GDEXUSER']: docc = False - if sender == self.PGLOG['GDEXUSER']: sender = self.PGLOG['GDEXEMAIL'] + if sender != self.PGLOG['COMMONUSER']: docc = False + if sender == self.PGLOG['COMMONUSER']: sender = self.PGLOG['ADMINUSER'] if sender.find('@') == -1: sender += "@ucar.edu" if not receiver: receiver = self.PGLOG['EMLADDR'] if self.PGLOG['EMLADDR'] else self.PGLOG['CURUID'] - if receiver == self.PGLOG['GDEXUSER']: receiver = self.PGLOG['GDEXEMAIL'] + if receiver == self.PGLOG['COMMONUSER']: receiver = self.PGLOG['ADMINUSER'] if receiver.find('@') == -1: receiver += "@ucar.edu" - if docc and not re.match(self.PGLOG['GDEXUSER'], sender): self.add_carbon_copy(sender, 1) + if docc and not re.match(self.PGLOG['COMMONUSER'], sender): self.add_carbon_copy(sender, 1) emlmsg = EmailMessage() emlmsg.set_content(msg) emlmsg['From'] = sender @@ -1324,10 +1326,17 @@ def reset_batch_host(self, bhost, logact=None): self.PGLOG['PGBATCH'] = '' self.PGLOG['CURBID'] = 0 - @staticmethod - def get_command(cmdstr=None): + def get_command(self, cmdstr=None): """Return the base command name, stripping directory and ``.py``/``.pl`` extension. + When invoked via the pywrapper setuid C wrapper, ``sys.argv[0]`` is the + resolved Python script path (e.g. ``/.../setuid_rdacp``) rather than the + alias the user typed (e.g. ``rdacp``); the kernel discards argv[0] when + handling the script's shebang. When the basename starts with + ``setuid_`` and the effective user equals ``self.PGLOG['COMMONUSER']``, + this process was started via pywrapper, so the ``setuid_`` prefix is + stripped to recover the logical command name. + Args: cmdstr: Path string. Defaults to ``sys.argv[0]``. @@ -1336,6 +1345,10 @@ def get_command(cmdstr=None): """ if not cmdstr: cmdstr = sys.argv[0] cmdstr = op.basename(cmdstr) + if cmdstr.startswith('setuid_'): + euser = pwd.getpwuid(os.geteuid()).pw_name + if euser == self.PGLOG['COMMONUSER']: + cmdstr = cmdstr[len('setuid_'):] ms = re.match(r'^(.+)\.(py|pl)$', cmdstr) if ms: return ms.group(1) @@ -1358,11 +1371,11 @@ def get_local_command(self, cmd, asuser=None): """ cuser = self.PGLOG['SETUID'] if self.PGLOG['SETUID'] else self.PGLOG['CURUID'] if not asuser or cuser == asuser: return cmd - if cuser == self.PGLOG['GDEXUSER']: + if cuser == self.PGLOG['COMMONUSER']: wrapper = "pgstart_" + asuser if self.valid_command(wrapper): return "{} {}".format(wrapper, cmd) - elif self.PGLOG['SUDOGDEX'] and asuser == self.PGLOG['GDEXUSER']: - return "sudo -u {} {}".format(self.PGLOG['GDEXUSER'], cmd) # sudo as user gdexdata + elif self.PGLOG['SUDOGDEX'] and asuser == self.PGLOG['COMMONUSER']: + return "sudo -u {} {}".format(self.PGLOG['COMMONUSER'], cmd) # sudo as user gdexdata return cmd def get_remote_command(self, cmd, host, asuser=None): @@ -1389,8 +1402,8 @@ def get_sync_command(self, host, asuser=None): Sync command string (e.g. ``"synccasper"`` or ``"casper-sync"``). """ host = self.get_short_host(host) - if (not (self.PGLOG['SETUID'] and self.PGLOG['SETUID'] == self.PGLOG['GDEXUSER']) and - (not asuser or asuser == self.PGLOG['GDEXUSER'])): + if (not (self.PGLOG['SETUID'] and self.PGLOG['SETUID'] == self.PGLOG['COMMONUSER']) and + (not asuser or asuser == self.PGLOG['COMMONUSER'])): return "sync" + host return host + "-sync" @@ -1407,7 +1420,7 @@ def set_suid(self, cuid=0): if cuid != self.PGLOG['EUID'] or cuid != self.PGLOG['RUID']: os.setreuid(cuid, cuid) self.PGLOG['SETUID'] = pwd.getpwuid(cuid).pw_name - if not (self.PGLOG['SETUID'] == self.PGLOG['GDEXUSER'] or cuid == self.PGLOG['RUID']): + if not (self.PGLOG['SETUID'] == self.PGLOG['COMMONUSER'] or cuid == self.PGLOG['RUID']): self.set_specialist_environments(self.PGLOG['SETUID']) self.PGLOG['CURUID'] == self.PGLOG['SETUID'] # set CURUID to a specific specialist @@ -1421,18 +1434,28 @@ def set_common_pglog(self): Called automatically by :meth:`__init__`. """ - self.PGLOG['CURDIR'] = os.getcwd() + # resolve common/admin user from environment (PGCOMMONUSER / PGADMINUSER) + self.SETPGLOG("COMMONUSER", "gdexdata") + self.SETPGLOG("ADMINUSER", "zji") + self.PGLOG['RDAUSER'] = self.PGLOG['COMMONUSER'] + self.PGLOG['RDAGRP'] = self.PGLOG['GDEXGRP'] + self.PGLOG['RDAEMAIL'] = self.PGLOG['ADMINUSER'] + self.PGLOG['SUDORDA'] = self.PGLOG['SUDOGDEX'] + # backwards-compat aliases (deprecated: use COMMONUSER / ADMINUSER) + self.PGLOG['GDEXUSER'] = self.PGLOG['COMMONUSER'] + self.PGLOG['GDEXEMAIL'] = self.PGLOG['ADMINUSER'] + self.PGLOG['CURDIR'] = os.getcwd() # set current user id self.PGLOG['RUID'] = os.getuid() self.PGLOG['EUID'] = os.geteuid() self.PGLOG['CURUID'] = pwd.getpwuid(self.PGLOG['RUID']).pw_name try: - self.PGLOG['RDAUID'] = self.PGLOG['GDEXUID'] = pwd.getpwnam(self.PGLOG['GDEXUSER']).pw_uid + self.PGLOG['RDAUID'] = self.PGLOG['GDEXUID'] = pwd.getpwnam(self.PGLOG['COMMONUSER']).pw_uid self.PGLOG['RDAGID'] = self.PGLOG['GDEXGID'] = grp.getgrnam(self.PGLOG['GDEXGRP']).gr_gid except KeyError: self.PGLOG['RDAUID'] = self.PGLOG['GDEXUID'] = 0 self.PGLOG['RDAGID'] = self.PGLOG['GDEXGID'] = 0 - if self.PGLOG['CURUID'] == self.PGLOG['GDEXUSER']: self.PGLOG['SETUID'] = self.PGLOG['GDEXUSER'] + if self.PGLOG['CURUID'] == self.PGLOG['COMMONUSER']: self.PGLOG['SETUID'] = self.PGLOG['COMMONUSER'] self.PGLOG['HOSTNAME'] = self.get_host() for htype in self.HOSTTYPES: ms = re.match(r'^{}(-|\d|$)'.format(htype), self.PGLOG['HOSTNAME']) diff --git a/src/rda_python_common/pg_opt.py b/src/rda_python_common/pg_opt.py index 6e5ffa6..bf41eab 100644 --- a/src/rda_python_common/pg_opt.py +++ b/src/rda_python_common/pg_opt.py @@ -649,7 +649,7 @@ def validate_dsowner(self, aname, dsid=None, logname=None, pgds=0, logact=0): dataset in params), or logs a fatal error and exits. """ if not logname: logname = (self.params['LN'] if 'LN' in self.params else self.PGLOG['CURUID']) - if logname == self.PGLOG['GDEXUSER']: return 1 + if logname == self.PGLOG['COMMONUSER']: return 1 dsids = {} if dsid: dsids[dsid] = 1 @@ -1847,7 +1847,7 @@ def send_request_email_notice(self, pgrqst, errmsg, fcount, rstat, readyfile=Non exclude = (einfo['SENDER'] if errmsg else einfo['RECEIVER']) if not errmsg and pgcntl and pgcntl['ccemail']: self.add_carbon_copy(pgcntl['ccemail'], 1, exclude, pgrqst['specialist']) - if self.PGLOG['CURUID'] != pgrqst['specialist'] and self.PGLOG['CURUID'] != self.PGLOG['GDEXUSER']: + if self.PGLOG['CURUID'] != pgrqst['specialist'] and self.PGLOG['CURUID'] != self.PGLOG['COMMONUSER']: self.add_carbon_copy(self.PGLOG['CURUID'], 1, exclude) if 'CC' in self.params: self.add_carbon_copy(self.params['CC'], 0, exclude) einfo['CCD'] = self.PGLOG['CCDADDR'] diff --git a/src/rda_python_common/pg_sig.py b/src/rda_python_common/pg_sig.py index d84eef1..3435f08 100644 --- a/src/rda_python_common/pg_sig.py +++ b/src/rda_python_common/pg_sig.py @@ -1183,14 +1183,14 @@ def record_background(self, bcmd, logact=None): ms = re.match(r'^(\S+)', bcmd) aname = ms.group(1) if ms else bcmd curuid = self.PGLOG['CURUID'] - gdexuser = self.PGLOG['GDEXUSER'] + commonuser = self.PGLOG['COMMONUSER'] for i in range(2): for proc in psutil.process_iter(['pid', 'ppid', 'username', 'cmdline']): try: info = proc.info if info.get('ppid') != 1: continue uid = info.get('username') - if uid != curuid and uid != gdexuser: continue + if uid != curuid and uid != commonuser: continue cmdline = info.get('cmdline') or [] if not cmdline: continue line = ' '.join(cmdline) @@ -1199,7 +1199,7 @@ def record_background(self, bcmd, logact=None): bid = info['pid'] if bid in self.CBIDS: return -1 acmd = line[idx+len(aname):] - if uid == gdexuser: + if uid == commonuser: acmd = re.sub(r'^\.(pl|py)\s+', '', acmd, 1) if re.match(r'^{}{}'.format(aname, acmd), bcmd): continue self.CBIDS[bid] = bcmd diff --git a/src/rda_python_common/pg_util.py b/src/rda_python_common/pg_util.py index 8586e31..a449664 100644 --- a/src/rda_python_common/pg_util.py +++ b/src/rda_python_common/pg_util.py @@ -14,6 +14,8 @@ import datetime import calendar import glob +import bisect +import functools from os import path as op from .pg_log import PgLOG @@ -419,11 +421,7 @@ def split_datetime(sdt, sep = r'\D'): list: Mixed int/str parts of the split datetime. """ if not isinstance(sdt, str): sdt = str(sdt) - adt = re.split(sep, sdt) - acnt = len(adt) - for i in range(acnt): - if adt[i].isdigit(): adt[i] = int(adt[i]) - return adt + return [int(x) if x.isdigit() else x for x in re.split(sep, sdt)] # date: given date in format of fromfmt # tofmt: date formats; ex. "Month D, YYYY" @@ -547,14 +545,8 @@ def fmtdatetime(self, yr, mn, dy, hr = None, nn = None, ss = None, tofmt = None) # adjust second/minute/hour values out of range for i in range(3): if tms[i] != None and tms[i+1] != None: - if tms[i] < 0: - while tms[i] < 0: - tms[i] += ups[i] - tms[i+1] -= 1 - elif tms[i] >= ups[i]: - while tms[i] >= ups[i]: - tms[i] -= ups[i] - tms[i+1] += 1 + carry, tms[i] = divmod(tms[i], ups[i]) + tms[i+1] += carry sdt = self.fmtdate(yr, mn, dy, tofmt) # format second/minute/hour values for i in range(3): @@ -592,14 +584,8 @@ def fmtdatehour(self, yr, mn, dy, hr, tofmt = None): """ if not tofmt: tofmt = "YYYY-MM-DD:HH" if hr != None and dy != None: # adjust hour value out of range - if hr < 0: - while hr < 0: - hr += 24 - dy -= 1 - elif hr > 23: - while hr > 23: - hr -= 24 - dy += 1 + carry, hr = divmod(hr, 24) + dy += carry datehour = self.fmtdate(yr, mn, dy, tofmt) if hr != None: ms = re.search(self.DATEFMTS['H'], datehour, re.I) @@ -733,9 +719,8 @@ def date_and_time(sdt): """ if not sdt: return [None, None] if not isinstance(sdt, str): sdt = str(sdt) - adt = re.split(' ', sdt) - acnt = len(adt) - if acnt == 1: adt.append('00:00:00') + adt = sdt.split(' ') + if len(adt) == 1: adt.append('00:00:00') return adt # convert given date/time to unix epoch time; -1 if cannot @@ -1209,10 +1194,8 @@ def joinarray(lst1, lst2, unique = None): cnt2 = len(lst2) if unique: for i in range(cnt2): - for j in range(cnt1): - if PgUtil.pgcmp(lst1[j], lst2[i]) != 0: break - if j >= cnt1: - lst1.append(lst2[i]) + if lst2[i] not in lst1: + lst1.append(lst2[i]) else: lst1.extend(lst2) return lst1 @@ -1312,7 +1295,8 @@ def sorthash(self, pgrecs, flds, hash, patterns = None): rec.append(val) rec.append(i) # add column to cache the row index srecs.append(rec) - srecs = self.quicksort(srecs, 0, count-1, desc, fcnt, nums) + srecs.sort(key=functools.cmp_to_key( + lambda a, b: PgUtil.cmp_records(a, b, desc, fcnt, nums))) # sort pgrecs according the cached row index column in ordered srecs rets = {} for fld in pgrecs: @@ -1335,10 +1319,13 @@ def diffdate(date1, date2): Returns: int: Positive when date1 > date2, negative when date1 < date2. """ - ut1 = ut2 = 0 - if date1: ut1 = PgUtil.unixtime(date1) - if date2: ut2 = PgUtil.unixtime(date2) - return round((ut1 - ut2)/86400) # 24*60*60 + epoch = datetime.date(1970, 1, 1) + def _to_date(d): + if not d: return epoch + ms = re.match(r'^(\d+)-(\d+)-(\d+)', str(d)) + if not ms: return epoch + return datetime.date(int(ms.group(1)), int(ms.group(2)), int(ms.group(3))) + return (_to_date(date1) - _to_date(date2)).days # Return: the number of seconds bewteen time1 and time2 @staticmethod @@ -1557,15 +1544,10 @@ def addyearmonth(ym, yr, mn): if ms: (syr, smn) = ms.groups() nyr = int(syr) + yr - nmn = int(smn) + mn - if nmn < 0: - while nmn < 0: - nyr -= 1 - nmn += 12 - else: - while nmn > 12: - nyr += 1 - nmn -= 12 + nmn = int(smn) + mn - 1 # shift to 0-indexed for divmod + extra, nmn = divmod(nmn, 12) + nyr += extra + nmn += 1 # back to 1-indexed ym = "{:04}{:02}".format(nyr, nmn) return ym @@ -1729,14 +1711,7 @@ def addhour(self, sdate, stime, nhour): if ms: shr = ms.group(1) hr = int(shr) + nhour - if hr < 0: - while hr < 0: - dy -= 1 - hr += 24 - else: - while hr > 23: - dy += 1 - hr -= 24 + dy, hr = divmod(hr, 24) shour = "{:02}".format(hr) if shr != shour: stime = re.sub(shr, shour, stime, 1) if dy: sdate = self.adddate(sdate, 0, 0, dy) @@ -1765,14 +1740,8 @@ def adddatehour(self, sdate, nhour, yr, mn, dy, hr = 0): if nhour != None: if isinstance(nhour, str): nhour = int(nhour) hr += nhour - if hr < 0: - while hr < 0: - dy -= 1 - hr += 24 - else: - while hr > 23: - dy += 1 - hr -= 24 + carry, hr = divmod(hr, 24) + dy += carry if nhour != None: nhour = hr if yr or mn or dy: sdate = self.adddate(sdate, yr, mn, dy) return [sdate, nhour] @@ -1799,7 +1768,7 @@ def adddatetime(self, sdatetime, yy, mm, dd, hh, nn, ss, nf = 0): str: Resulting datetime string in 'YYYY-MM-DD HH:MM:SS' format. """ if sdatetime and not isinstance(sdatetime, str): sdatetime = str(sdatetime) - (sdate, stime) = re.split(' ', sdatetime) + (sdate, stime) = sdatetime.split(' ', 1) if hh or nn or ss: (sdate, stime) = self.addtime(sdate, stime, hh, nn, ss) if nf: sdate = self.addmonth(sdate, mm, nf) @@ -1837,14 +1806,8 @@ def addtime(self, sdate, stime, h, m, s): tms[1] += int(ms.group(2)) tms[0] += int(ms.group(3)) for i in range(3): - if tms[i] < 0: - while tms[i] < 0: - tms[i] += ups[i] - tms[i+1] -= 1 - elif tms[i] >= ups[i]: - while tms[i] >= ups[i]: - tms[i] -= ups[i] - tms[i+1] += 1 + carry, tms[i] = divmod(tms[i], ups[i]) + tms[i+1] += carry stime = "{:02}:{:02}:{:02}".format(tms[2], tms[1], tms[0]) if tms[3]: sdate = self.adddate(sdate, 0, 0, tms[3]) return [sdate, stime] @@ -1985,7 +1948,7 @@ def enddatetime(self, sdatetime, unit, days = 0, nf = 0): """ if sdatetime and not isinstance(sdatetime, str): sdatetime = str(sdatetime) if not (unit and unit in 'YMWDHNS'): return sdatetime - (sdate, stime) = re.split(' ', sdatetime) + (sdate, stime) = sdatetime.split(' ', 1) if unit in 'HNS': stime = self.endtime(stime, unit) else: @@ -2011,7 +1974,7 @@ def get_column_length(colname, values): for val in values: if val is None: continue sval = str(val) - if sval and not re.search(r'\n', sval): + if sval and '\n' not in sval: slen = len(sval) if slen > clen: clen = slen return clen @@ -2199,21 +2162,8 @@ def asearch(lidx, hidx, key, list): Returns: int: Index of the matching element, or -1 when not found. """ - ret = -1 - if (hidx - lidx) < 11: # use linear search for less than 11 items - for midx in range(lidx, hidx): - if key == list[midx]: - ret = midx - break - else: - midx = (lidx + hidx) // 2 - if key == list[midx]: - ret = midx - elif key < list[midx]: - ret = PgUtil.asearch(lidx, midx, key, list) - else: - ret = PgUtil.asearch(midx + 1, hidx, key, list) - return ret + idx = bisect.bisect_left(list, key, lidx, hidx) + return idx if idx < hidx and list[idx] == key else -1 # lidx: lower index limit (including) # hidx: higher index limit (excluding) @@ -2362,14 +2312,11 @@ def is_text_file(fname, blocksize = 256, threshhold = 0.1): buffer = f.read(blocksize) # Check for null bytes (a strong indicator of a binary file) if not buffer or b'\0' in buffer: return 0 - text_characters = ( + text_set = frozenset( b'\t\n\r\f\v' + # Whitespace characters bytes(range(32, 127)) # Printable ASCII characters ) - non_text_count = 0 - for byte in buffer: - if byte not in text_characters: - non_text_count += 1 # Count non-text characters + non_text_count = sum(b not in text_set for b in buffer) # If a significant portion of the buffer consists of non-text characters, # it's likely a binary file. return 1 if((non_text_count/len(buffer)) < threshhold) else 0