SQL Postgres¶
PgSqlHandler(connection_properties, max_retries=8, max_conn_retries=3, retry_wait_time=30, exponential_min_retry_wait_time=0.1, exponential_multiplier=1, skip_connect=True, pg_settings=None, **kwargs)
¶
Subclass of SqlHandler used for PostgreSQL.
This already connects to the SQL server.
This handler implements a robust retry strategy using exponential backoff to gracefully manage connection or query failures. By default, operations are attempted up to 8 times; between failures, the wait time doubles exponentially—starting at 1 second (1s, 2s, 4s, 8s, 16s)—until it reaches a configured ceiling of 30 seconds. This approach allows for rapid recovery from momentary glitches while preventing server overload during persistent outages.
Using the default parameters the connection will fail after around 90 seconds if the server is not reachable.
Parameters:
-
(connection_properties¶SqlConnProperties) –Object containing connection parameters.
-
(max_retries¶int, default:8) –Number of retries that will be attempted when doing queries. Will be used in
stopparameter of tenacity.stop_after_attempt, by default 8 -
(max_conn_retries¶int, default:3) –Number of retries that will be attempted when reconnecting. Will be used in
stopparameter of tenacity.stop_after_attempt, by default 3 -
(retry_wait_time¶float, default:30) –Max time to wait between retries when reconnecting or doing queries. Will be used in
maxparameter of tenacity.wait_exponential, by default 30 -
(exponential_min_retry_wait_time¶float, default:0.1) –Min time to wait between retries when reconnecting or doing queries. Will be used in
minparameter of tenacity.wait_exponential, by default 1 -
(exponential_multiplier¶float, default:1) –Multiplier to use when calculating wait time between retries when reconnecting or doing queries. Will be used in
multiplierparameter of tenacity.wait_exponential, by default 1 -
(skip_connect¶bool, default:True) –If True, the connection will not be established when the object is created.
If this is set toTrue, the user will need to manually call the reconnect() method when they want to connect to the server.
By default True
-
(pg_settings¶dict[str, str] | None, default:None) –Dictionary with additional PostgreSQL settings to be set at the start of the connection.
An example would be {"work_mem": "64MB"} to set the work_mem parameter.
By default idle_session_timeout=30000 (30 seconds)
If None, no additional settings will be set. By default None.
-
(**kwargs¶dict, default:{}) –Just kept here for compatibility.
Source code in echo_connhandler/sql_postgres.py
@validate_call
def __init__(
self,
connection_properties: SqlConnProperties,
max_retries: int = 8,
max_conn_retries: int = 3,
retry_wait_time: float = 30,
exponential_min_retry_wait_time: float = 0.1,
exponential_multiplier: float = 1,
skip_connect: bool = True,
pg_settings: dict[str, str] | None = None,
**kwargs, # pylint: disable=unused-argument # noqa
) -> None:
"""Method that initializes the SQL handler.
This already connects to the SQL server.
This handler implements a robust retry strategy using exponential backoff to gracefully manage connection or query failures. By default, operations are attempted up to 8 times; between failures, the wait time doubles exponentially—starting at 1 second (1s, 2s, 4s, 8s, 16s)—until it reaches a configured ceiling of 30 seconds. This approach allows for rapid recovery from momentary glitches while preventing server overload during persistent outages.
Using the default parameters the connection will fail after around 90 seconds if the server is not reachable.
Parameters
----------
connection_properties : SqlConnProperties
Object containing connection parameters.
max_retries : int, optional
Number of retries that will be attempted when doing queries. Will be used in `stop` parameter of tenacity.stop_after_attempt, by default 8
max_conn_retries : int, optional
Number of retries that will be attempted when reconnecting. Will be used in `stop` parameter of tenacity.stop_after_attempt, by default 3
retry_wait_time : float, optional
Max time to wait between retries when reconnecting or doing queries. Will be used in `max` parameter of tenacity.wait_exponential, by default 30
exponential_min_retry_wait_time: float, optional
Min time to wait between retries when reconnecting or doing queries. Will be used in `min` parameter of tenacity.wait_exponential, by default 1
exponential_multiplier: float, optional
Multiplier to use when calculating wait time between retries when reconnecting or doing queries. Will be used in `multiplier` parameter of tenacity.wait_exponential, by default 1
skip_connect : bool, optional
If True, the connection will not be established when the object is created.
If this is set toTrue, the user will need to manually call the reconnect() method when they want to connect to the server.
By default True
pg_settings : dict[str, str] | None, optional
Dictionary with additional PostgreSQL settings to be set at the start of the connection.
An example would be {"work_mem": "64MB"} to set the work_mem parameter.
By default idle_session_timeout=30000 (30 seconds)
If None, no additional settings will be set. By default None.
**kwargs : dict, optional
Just kept here for compatibility.
"""
# max number of notices to store
self._max_notices = 10
# var to store notices
self.notices = deque(maxlen=self._max_notices)
# setting work_mem if set
# validating
if pg_settings is not None:
if not isinstance(pg_settings, dict):
raise TypeError(f"pg_settings should be a dict, not {type(pg_settings)}")
for key, value in pg_settings.items():
if not isinstance(key, str):
raise TypeError(f"all keys in pg_settings should be strings, not {type(key)}")
if not isinstance(value, str):
raise TypeError(f"all values in pg_settings should be strings, not {type(value)}")
# adding default value for idle_session_timeout if not set
if pg_settings is None:
pg_settings = {"idle_session_timeout": "30000"}
elif "idle_session_timeout" not in pg_settings:
pg_settings["idle_session_timeout"] = "30000"
self.pg_settings = pg_settings
super().__init__(
connection_properties=connection_properties,
max_retries=max_retries,
max_conn_retries=max_conn_retries,
retry_wait_time=retry_wait_time,
exponential_min_retry_wait_time=exponential_min_retry_wait_time,
exponential_multiplier=exponential_multiplier,
skip_connect=skip_connect,
)
create_table(schema, table_name, columns, pkey_cols=None, if_exists='fail', temporary=False)
¶
Method to create a table in the database.
Parameters:
-
(schema¶str | None) –Name of the schema to create the table in. If temporary is set to True, this must be set to None.
-
(table_name¶str) –Name of the table to be created.
-
(columns¶dict[str, str]) –Dict in the format {column_name: column_type, ...}.
-
(pkey_cols¶list[str] | None, default:None) –List of columns to be set as primary key. If None, no primary key will be set. By default None.
-
(if_exists¶Literal['fail', 'replace'], default:'fail') –What to do if the table already exists.
- fail: Will raise a ValueError.
- replace: Will drop the table and create a new one.
By default "fail".
-
(temporary¶bool, default:False) –If set to True, the table will be created as temporary. By default False.
Source code in echo_connhandler/sql_postgres.py
@validate_call
def create_table(
self,
schema: str | None,
table_name: str,
columns: dict[str, str],
pkey_cols: list[str] | None = None,
if_exists: Literal["fail", "replace"] = "fail",
temporary: bool = False,
) -> None:
"""Method to create a table in the database.
Parameters
----------
schema : str | None
Name of the schema to create the table in. If temporary is set to True, this must be set to None.
table_name : str
Name of the table to be created.
columns : dict[str, str]
Dict in the format {column_name: column_type, ...}.
pkey_cols : list[str] | None, optional
List of columns to be set as primary key. If None, no primary key will be set. By default None.
if_exists : Literal["fail", "replace"], optional
What to do if the table already exists.
- fail: Will raise a ValueError.
- replace: Will drop the table and create a new one.
By default "fail".
temporary : bool, optional
If set to True, the table will be created as temporary. By default False.
"""
# checking arguments
if not isinstance(schema, str | type(None)):
raise TypeError(f"schema should be a string or none, not {type(schema)}")
if not isinstance(table_name, str):
raise TypeError(f"table_name should be a string, not {type(table_name)}")
if not isinstance(columns, dict):
raise TypeError(f"columns should be a dict, not {type(columns)}")
if not all(isinstance(col, str) for col in columns):
raise TypeError(f"all column names should be strings, not {type(columns)}")
if not all(isinstance(col_type, str) for col_type in columns.values()):
raise TypeError(f"all column types should be strings, not {type(columns)}")
if pkey_cols is not None:
if not isinstance(pkey_cols, list):
raise TypeError(f"pkey_cols should be a list, not {type(pkey_cols)}")
if not all(isinstance(col, str) for col in pkey_cols):
raise TypeError(f"all primary key columns should be strings, not {type(pkey_cols)}")
if any(col not in columns for col in pkey_cols):
raise ValueError(f"all primary key columns should be in columns, not {pkey_cols}")
if if_exists not in ["fail", "replace"]:
raise ValueError(f"if_exists should be one of ['fail', 'replace'], not {if_exists}")
if not isinstance(temporary, bool):
raise TypeError(f"temporary should be a bool, not {type(temporary)}")
if temporary and schema is not None:
raise ValueError("If temporary is set to True, schema must be set to None.")
# checking if table already exists
table_already_exists = self.table_exists(schema=schema, table_name=table_name)
if table_already_exists and if_exists == "fail":
raise ValueError(f"Table {table_name} already exists in the database.")
if (table_already_exists or temporary) and if_exists == "replace":
self.delete_table(schema=schema, table_name=table_name, on_error="ignore")
# creating table
query = sql.SQL("CREATE {temporary} TABLE {schema}{table_name} ({columns} {pkey})").format(
temporary=sql.SQL("TEMPORARY" if temporary else ""),
schema=sql.SQL("{schema}.").format(schema=sql.Identifier(schema)) if schema else sql.SQL(""),
table_name=sql.Identifier(table_name),
columns=sql.SQL(",").join(
[
sql.SQL("{col_name} {col_type}").format(col_name=sql.Identifier(col), col_type=sql.SQL(col_type))
for col, col_type in columns.items()
],
),
pkey=sql.SQL(", PRIMARY KEY ({pkey_cols})").format(pkey_cols=sql.SQL(",").join([sql.Identifier(col) for col in pkey_cols]))
if pkey_cols
else sql.SQL(""),
)
# executing query
self.execute(query)
logger.debug(f"Table {table_name} created in schema {schema}.")
delete_table(schema, table_name, on_error='fail')
¶
Method to delete a table in the database.
Parameters:
-
(schema¶str | None) –Name of the schema of the table to delete.
-
(table_name¶str) –Name of the table to be deleted.
-
(on_error¶Literal['fail', 'ignore'], default:'fail') –What to do if the table does not exist. - fail: Will raise a ValueError. - ignore: Will ignore the error. By default "fail".
Source code in echo_connhandler/sql_postgres.py
@validate_call
def delete_table(
self,
schema: str | None,
table_name: str,
on_error: Literal["fail", "ignore"] = "fail",
) -> None:
"""Method to delete a table in the database.
Parameters
----------
schema : str | None
Name of the schema of the table to delete.
table_name: str
Name of the table to be deleted.
on_error : Literal["fail", "ignore"], optional
What to do if the table does not exist.
- fail: Will raise a ValueError.
- ignore: Will ignore the error.
By default "fail".
"""
# checking arguments
if not isinstance(schema, str | type(None)):
raise TypeError(f"schema should be a string or None, not {type(schema)}")
if not isinstance(table_name, str):
raise TypeError(f"table_name should be a string, not {type(table_name)}")
if on_error not in ["fail", "ignore"]:
raise ValueError(f"on_error should be one of ['fail', 'ignore'], not {on_error}")
# checking if table exists
if not self.table_exists(schema=schema, table_name=table_name) and schema is not None:
if on_error == "fail":
raise ValueError(f"Table {table_name} does not exist in the database.")
return
# deleting table
query = sql.SQL("DROP TABLE IF EXISTS {schema}{table_name}").format(
schema=sql.SQL("{schema}.").format(schema=sql.Identifier(schema)) if schema else sql.SQL(""),
table_name=sql.Identifier(table_name),
)
self.execute(query)
logger.debug(f"Table {table_name} deleted from schema {schema}.")
get_table_columns(schema, table_names, table_types=None)
¶
Method to get the columns in a table and it's data types.
Parameters:
-
(schema¶str | None) –Name of the schema where the tables are located.
-
(table_names¶list[str]) –Name of the desired tables.
-
(table_types¶list[Literal['table', 'view', 'materialized_view']], default:None) –List of table types to be returned. Can be one of ["table", "view", "materialized_view"]. If set to None, will be equal to ["table"], by default None.
Returns:
-
dict[str, dict[str, str]]–Dict in the format {table_name: {column_name: column_type, ...}, ...}.
Source code in echo_connhandler/sql_postgres.py
@validate_call
def get_table_columns(
self,
schema: str | None,
table_names: list[str],
table_types: list[Literal["table", "view", "materialized_view"]] | None = None,
) -> dict[str, dict[str, str]]:
# sourcery skip: class-extract-method
"""Method to get the columns in a table and it's data types.
Parameters
----------
schema : str | None
Name of the schema where the tables are located.
table_names : list[str]
Name of the desired tables.
table_types : list[Literal["table", "view", "materialized_view"]], optional
List of table types to be returned. Can be one of ["table", "view", "materialized_view"]. If set to None, will be equal to ["table"], by default None.
Returns
-------
dict[str, dict[str, str]]
Dict in the format {table_name: {column_name: column_type, ...}, ...}.
"""
# checking arguments
if not isinstance(table_names, list):
raise TypeError(f"table_names should be a list, not {type(table_names)}")
if not all(isinstance(table, str) for table in table_names):
raise TypeError(f"all table names should be strings, not {type(table_names)}")
if not isinstance(schema, str | type(None)):
raise TypeError(f"schema should be a string or none, not {type(schema)}")
if table_types is not None:
if not isinstance(table_types, list):
raise TypeError(f"table_types should be a list, not {type(table_types)}")
if not all(isinstance(table_type, str) for table_type in table_types):
raise TypeError(f"all table types should be strings, not {type(table_types)}")
if any(table_type not in ["table", "view", "materialized_view"] for table_type in table_types):
raise ValueError(f"all table types should be one of ['table', 'view', 'materialized_view'], not {table_types}")
# setting default value for table_types
if table_types is None:
table_types = ["table"]
_prev_suppress = self._suppress_auto_close
self._suppress_auto_close = True
try:
# checking if all tables exist
existing_tables = self.get_table_names(schema=schema, table_types=table_types)
if missing_tables := set(table_names) - set(existing_tables):
raise ValueError(f"Tables {missing_tables} do not exist in the database.")
query = sql.SQL(
r"""SELECT
c.relname AS table_name,
a.attname AS column_name,
REGEXP_REPLACE(
REGEXP_REPLACE(
format_type(a.atttypid, a.atttypmod),
'character varying(\(.*\))?',
'TEXT',
'g'
),
'^jsonb?$',
'JSON',
'gi'
) AS data_type
FROM pg_attribute a
JOIN pg_class c ON a.attrelid = c.oid
JOIN pg_namespace n ON c.relnamespace = n.oid
WHERE c.relname IN ({tables})
{schema}
AND a.attnum > 0 -- Exclude system columns (oid, ctid, etc.)
AND NOT a.attisdropped -- Exclude deleted columns
ORDER BY a.attnum;""",
).format(
schema=sql.SQL("AND n.nspname = {schema}").format(schema=sql.Literal(schema)) if schema else sql.SQL(""),
tables=sql.SQL(",").join([sql.Literal(table) for table in table_names]),
)
df = self.read_to_pandas(query)
return {
table: dict(
df[df["table_name"] == table][["column_name", "data_type"]].values,
)
for table in table_names
}
finally:
self._suppress_auto_close = _prev_suppress
if not _prev_suppress:
self._auto_close_if_needed()
get_table_names(schema, table_types=None)
¶
Method to get the names of all tables in the database.
Parameters:
-
(schema¶str | None) –Name of the schema where the tables are located.
-
(table_types¶list[Literal['table', 'view', 'materialized_view']], default:None) –List of table types to be returned. Can be one of ["table", "view", "materialized_view"]. If set to None, will be equal to ["table"], by default None.
Returns:
-
list[str]–List of table names.
Source code in echo_connhandler/sql_postgres.py
@validate_call
def get_table_names(
self,
schema: str | None,
table_types: list[Literal["table", "view", "materialized_view"]] | None = None,
) -> list[str]:
"""Method to get the names of all tables in the database.
Parameters
----------
schema : str | None
Name of the schema where the tables are located.
table_types : list[Literal["table", "view", "materialized_view"]], optional
List of table types to be returned. Can be one of ["table", "view", "materialized_view"]. If set to None, will be equal to ["table"], by default None.
Returns
-------
list[str]
List of table names.
"""
# checking arguments
if not isinstance(schema, str | type(None)):
raise TypeError(f"schema should be a string or none, not {type(schema)}")
if table_types is not None:
if not isinstance(table_types, list):
raise TypeError(f"table_types should be a list, not {type(table_types)}")
if not all(isinstance(table_type, str) for table_type in table_types):
raise TypeError(f"all table types should be strings, not {type(table_types)}")
if any(table_type not in ["table", "view", "materialized_view"] for table_type in table_types):
raise ValueError(f"all table types should be one of ['table', 'view', 'materialized_view'], not {table_types}")
# setting default value for table_types
if table_types is None:
table_types = ["table"]
types_mapping = {
"table": "BASE TABLE",
"view": "VIEW",
"materialized_view": "MATERIALIZED VIEW",
}
_prev_suppress = self._suppress_auto_close
self._suppress_auto_close = True
try:
# getting table names (except materialized_view)
query = sql.SQL(
"SELECT table_name FROM information_schema.tables WHERE table_type IN ({types}) {schema}",
).format(
schema=sql.SQL("AND table_schema = {schema}").format(schema=sql.Literal(schema)) if schema else sql.SQL(""),
types=sql.SQL(",").join([sql.Literal(types_mapping[table_type]) for table_type in table_types]),
)
df = self.read_to_pandas(query)
table_list = df["table_name"].tolist()
# getting materialized_view names
if "materialized_view" in table_types:
query = sql.SQL(
"SELECT matviewname FROM pg_matviews {schema}",
).format(
schema=sql.SQL("WHERE schemaname = {schema}").format(schema=sql.Literal(schema)) if schema else sql.SQL(""),
)
df = self.read_to_pandas(query)
table_list += df["matviewname"].tolist()
return table_list
finally:
self._suppress_auto_close = _prev_suppress
if not _prev_suppress:
self._auto_close_if_needed()
get_table_pkey(schema, table_names, return_type='columns')
¶
Method to get the primary key of a table.
Can either return the name of the primary key or the columns associated with the primary key (depending on the type of SQL)
Parameters:
-
(schema¶str | None) –Name of the schema where the tables are located.
-
(table_names¶list[str]) –Name of the desired tables.
-
(return_type¶Literal['name', 'columns'], default:'columns') –What to return. Can be either "name" of the primary key or "columns" for the columns involved, by default "columns".
Returns:
-
dict[str, list[str] | str]–In case of return_type="columns", returns a dict in the format {table_name: [pk_col1, pk_col2, ...], ...}. In case of return_type="name", returns a dict in the format {table_name: pkey_name, ...}.
Source code in echo_connhandler/sql_postgres.py
@validate_call
def get_table_pkey(
self,
schema: str | None,
table_names: list[str],
return_type: Literal["name", "columns"] = "columns",
) -> dict[str, list[str] | str]:
"""Method to get the primary key of a table.
Can either return the name of the primary key or the columns associated with the primary key (depending on the type of SQL)
Parameters
----------
schema : str | None
Name of the schema where the tables are located.
table_names : list[str]
Name of the desired tables.
return_type : Literal["name", "columns"], optional
What to return. Can be either "name" of the primary key or "columns" for the columns involved, by default "columns".
Returns
-------
dict[str, list[str] | str]
In case of return_type="columns", returns a dict in the format {table_name: [pk_col1, pk_col2, ...], ...}.
In case of return_type="name", returns a dict in the format {table_name: pkey_name, ...}.
"""
# checking arguments
if not isinstance(table_names, list):
raise TypeError(f"table_names should be a list, not {type(table_names)}")
if not all(isinstance(table, str) for table in table_names):
raise TypeError(f"all table names should be strings, not {type(table_names)}")
if not isinstance(schema, str | type(None)):
raise TypeError(f"schema should be a string or none, not {type(schema)}")
if return_type not in ["name", "columns"]:
raise ValueError(f"return_type should be one of ['name', 'columns'], not {return_type}")
_prev_suppress = self._suppress_auto_close
self._suppress_auto_close = True
try:
# checking if all tables exist
existing_tables = self.get_table_names(schema=schema)
if missing_tables := set(table_names) - set(existing_tables):
raise ValueError(f"Tables {missing_tables} do not exist in the database.")
# getting primary keys names
query = sql.SQL(
"SELECT table_name, constraint_name FROM information_schema.table_constraints WHERE table_name IN ({tables}) AND constraint_type = 'PRIMARY KEY' {schema}",
).format(
schema=sql.SQL("AND table_schema = {schema}").format(schema=sql.Literal(schema)) if schema else sql.SQL(""),
tables=sql.SQL(",").join([sql.Literal(table) for table in table_names]),
)
df = self.read_to_pandas(query)
pkey_names = dict(df[["table_name", "constraint_name"]].values)
if return_type == "name":
return pkey_names
if not pkey_names:
return {table: [] for table in table_names}
# getting primary keys columns
query = sql.SQL(
"SELECT table_name, column_name FROM information_schema.constraint_column_usage WHERE table_name IN ({tables}) AND constraint_name IN ({pkey_names}) {schema}",
).format(
schema=sql.SQL("AND table_schema = {schema}").format(schema=sql.Literal(schema)) if schema else sql.SQL(""),
tables=sql.SQL(",").join([sql.Literal(table) for table in table_names]),
pkey_names=sql.SQL(",").join([sql.Literal(pkey) for pkey in pkey_names.values()]),
)
df = self.read_to_pandas(query)
return df.groupby("table_name")["column_name"].apply(list).to_dict()
finally:
self._suppress_auto_close = _prev_suppress
if not _prev_suppress:
self._auto_close_if_needed()
table_exists(schema, table_name)
¶
Method to check if a table, view, or materialized view exists in the PostgreSQL database.
Parameters:
-
(schema¶str | None) –Name of the schema where the table should be located.
-
(table_name¶str) –Name of the table to check.
Returns:
-
bool–True if the table exists, False otherwise.
Source code in echo_connhandler/sql_postgres.py
@validate_call
def table_exists(self, schema: str | None, table_name: str) -> bool:
"""Method to check if a table, view, or materialized view exists in the PostgreSQL database.
Parameters
----------
schema : str | None
Name of the schema where the table should be located.
table_name : str
Name of the table to check.
Returns
-------
bool
True if the table exists, False otherwise.
"""
# checking arguments
if not isinstance(schema, str | type(None)):
raise TypeError(f"schema should be a string or None, not {type(schema)}")
if not isinstance(table_name, str):
raise TypeError(f"table_name should be a string, not {type(table_name)}")
_prev_suppress = self._suppress_auto_close
self._suppress_auto_close = True
try:
if schema:
# Direct check in pg_class and pg_namespace (extremely fast)
query = sql.SQL(
"SELECT EXISTS (SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = {schema} AND c.relname = {table_name})",
).format(
schema=sql.Literal(schema),
table_name=sql.Literal(table_name),
)
else:
# If no schema is passed, check if the table is visible in the current search_path
query = sql.SQL(
"SELECT EXISTS (SELECT 1 FROM pg_class WHERE relname = {table_name} AND pg_table_is_visible(oid))",
).format(
table_name=sql.Literal(table_name),
)
with self.cursor() as cur:
cur.execute(query)
return cur.fetchone()[0]
finally:
self._suppress_auto_close = _prev_suppress
if not _prev_suppress:
self._auto_close_if_needed()