4119519683

Build # 4119519683

Build Type

push

github

Committed by GitHub

Commit Message

Lintify (#287)

Run Details

36 of 36 new or added lines in 9 files covered. (100.0%)

5321 of 6332 relevant lines covered (84.03%)

1.68 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

74.07

/winterdrp/processors/database/postgres.py

"""
Module containing postgres util functions
"""
# pylint: disable=not-context-manager
import logging
import os
from glob import glob
from pathlib import Path
from typing import Optional

import numpy as np
import psycopg
from psycopg import errors
from psycopg.rows import Row

from winterdrp.data import DataBlock
from winterdrp.errors import ProcessorError
from winterdrp.processors.database.constraints import DBQueryConstraints
from winterdrp.processors.database.utils import get_ordered_schema_list

logger = logging.getLogger(__name__)

DB_USER_KEY = "DB_USER"
DB_PASSWORD_KEY = "DB_PWD"

PG_ADMIN_USER_KEY = "PG_ADMIN_USER"
PG_ADMIN_PWD_KEY = "PG_ADMIN_PWD"

DB_USER = os.getenv(DB_USER_KEY)
DB_PASSWORD = os.getenv(DB_PASSWORD_KEY)

ADMIN_USER = os.getenv(PG_ADMIN_USER_KEY, DB_USER)
ADMIN_PASSWORD = os.getenv(PG_ADMIN_PWD_KEY, DB_PASSWORD)

POSTGRES_DUPLICATE_PROTOCOLS = ["fail", "ignore", "replace"]


class DataBaseError(ProcessorError):
    """Error relating to postgres interactions"""


class PostgresUser:
    """
    Basic Postgres user class for executing functions
    """

    user_env_varaiable = DB_USER_KEY
    pass_env_variable = DB_PASSWORD_KEY

    def __init__(self, db_user: str = DB_USER, db_password: str = DB_PASSWORD):
        self.db_user = db_user
        self.db_password = db_password

    def validate_credentials(self):
        """
        Checks that user credentials exist
        :return: None
        """
        if self.db_user is None:
            err = (
                f"'db_user' is set as None. Please pass a db_user as an argument, "
                f"or set the environment variable '{self.user_env_varaiable}'."
            )
            logger.error(err)
            raise DataBaseError(err)

        if self.db_password is None:
            err = (
                f"'db_password' is set as None. Please pass a password as an argument, "
                f"or set the environment variable '{self.pass_env_variable}'."
            )
            logger.error(err)
            raise DataBaseError(err)

        # TODO check user exists

    def run_sql_command_from_file(self, file_path: str | Path, db_name: str):
        """
        Execute SQL command from file

        :param file_path: File to execute
        :param db_name: name of database
        :return: False
        """
        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            with open(file_path, "r", encoding="utf8") as sql_file:
                conn.execute(sql_file.read())

            logger.info(f"Executed sql commands from file {file_path}")

    def create_table(self, schema_path: str | Path, db_name: str):
        """
        Create a database table

        :param schema_path: File to execute
        :param db_name: name of database
        :return: None
        """
        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            with open(schema_path, "r", encoding="utf8") as schema_file:
                conn.execute(schema_file.read())

        logger.info(f"Created table from schema path {schema_path}")

    def create_tables_from_schema(
        self,
        schema_dir: str | Path,
        db_name: str,
    ):
        """
        Creates a db with tables, as described by .sql files in a directory

        :param schema_dir: Directory containing schema files
        :param db_name: name of DB
        :return: None
        """
        schema_files = glob(f"{schema_dir}/*.sql")
        ordered_schema_files = get_ordered_schema_list(schema_files)
        logger.info(f"Creating the following tables - {ordered_schema_files}")
        for schema_file in ordered_schema_files:
            self.create_table(schema_path=schema_file, db_name=db_name)

    def export_to_db(
        self,
        value_dict: dict | DataBlock,
        db_name: str,
        db_table: str,
        duplicate_protocol: str = "fail",
    ) -> tuple[list, list]:
        """
        Export a list of fields in value dict to a batabase table

        :param value_dict: dictionary/DataBlock/other dictonary-like object to export
        :param db_name: name of db to export to
        :param db_table: table of DB to export to
        :param duplicate_protocol: protocol for handling duplicates,
            in "fail"/"ignore"/"replace"
        :return:
        """

        assert duplicate_protocol in POSTGRES_DUPLICATE_PROTOCOLS

        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True

            sql_query = f"""
            SELECT Col.Column_Name from
                INFORMATION_SCHEMA.TABLE_CONSTRAINTS Tab,
                INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE Col
            WHERE
                Col.Constraint_Name = Tab.Constraint_Name
                AND Col.Table_Name = Tab.Table_Name
                AND Constraint_Type = 'PRIMARY KEY'
                AND Col.Table_Name = '{db_table}'
            """
            serial_keys, serial_key_values = [], []
            with conn.execute(sql_query) as cursor:
                primary_key = [x[0] for x in cursor.fetchall()]
                serial_keys = list(self.get_sequence_keys_from_table(db_table, db_name))
                logger.debug(serial_keys)
                colnames = [
                    desc[0]
                    for desc in conn.execute(
                        f"SELECT * FROM {db_table} LIMIT 1"
                    ).description
                    if desc[0] not in serial_keys
                ]

                colnames_str = ""
                for column in colnames:
                    colnames_str += f'"{column}",'
                colnames_str = colnames_str[:-1]
                txt = f"INSERT INTO {db_table} ({colnames_str}) VALUES ("

                for char in ["[", "]", "'"]:
                    txt = txt.replace(char, "")

                for column in colnames:
                    txt += f"'{str(value_dict[column])}', "

                txt = txt + ") "
                txt = txt.replace(", )", ")")

                if len(serial_keys) > 0:
                    txt += "RETURNING "
                    for key in serial_keys:
                        txt += f"{key},"
                    txt += ";"
                    txt = txt.replace(",;", ";")

                logger.debug(txt)
                command = txt

                try:
                    cursor.execute(command)
                    if len(serial_keys) > 0:
                        serial_key_values = cursor.fetchall()[0]
                    else:
                        serial_key_values = []

                except errors.UniqueViolation as exc:
                    primary_key_values = [value_dict[x] for x in primary_key]

                    if duplicate_protocol == "fail":
                        err = (
                            f"Duplicate error, entry with "
                            f"{primary_key}={primary_key_values} "
                            f"already exists in {db_name}."
                        )
                        logger.error(err)
                        raise errors.UniqueViolation from exc

                    if duplicate_protocol == "ignore":
                        logger.debug(
                            f"Found duplicate entry with "
                            f"{primary_key}={primary_key_values} in {db_name}. "
                            f"Ignoring."
                        )
                    elif duplicate_protocol == "replace":
                        logger.debug(
                            f"Updating duplicate entry with "
                            f"{primary_key}={primary_key_values} in {db_name}."
                        )

                        db_constraints = DBQueryConstraints(
                            columns=primary_key,
                            accepted_values=primary_key_values,
                        )

                        update_colnames = []
                        for column in colnames:
                            if column not in primary_key:
                                update_colnames.append(column)

                        serial_key_values = self.modify_db_entry(
                            db_constraints=db_constraints,
                            value_dict=value_dict,
                            db_alter_columns=update_colnames,
                            db_table=db_table,
                            db_name=db_name,
                            return_columns=serial_keys,
                        )

        return serial_keys, serial_key_values

    def modify_db_entry(
        self,
        db_name: str,
        db_table: str,
        db_constraints: DBQueryConstraints,
        value_dict: dict | DataBlock,
        db_alter_columns: str | list[str],
        return_columns: Optional[str | list[str]] = None,
    ) -> list[Row]:
        """
        Modify a db entry

        :param db_name: name of db
        :param db_table: Name of table
        :param value_dict: dict-like object to provide updated values
        :param db_alter_columns: columns to alter in db
        :param return_columns: columns to return
        :return: db query (return columns)
        """

        if not isinstance(db_alter_columns, list):
            db_alter_columns = [db_alter_columns]

        if return_columns is None:
            return_columns = db_alter_columns
        if not isinstance(return_columns, list):
            return_columns = [return_columns]

        constraints = db_constraints.parse_constraints()

        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True

            db_alter_values = [str(value_dict[c]) for c in db_alter_columns]

            alter_values_txt = [
                f"{db_alter_columns[ind]}='{db_alter_values[ind]}'"
                for ind in range(len(db_alter_columns))
            ]

            sql_query = (
                f"UPDATE {db_table} SET {', '.join(alter_values_txt)} "
                f"WHERE {constraints}"
            )

            if len(return_columns) > 0:
                logger.debug(return_columns)
                sql_query += f""" RETURNING {', '.join(return_columns)}"""
            sql_query += ";"
            query_output = self.execute_query(sql_query, db_name)

        return query_output

    def get_sequence_keys_from_table(self, db_table: str, db_name: str) -> np.ndarray:
        """
        Gets sequence keys of db table

        :param db_table: database table to use
        :param db_name: dataname name
        :return: numpy array of keys
        """
        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            sequences = [
                x[0]
                for x in conn.execute(
                    "SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';"
                ).fetchall()
            ]
            seq_tables = np.array([x.split("_")[0] for x in sequences])
            seq_columns = np.array([x.split("_")[1] for x in sequences])
            table_sequence_keys = seq_columns[(seq_tables == db_table)]
        return table_sequence_keys

    def import_from_db(
        self,
        db_name: str,
        db_table: str,
        db_output_columns: str | list[str],
        output_alias_map: Optional[str | list[str]] = None,
        max_num_results: Optional[int] = None,
        db_constraints: Optional[DBQueryConstraints] = None,
    ) -> list[dict]:
        """Query an SQL database with constraints, and return a list of dictionaries.
        One dictionary per entry returned from the query.

        #TODO check admin

        Parameters
        ----------
        db_name: Name of database to query
        db_table: Name of database table to query
        db_output_columns: Name(s) of columns to return for matched database entries
        output_alias_map: Alias to assign for each output column
        max_num_results: Maximum number of results to return

        Returns
        -------
        A list of dictionaries (one per entry)
        """

        if not isinstance(db_output_columns, list):
            db_output_columns = [db_output_columns]

        if output_alias_map is None:
            output_alias_map = db_output_columns

        if not isinstance(output_alias_map, list):
            output_alias_map = [output_alias_map]

        assert len(output_alias_map) == len(db_output_columns)

        all_query_res = []

        if db_constraints is not None:
            constraints = db_constraints.parse_constraints()
        else:
            constraints = ""

        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            sql_query = f"""
            SELECT {', '.join(db_output_columns)} from {db_table}
                WHERE {constraints}
            """

            if max_num_results is not None:
                sql_query += f" LIMIT {max_num_results}"

            sql_query += ";"

            logger.debug(f"Query: {sql_query}")

            with conn.execute(sql_query) as cursor:
                query_output = cursor.fetchall()

            for entry in query_output:
                assert len(entry) == len(db_output_columns)

                query_res = {}

                for i, key in enumerate(output_alias_map):
                    query_res[key] = entry[i]

                all_query_res.append(query_res)

        return all_query_res

    def execute_query(self, sql_query: str, db_name: str) -> list[Row]:
        """
        Generically execute SQL query

        :param sql_query: SQL query to execute
        :param db_name: db name
        :return: rows from db
        """
        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            logger.debug(f"Query: {sql_query}")

            with conn.execute(sql_query) as cursor:
                query_output = cursor.fetchall()

        return query_output

    def crossmatch_with_database(
        self,
        db_name: str,
        db_table: str,
        db_output_columns: str | list[str],
        ra: float,
        dec: float,
        crossmatch_radius_arcsec: float,
        output_alias_map: Optional[dict] = None,
        ra_field_name: str = "ra",
        dec_field_name: str = "dec",
        query_distance_bool: bool = False,
        q3c_bool: bool = False,
        query_constraints: Optional[DBQueryConstraints] = None,
        order_field_name: Optional[str] = None,
        num_limit: Optional[int] = None,
    ) -> list[dict]:
        """
        Crossmatch a given spatial position (ra/dec) with sources in a database,
        and returns a list of matches

        #TODO: check admin

        :param db_name: name of db to query
        :param db_table: name of db table
        :param db_output_columns: columns to return
        :param output_alias_map: mapping for renaming columns
        :param ra: RA
        :param dec: dec
        :param crossmatch_radius_arcsec: radius for crossmatch
        :param ra_field_name: name of ra column in database
        :param dec_field_name: name of dec column in database
        :param query_distance_bool: boolean where to return crossmatch distance
        :param q3c_bool: boolean whether to use q3c_bool
        :param order_field_name: field to order result by
        :param num_limit: limit on sql query
        :return: list of query result dictionaries
        """

        if output_alias_map is None:
            output_alias_map = {}
            for col in db_output_columns:
                output_alias_map[col] = col

        crossmatch_radius_deg = crossmatch_radius_arcsec / 3600.0

        if q3c_bool:
            constraints = (
                f"q3c_radial_query({ra_field_name},{dec_field_name},"
                f"{ra},{dec},{crossmatch_radius_deg}) "
            )
        else:
            ra_min = ra - crossmatch_radius_deg
            ra_max = ra + crossmatch_radius_deg
            dec_min = dec - crossmatch_radius_deg
            dec_max = dec + crossmatch_radius_deg
            constraints = (
                f" {ra_field_name} between {ra_min} and {ra_max} AND "
                f"{dec_field_name} between {dec_min} and {dec_max} "
            )

        if query_constraints is not None:
            constraints += f"""AND {query_constraints.parse_constraints()}"""

        select = f""" {'"' + '","'.join(db_output_columns) + '"'}"""
        if query_distance_bool:
            if q3c_bool:
                select = (
                    f"q3c_dist({ra_field_name},{dec_field_name},{ra},{dec}) AS xdist,"
                    + select
                )
            else:
                select = f"""{ra_field_name} - ra AS xdist,""" + select

        query = f"""SELECT {select} FROM {db_table} WHERE {constraints}"""

        if order_field_name is not None:
            query += f""" ORDER BY {order_field_name}"""
        if num_limit is not None:
            query += f""" LIMIT {num_limit}"""

        query += ";"

        query_output = self.execute_query(query, db_name)
        all_query_res = []

        for entry in query_output:
            if not query_distance_bool:
                assert len(entry) == len(db_output_columns)
            else:
                assert len(entry) == len(db_output_columns) + 1
            query_res = {}
            for i, key in enumerate(output_alias_map):
                query_res[key] = entry[i]
                if query_distance_bool:
                    query_res["dist"] = entry["xdist"]
            all_query_res.append(query_res)
        return all_query_res

    def check_if_exists(
        self, check_command: str, check_value: str, db_name: str = "postgres"
    ) -> bool:
        """
        Check if a user account exists

        :param check_command if a user/database/table exists
        :param check_value: username to check
        :param db_name: name of database to query
        :return: boolean
        """
        with psycopg.connect(
            f"dbname={db_name} user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            data = conn.execute(check_command).fetchall()
        existing_user_names = [x[0] for x in data]
        logger.debug(f"Found the following values: {existing_user_names}")

        return check_value in existing_user_names

    def create_db(self, db_name: str):
        """
        Creates a database using credentials

        :param db_name: DB to create
        :return: None
        """

        with psycopg.connect(
            f"dbname=postgres user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            sql = f"""CREATE database {db_name}"""
            conn.execute(sql)
            logger.info(f"Created db {db_name}")

    def check_if_db_exists(self, db_name: str) -> bool:
        """
        Check if a user account exists

        :param db_name: database to check
        :return: boolean
        """

        check_command = """SELECT datname FROM pg_database;"""

        db_exist_bool = self.check_if_exists(
            check_command=check_command,
            check_value=db_name,
            db_name="postgres",
        )

        logger.debug(f"Database '{db_name}' does {['not ', ''][db_exist_bool]} exist")

        return db_exist_bool

    def check_if_table_exists(self, db_name: str, db_table: str) -> bool:
        """
        Check if a db table account exists

        :param db_name: database to check
        :param db_table: table to check
        :return: boolean
        """

        check_command = (
            "SELECT table_name FROM information_schema.tables "
            "WHERE table_schema='public';"
        )

        table_exist_bool = self.check_if_exists(
            check_command=check_command,
            check_value=db_table,
            db_name=db_name,
        )

        logger.debug(f"Table '{db_table}' does {['not ', ''][table_exist_bool]} exist")

        return table_exist_bool


class PostgresAdmin(PostgresUser):
    """
    An Admin postgres user, with additional functionality for creatying new users
    """

    user_env_varaiable = PG_ADMIN_USER_KEY
    pass_env_variable = PG_ADMIN_PWD_KEY

    def __init__(self, db_user: str = ADMIN_USER, db_password: str = ADMIN_PASSWORD):
        super().__init__(db_user=db_user, db_password=db_password)

    def create_new_user(self, new_db_user: str, new_password: str):
        """
        Create a new postgres user

        :param new_db_user: new username
        :param new_password: new user password
        :return: None
        """

        with psycopg.connect(
            f"dbname=postgres user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            command = f"CREATE ROLE {new_db_user} WITH password '{new_password}' LOGIN;"
            conn.execute(command)

    def grant_privileges(self, db_name: str, db_user: str):
        """
        Grant privilege to user on database

        :param db_name: name of database
        :param db_user: username to grant privileges for db_user
        :return: None
        """
        with psycopg.connect(
            f"dbname=postgres user={self.db_user} password={self.db_password}"
        ) as conn:
            conn.autocommit = True
            command = f"""GRANT ALL PRIVILEGES ON DATABASE {db_name} TO {db_user};"""
            conn.execute(command)

    def check_if_user_exists(self, user_name: str) -> bool:
        """
        Check if a user account exists

        :param user_name: username to check
        :return: boolean
        """
        check_command = """SELECT usename FROM pg_user;"""

        user_exist_bool = self.check_if_exists(
            check_command=check_command,
            check_value=user_name,
        )

        logger.debug(f"User '{user_name}' does {['not ', ''][user_exist_bool]} exist")

        return user_exist_bool

1	"""
2	Module containing postgres util functions
3	"""
4	# pylint: disable=not-context-manager
5	import logging	2✔
6	import os	2✔
7	from glob import glob	2✔
8	from pathlib import Path	2✔
9	from typing import Optional	2✔
10
11	import numpy as np	2✔
12	import psycopg	2✔
13	from psycopg import errors	2✔
14	from psycopg.rows import Row	2✔
15
16	from winterdrp.data import DataBlock	2✔
17	from winterdrp.errors import ProcessorError	2✔
18	from winterdrp.processors.database.constraints import DBQueryConstraints	2✔
19	from winterdrp.processors.database.utils import get_ordered_schema_list	2✔
20
21	logger = logging.getLogger(__name__)	2✔
22
23	DB_USER_KEY = "DB_USER"	2✔
24	DB_PASSWORD_KEY = "DB_PWD"	2✔
25
26	PG_ADMIN_USER_KEY = "PG_ADMIN_USER"	2✔
27	PG_ADMIN_PWD_KEY = "PG_ADMIN_PWD"	2✔
28
29	DB_USER = os.getenv(DB_USER_KEY)	2✔
30	DB_PASSWORD = os.getenv(DB_PASSWORD_KEY)	2✔
31
32	ADMIN_USER = os.getenv(PG_ADMIN_USER_KEY, DB_USER)	2✔
33	ADMIN_PASSWORD = os.getenv(PG_ADMIN_PWD_KEY, DB_PASSWORD)	2✔
34
35	POSTGRES_DUPLICATE_PROTOCOLS = ["fail", "ignore", "replace"]	2✔
36
37
38	class DataBaseError(ProcessorError):	2✔
39	"""Error relating to postgres interactions"""
40
41
42	class PostgresUser:	2✔
43	"""
44	Basic Postgres user class for executing functions
45	"""
46
47	user_env_varaiable = DB_USER_KEY	2✔
48	pass_env_variable = DB_PASSWORD_KEY	2✔
49
50	def __init__(self, db_user: str = DB_USER, db_password: str = DB_PASSWORD):	2✔
51	self.db_user = db_user	2✔
52	self.db_password = db_password	2✔
53
54	def validate_credentials(self):	2✔
55	"""
56	Checks that user credentials exist
57	:return: None
58	"""
59	if self.db_user is None:	2✔
60	err = (
61	f"'db_user' is set as None. Please pass a db_user as an argument, "
62	f"or set the environment variable '{self.user_env_varaiable}'."
63	)
64	logger.error(err)
65	raise DataBaseError(err)
66
67	if self.db_password is None:	2✔
68	err = (
69	f"'db_password' is set as None. Please pass a password as an argument, "
70	f"or set the environment variable '{self.pass_env_variable}'."
71	)
72	logger.error(err)
73	raise DataBaseError(err)
74
75	# TODO check user exists
76
77	def run_sql_command_from_file(self, file_path: str \| Path, db_name: str):	2✔
78	"""
79	Execute SQL command from file
80
81	:param file_path: File to execute
82	:param db_name: name of database
83	:return: False
84	"""
85	with psycopg.connect(	×
86	f"dbname={db_name} user={self.db_user} password={self.db_password}"
87	) as conn:
88	with open(file_path, "r", encoding="utf8") as sql_file:	×
89	conn.execute(sql_file.read())	×
90
91	logger.info(f"Executed sql commands from file {file_path}")	×
92
93	def create_table(self, schema_path: str \| Path, db_name: str):	2✔
94	"""
95	Create a database table
96
97	:param schema_path: File to execute
98	:param db_name: name of database
99	:return: None
100	"""
101	with psycopg.connect(	2✔
102	f"dbname={db_name} user={self.db_user} password={self.db_password}"
103	) as conn:
104	conn.autocommit = True	2✔
105	with open(schema_path, "r", encoding="utf8") as schema_file:	2✔
106	conn.execute(schema_file.read())	2✔
107
108	logger.info(f"Created table from schema path {schema_path}")	2✔
109
110	def create_tables_from_schema(	2✔
111	self,
112	schema_dir: str \| Path,
113	db_name: str,
114	):
115	"""
116	Creates a db with tables, as described by .sql files in a directory
117
118	:param schema_dir: Directory containing schema files
119	:param db_name: name of DB
120	:return: None
121	"""
122	schema_files = glob(f"{schema_dir}/*.sql")	2✔
123	ordered_schema_files = get_ordered_schema_list(schema_files)	2✔
124	logger.info(f"Creating the following tables - {ordered_schema_files}")	2✔
125	for schema_file in ordered_schema_files:	2✔
126	self.create_table(schema_path=schema_file, db_name=db_name)	2✔
127
128	def export_to_db(	2✔
129	self,
130	value_dict: dict \| DataBlock,
131	db_name: str,
132	db_table: str,
133	duplicate_protocol: str = "fail",
134	) -> tuple[list, list]:
135	"""
136	Export a list of fields in value dict to a batabase table
137
138	:param value_dict: dictionary/DataBlock/other dictonary-like object to export
139	:param db_name: name of db to export to
140	:param db_table: table of DB to export to
141	:param duplicate_protocol: protocol for handling duplicates,
142	in "fail"/"ignore"/"replace"
143	:return:
144	"""
145
146	assert duplicate_protocol in POSTGRES_DUPLICATE_PROTOCOLS	2✔
147
148	with psycopg.connect(	2✔
149	f"dbname={db_name} user={self.db_user} password={self.db_password}"
150	) as conn:
151	conn.autocommit = True	2✔
152
153	sql_query = f"""	2✔
154	SELECT Col.Column_Name from
155	INFORMATION_SCHEMA.TABLE_CONSTRAINTS Tab,
156	INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE Col
157	WHERE
158	Col.Constraint_Name = Tab.Constraint_Name
159	AND Col.Table_Name = Tab.Table_Name
160	AND Constraint_Type = 'PRIMARY KEY'
161	AND Col.Table_Name = '{db_table}'
162	"""
163	serial_keys, serial_key_values = [], []	2✔
164	with conn.execute(sql_query) as cursor:	2✔
165	primary_key = [x[0] for x in cursor.fetchall()]	2✔
166	serial_keys = list(self.get_sequence_keys_from_table(db_table, db_name))	2✔
167	logger.debug(serial_keys)	2✔
168	colnames = [	2✔
169	desc[0]
170	for desc in conn.execute(
171	f"SELECT * FROM {db_table} LIMIT 1"
172	).description
173	if desc[0] not in serial_keys
174	]
175
176	colnames_str = ""	2✔
177	for column in colnames:	2✔
178	colnames_str += f'"{column}",'	2✔
179	colnames_str = colnames_str[:-1]	2✔
180	txt = f"INSERT INTO {db_table} ({colnames_str}) VALUES ("	2✔
181
182	for char in ["[", "]", "'"]:	2✔
183	txt = txt.replace(char, "")	2✔
184
185	for column in colnames:	2✔
186	txt += f"'{str(value_dict[column])}', "	2✔
187
188	txt = txt + ") "	2✔
189	txt = txt.replace(", )", ")")	2✔
190
191	if len(serial_keys) > 0:	2✔
192	txt += "RETURNING "	2✔
193	for key in serial_keys:	2✔
194	txt += f"{key},"	2✔
195	txt += ";"	2✔
196	txt = txt.replace(",;", ";")	2✔
197
198	logger.debug(txt)	2✔
199	command = txt	2✔
200
201	try:	2✔
202	cursor.execute(command)	2✔
203	if len(serial_keys) > 0:	2✔
204	serial_key_values = cursor.fetchall()[0]	2✔
205	else:
206	serial_key_values = []	2✔
207
208	except errors.UniqueViolation as exc:	×
209	primary_key_values = [value_dict[x] for x in primary_key]	×
210
211	if duplicate_protocol == "fail":	×
212	err = (
213	f"Duplicate error, entry with "
214	f"{primary_key}={primary_key_values} "
215	f"already exists in {db_name}."
216	)
217	logger.error(err)
218	raise errors.UniqueViolation from exc
219
220	if duplicate_protocol == "ignore":	×
221	logger.debug(	×
222	f"Found duplicate entry with "
223	f"{primary_key}={primary_key_values} in {db_name}. "
224	f"Ignoring."
225	)
226	elif duplicate_protocol == "replace":	×
227	logger.debug(	×
228	f"Updating duplicate entry with "
229	f"{primary_key}={primary_key_values} in {db_name}."
230	)
231
232	db_constraints = DBQueryConstraints(	×
233	columns=primary_key,
234	accepted_values=primary_key_values,
235	)
236
237	update_colnames = []	×
238	for column in colnames:	×
239	if column not in primary_key:	×
240	update_colnames.append(column)	×
241
242	serial_key_values = self.modify_db_entry(	×
243	db_constraints=db_constraints,
244	value_dict=value_dict,
245	db_alter_columns=update_colnames,
246	db_table=db_table,
247	db_name=db_name,
248	return_columns=serial_keys,
249	)
250
251	return serial_keys, serial_key_values	2✔
252
253	def modify_db_entry(	2✔
254	self,
255	db_name: str,
256	db_table: str,
257	db_constraints: DBQueryConstraints,
258	value_dict: dict \| DataBlock,
259	db_alter_columns: str \| list[str],
260	return_columns: Optional[str \| list[str]] = None,
261	) -> list[Row]:
262	"""
263	Modify a db entry
264
265	:param db_name: name of db
266	:param db_table: Name of table
267	:param value_dict: dict-like object to provide updated values
268	:param db_alter_columns: columns to alter in db
269	:param return_columns: columns to return
270	:return: db query (return columns)
271	"""
272
273	if not isinstance(db_alter_columns, list):	2✔
274	db_alter_columns = [db_alter_columns]	2✔
275
276	if return_columns is None:	2✔
277	return_columns = db_alter_columns	2✔
278	if not isinstance(return_columns, list):	2✔
279	return_columns = [return_columns]	×
280
281	constraints = db_constraints.parse_constraints()	2✔
282
283	with psycopg.connect(	2✔
284	f"dbname={db_name} user={self.db_user} password={self.db_password}"
285	) as conn:
286	conn.autocommit = True	2✔
287
288	db_alter_values = [str(value_dict[c]) for c in db_alter_columns]	2✔
289
290	alter_values_txt = [	2✔
291	f"{db_alter_columns[ind]}='{db_alter_values[ind]}'"
292	for ind in range(len(db_alter_columns))
293	]
294
295	sql_query = (	2✔
296	f"UPDATE {db_table} SET {', '.join(alter_values_txt)} "
297	f"WHERE {constraints}"
298	)
299
300	if len(return_columns) > 0:	2✔
301	logger.debug(return_columns)	2✔
302	sql_query += f""" RETURNING {', '.join(return_columns)}"""	2✔
303	sql_query += ";"	2✔
304	query_output = self.execute_query(sql_query, db_name)	2✔
305
306	return query_output	2✔
307
308	def get_sequence_keys_from_table(self, db_table: str, db_name: str) -> np.ndarray:	2✔
309	"""
310	Gets sequence keys of db table
311
312	:param db_table: database table to use
313	:param db_name: dataname name
314	:return: numpy array of keys
315	"""
316	with psycopg.connect(	2✔
317	f"dbname={db_name} user={self.db_user} password={self.db_password}"
318	) as conn:
319	conn.autocommit = True	2✔
320	sequences = [	2✔
321	x[0]
322	for x in conn.execute(
323	"SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';"
324	).fetchall()
325	]
326	seq_tables = np.array([x.split("_")[0] for x in sequences])	2✔
327	seq_columns = np.array([x.split("_")[1] for x in sequences])	2✔
328	table_sequence_keys = seq_columns[(seq_tables == db_table)]	2✔
329	return table_sequence_keys	2✔
330
331	def import_from_db(	2✔
332	self,
333	db_name: str,
334	db_table: str,
335	db_output_columns: str \| list[str],
336	output_alias_map: Optional[str \| list[str]] = None,
337	max_num_results: Optional[int] = None,
338	db_constraints: Optional[DBQueryConstraints] = None,
339	) -> list[dict]:
340	"""Query an SQL database with constraints, and return a list of dictionaries.
341	One dictionary per entry returned from the query.
342
343	#TODO check admin
344
345	Parameters
346	----------
347	db_name: Name of database to query
348	db_table: Name of database table to query
349	db_output_columns: Name(s) of columns to return for matched database entries
350	output_alias_map: Alias to assign for each output column
351	max_num_results: Maximum number of results to return
352
353	Returns
354	-------
355	A list of dictionaries (one per entry)
356	"""
357
358	if not isinstance(db_output_columns, list):	×
359	db_output_columns = [db_output_columns]	×
360
361	if output_alias_map is None:	×
362	output_alias_map = db_output_columns	×
363
364	if not isinstance(output_alias_map, list):	×
365	output_alias_map = [output_alias_map]	×
366
367	assert len(output_alias_map) == len(db_output_columns)	×
368
369	all_query_res = []	×
370
371	if db_constraints is not None:	×
372	constraints = db_constraints.parse_constraints()	×
373	else:
374	constraints = ""	×
375
376	with psycopg.connect(	×
377	f"dbname={db_name} user={self.db_user} password={self.db_password}"
378	) as conn:
379	conn.autocommit = True	×
380	sql_query = f"""	×
381	SELECT {', '.join(db_output_columns)} from {db_table}
382	WHERE {constraints}
383	"""
384
385	if max_num_results is not None:	×
386	sql_query += f" LIMIT {max_num_results}"	×
387
388	sql_query += ";"	×
389
390	logger.debug(f"Query: {sql_query}")	×
391
392	with conn.execute(sql_query) as cursor:	×
393	query_output = cursor.fetchall()	×
394
395	for entry in query_output:	×
396	assert len(entry) == len(db_output_columns)	×
397
398	query_res = {}	×
399
400	for i, key in enumerate(output_alias_map):	×
401	query_res[key] = entry[i]	×
402
403	all_query_res.append(query_res)	×
404
405	return all_query_res	×
406
407	def execute_query(self, sql_query: str, db_name: str) -> list[Row]:	2✔
408	"""
409	Generically execute SQL query
410
411	:param sql_query: SQL query to execute
412	:param db_name: db name
413	:return: rows from db
414	"""
415	with psycopg.connect(	2✔
416	f"dbname={db_name} user={self.db_user} password={self.db_password}"
417	) as conn:
418	conn.autocommit = True	2✔
419	logger.debug(f"Query: {sql_query}")	2✔
420
421	with conn.execute(sql_query) as cursor:	2✔
422	query_output = cursor.fetchall()	2✔
423
424	return query_output	2✔
425
426	def crossmatch_with_database(	2✔
427	self,
428	db_name: str,
429	db_table: str,
430	db_output_columns: str \| list[str],
431	ra: float,
432	dec: float,
433	crossmatch_radius_arcsec: float,
434	output_alias_map: Optional[dict] = None,
435	ra_field_name: str = "ra",
436	dec_field_name: str = "dec",
437	query_distance_bool: bool = False,
438	q3c_bool: bool = False,
439	query_constraints: Optional[DBQueryConstraints] = None,
440	order_field_name: Optional[str] = None,
441	num_limit: Optional[int] = None,
442	) -> list[dict]:
443	"""
444	Crossmatch a given spatial position (ra/dec) with sources in a database,
445	and returns a list of matches
446
447	#TODO: check admin
448
449	:param db_name: name of db to query
450	:param db_table: name of db table
451	:param db_output_columns: columns to return
452	:param output_alias_map: mapping for renaming columns
453	:param ra: RA
454	:param dec: dec
455	:param crossmatch_radius_arcsec: radius for crossmatch
456	:param ra_field_name: name of ra column in database
457	:param dec_field_name: name of dec column in database
458	:param query_distance_bool: boolean where to return crossmatch distance
459	:param q3c_bool: boolean whether to use q3c_bool
460	:param order_field_name: field to order result by
461	:param num_limit: limit on sql query
462	:return: list of query result dictionaries
463	"""
464
465	if output_alias_map is None:	2✔
466	output_alias_map = {}	2✔
467	for col in db_output_columns:	2✔
468	output_alias_map[col] = col	2✔
469
470	crossmatch_radius_deg = crossmatch_radius_arcsec / 3600.0	2✔
471
472	if q3c_bool:	2✔
473	constraints = (	×
474	f"q3c_radial_query({ra_field_name},{dec_field_name},"
475	f"{ra},{dec},{crossmatch_radius_deg}) "
476	)
477	else:
478	ra_min = ra - crossmatch_radius_deg	2✔
479	ra_max = ra + crossmatch_radius_deg	2✔
480	dec_min = dec - crossmatch_radius_deg	2✔
481	dec_max = dec + crossmatch_radius_deg	2✔
482	constraints = (	2✔
483	f" {ra_field_name} between {ra_min} and {ra_max} AND "
484	f"{dec_field_name} between {dec_min} and {dec_max} "
485	)
486
487	if query_constraints is not None:	2✔
488	constraints += f"""AND {query_constraints.parse_constraints()}"""	2✔
489
490	select = f""" {'"' + '","'.join(db_output_columns) + '"'}"""	2✔
491	if query_distance_bool:	2✔
492	if q3c_bool:	×
493	select = (	×
494	f"q3c_dist({ra_field_name},{dec_field_name},{ra},{dec}) AS xdist,"
495	+ select
496	)
497	else:
498	select = f"""{ra_field_name} - ra AS xdist,""" + select	×
499
500	query = f"""SELECT {select} FROM {db_table} WHERE {constraints}"""	2✔
501
502	if order_field_name is not None:	2✔
503	query += f""" ORDER BY {order_field_name}"""	×
504	if num_limit is not None:	2✔
505	query += f""" LIMIT {num_limit}"""	2✔
506
507	query += ";"	2✔
508
509	query_output = self.execute_query(query, db_name)	2✔
510	all_query_res = []	2✔
511
512	for entry in query_output:	2✔
513	if not query_distance_bool:	×
514	assert len(entry) == len(db_output_columns)	×
515	else:
516	assert len(entry) == len(db_output_columns) + 1	×
517	query_res = {}	×
518	for i, key in enumerate(output_alias_map):	×
519	query_res[key] = entry[i]	×
520	if query_distance_bool:	×
521	query_res["dist"] = entry["xdist"]	×
522	all_query_res.append(query_res)	×
523	return all_query_res	2✔
524
525	def check_if_exists(	2✔
526	self, check_command: str, check_value: str, db_name: str = "postgres"
527	) -> bool:
528	"""
529	Check if a user account exists
530
531	:param check_command if a user/database/table exists
532	:param check_value: username to check
533	:param db_name: name of database to query
534	:return: boolean
535	"""
536	with psycopg.connect(	2✔
537	f"dbname={db_name} user={self.db_user} password={self.db_password}"
538	) as conn:
539	conn.autocommit = True	2✔
540	data = conn.execute(check_command).fetchall()	2✔
541	existing_user_names = [x[0] for x in data]	2✔
542	logger.debug(f"Found the following values: {existing_user_names}")	2✔
543
544	return check_value in existing_user_names	2✔
545
546	def create_db(self, db_name: str):	2✔
547	"""
548	Creates a database using credentials
549
550	:param db_name: DB to create
551	:return: None
552	"""
553
554	with psycopg.connect(	2✔
555	f"dbname=postgres user={self.db_user} password={self.db_password}"
556	) as conn:
557	conn.autocommit = True	2✔
558	sql = f"""CREATE database {db_name}"""	2✔
559	conn.execute(sql)	2✔
560	logger.info(f"Created db {db_name}")	2✔
561
562	def check_if_db_exists(self, db_name: str) -> bool:	2✔
563	"""
564	Check if a user account exists
565
566	:param db_name: database to check
567	:return: boolean
568	"""
569
570	check_command = """SELECT datname FROM pg_database;"""	2✔
571
572	db_exist_bool = self.check_if_exists(	2✔
573	check_command=check_command,
574	check_value=db_name,
575	db_name="postgres",
576	)
577
578	logger.debug(f"Database '{db_name}' does {['not ', ''][db_exist_bool]} exist")	2✔
579
580	return db_exist_bool	2✔
581
582	def check_if_table_exists(self, db_name: str, db_table: str) -> bool:	2✔
583	"""
584	Check if a db table account exists
585
586	:param db_name: database to check
587	:param db_table: table to check
588	:return: boolean
589	"""
590
591	check_command = (	2✔
592	"SELECT table_name FROM information_schema.tables "
593	"WHERE table_schema='public';"
594	)
595
596	table_exist_bool = self.check_if_exists(	2✔
597	check_command=check_command,
598	check_value=db_table,
599	db_name=db_name,
600	)
601
602	logger.debug(f"Table '{db_table}' does {['not ', ''][table_exist_bool]} exist")	2✔
603
604	return table_exist_bool	2✔
605
606
607	class PostgresAdmin(PostgresUser):	2✔
608	"""
609	An Admin postgres user, with additional functionality for creatying new users
610	"""
611
612	user_env_varaiable = PG_ADMIN_USER_KEY	2✔
613	pass_env_variable = PG_ADMIN_PWD_KEY	2✔
614
615	def __init__(self, db_user: str = ADMIN_USER, db_password: str = ADMIN_PASSWORD):	2✔
616	super().__init__(db_user=db_user, db_password=db_password)	2✔
617
618	def create_new_user(self, new_db_user: str, new_password: str):	2✔
619	"""
620	Create a new postgres user
621
622	:param new_db_user: new username
623	:param new_password: new user password
624	:return: None
625	"""
626
627	with psycopg.connect(	×
628	f"dbname=postgres user={self.db_user} password={self.db_password}"
629	) as conn:
630	conn.autocommit = True	×
631	command = f"CREATE ROLE {new_db_user} WITH password '{new_password}' LOGIN;"	×
632	conn.execute(command)	×
633
634	def grant_privileges(self, db_name: str, db_user: str):	2✔
635	"""
636	Grant privilege to user on database
637
638	:param db_name: name of database
639	:param db_user: username to grant privileges for db_user
640	:return: None
641	"""
642	with psycopg.connect(	2✔
643	f"dbname=postgres user={self.db_user} password={self.db_password}"
644	) as conn:
645	conn.autocommit = True	2✔
646	command = f"""GRANT ALL PRIVILEGES ON DATABASE {db_name} TO {db_user};"""	2✔
647	conn.execute(command)	2✔
648
649	def check_if_user_exists(self, user_name: str) -> bool:	2✔
650	"""
651	Check if a user account exists
652
653	:param user_name: username to check
654	:return: boolean
655	"""
656	check_command = """SELECT usename FROM pg_user;"""	2✔
657
658	user_exist_bool = self.check_if_exists(	2✔
659	check_command=check_command,
660	check_value=user_name,
661	)
662
663	logger.debug(f"User '{user_name}' does {['not ', ''][user_exist_bool]} exist")	2✔
664
665	return user_exist_bool	2✔

winter-telescope / winterdrp / 4119519683

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous