import psycopg2 import time from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT postgres_table_template = """ CREATE TABLE IF NOT EXISTS "{}" ( key Integer NOT NULL, value Integer, PRIMARY KEY(key)) """ postgres_table_template_2 = """ CREATE TABLE IF NOT EXISTS "{}" ( key Integer NOT NULL, value1 Integer, value2 Integer, value3 Integer, PRIMARY KEY(key)) """ postgres_table_template_3 = """ CREATE TABLE IF NOT EXISTS "{}" ( key1 Integer NOT NULL, value1 Integer, key2 Integer NOT NULL, value2 Integer NOT NULL) """ postgres_table_template_4 = """ CREATE TABLE IF NOT EXISTS "{}"."{}" ( key Integer NOT NULL, value Integer, PRIMARY KEY(key)) """ postgres_table_template_5 = """ CREATE TABLE IF NOT EXISTS "{}" ( key Integer NOT NULL, value UUID, PRIMARY KEY(key)) """ def get_postgres_conn( ip, port, database=False, auto_commit=True, database_name="postgres_database", replication=False, ): if database == True: conn_string = f"host={ip} port={port} dbname='{database_name}' user='postgres' password='mysecretpassword'" else: conn_string = ( f"host={ip} port={port} user='postgres' password='mysecretpassword'" ) if replication: conn_string += " replication='database'" conn = psycopg2.connect(conn_string) if auto_commit: conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) conn.autocommit = True return conn def create_replication_slot(conn, slot_name="user_slot"): cursor = conn.cursor() cursor.execute( f"CREATE_REPLICATION_SLOT {slot_name} LOGICAL pgoutput EXPORT_SNAPSHOT" ) result = cursor.fetchall() print(result[0][0]) # slot name print(result[0][1]) # start lsn print(result[0][2]) # snapshot return result[0][2] def drop_replication_slot(conn, slot_name="user_slot"): cursor = conn.cursor() cursor.execute(f"select pg_drop_replication_slot('{slot_name}')") def create_postgres_schema(cursor, schema_name): drop_postgres_schema(cursor, schema_name) cursor.execute(f"CREATE SCHEMA {schema_name}") def drop_postgres_schema(cursor, schema_name): cursor.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE") def create_postgres_table( cursor, table_name, database_name="", replica_identity_full=False, template=postgres_table_template, ): if database_name == "": name = table_name else: name = f"{database_name}.{table_name}" drop_postgres_table(cursor, name) query = template.format(name) cursor.execute(query) print(f"Query: {query}") if replica_identity_full: cursor.execute(f"ALTER TABLE {name} REPLICA IDENTITY FULL;") def drop_postgres_table(cursor, name): cursor.execute(f"""DROP TABLE IF EXISTS "{name}" """) def create_postgres_table_with_schema(cursor, schema_name, table_name): drop_postgres_table_with_schema(cursor, schema_name, table_name) cursor.execute(postgres_table_template_4.format(schema_name, table_name)) def drop_postgres_table_with_schema(cursor, schema_name, table_name): cursor.execute(f"""DROP TABLE IF EXISTS "{schema_name}"."{table_name}" """) class PostgresManager: def __init__(self): self.created_postgres_db_list = set() self.created_materialized_postgres_db_list = set() self.created_ch_postgres_db_list = set() def init( self, instance, ip, port, default_database="postgres_database", postgres_db_exists=False, ): self.instance = instance self.ip = ip self.port = port self.default_database = default_database self.postgres_db_exists = postgres_db_exists self.prepare() def get_default_database(self): return self.default_database def restart(self): try: self.clear() self.prepare() except Exception as ex: self.prepare() raise ex def execute(self, query): self.cursor.execute(query) def prepare(self): self.conn = get_postgres_conn(ip=self.ip, port=self.port) self.cursor = self.conn.cursor() if self.default_database != "": if not self.postgres_db_exists: self.create_postgres_db(self.default_database) self.conn = get_postgres_conn( ip=self.ip, port=self.port, database=True, database_name=self.default_database, ) self.cursor = self.conn.cursor() self.create_clickhouse_postgres_db() def clear(self): if self.conn.closed == 0: self.conn.close() for db in self.created_materialized_postgres_db_list.copy(): self.drop_materialized_db(db) for db in self.created_ch_postgres_db_list.copy(): self.drop_clickhouse_postgres_db(db) if len(self.created_postgres_db_list) > 0: self.conn = get_postgres_conn(ip=self.ip, port=self.port) self.cursor = self.conn.cursor() for db in self.created_postgres_db_list.copy(): self.drop_postgres_db(db) def get_db_cursor(self, database_name=""): if database_name == "": database_name = self.default_database self.conn = get_postgres_conn( ip=self.ip, port=self.port, database=True, database_name=database_name ) return self.conn.cursor() def database_or_default(self, database_name): if database_name != "": return database_name if self.default_database != "": return self.default_database raise Exception("Database name is empty") def create_postgres_db(self, database_name=""): database_name = self.database_or_default(database_name) self.drop_postgres_db(database_name) self.created_postgres_db_list.add(database_name) self.cursor.execute(f"CREATE DATABASE {database_name}") def drop_postgres_db(self, database_name=""): database_name = self.database_or_default(database_name) self.cursor.execute(f"DROP DATABASE IF EXISTS {database_name} WITH (FORCE)") if database_name in self.created_postgres_db_list: self.created_postgres_db_list.remove(database_name) def create_clickhouse_postgres_db( self, database_name="", schema_name="", postgres_database="", ): database_name = self.database_or_default(database_name) if postgres_database == "": postgres_database = database_name self.drop_clickhouse_postgres_db(database_name) self.created_ch_postgres_db_list.add(database_name) if len(schema_name) == 0: self.instance.query( f""" CREATE DATABASE {database_name} ENGINE = PostgreSQL('{self.ip}:{self.port}', '{postgres_database}', 'postgres', 'mysecretpassword')""" ) else: self.instance.query( f""" CREATE DATABASE {database_name} ENGINE = PostgreSQL('{self.ip}:{self.port}', '{postgres_database}', 'postgres', 'mysecretpassword', '{schema_name}')""" ) def drop_clickhouse_postgres_db(self, database_name=""): database_name = self.database_or_default(database_name) self.instance.query(f"DROP DATABASE IF EXISTS {database_name}") if database_name in self.created_ch_postgres_db_list: self.created_ch_postgres_db_list.remove(database_name) def create_materialized_db( self, ip, port, materialized_database="test_database", postgres_database="", settings=[], table_overrides="", ): postgres_database = self.database_or_default(postgres_database) self.created_materialized_postgres_db_list.add(materialized_database) self.instance.query(f"DROP DATABASE IF EXISTS {materialized_database}") create_query = f"CREATE DATABASE {materialized_database} ENGINE = MaterializedPostgreSQL('{ip}:{port}', '{postgres_database}', 'postgres', 'mysecretpassword')" if len(settings) > 0: create_query += " SETTINGS " for i in range(len(settings)): if i != 0: create_query += ", " create_query += settings[i] create_query += table_overrides self.instance.query(create_query) assert materialized_database in self.instance.query("SHOW DATABASES") def drop_materialized_db(self, materialized_database="test_database"): self.instance.query(f"DROP DATABASE IF EXISTS {materialized_database} SYNC") if materialized_database in self.created_materialized_postgres_db_list: self.created_materialized_postgres_db_list.remove(materialized_database) def create_postgres_schema(self, name): create_postgres_schema(self.cursor, name) def create_postgres_table( self, table_name, database_name="", template=postgres_table_template ): create_postgres_table( self.cursor, table_name, database_name=database_name, template=template ) def create_and_fill_postgres_table(self, table_name, database_name=""): create_postgres_table(self.cursor, table_name, database_name) database_name = self.database_or_default(database_name) self.instance.query( f"INSERT INTO {database_name}.{table_name} SELECT number, number from numbers(50)" ) def create_and_fill_postgres_tables(self, tables_num, numbers=50, database_name=""): for i in range(tables_num): table_name = f"postgresql_replica_{i}" create_postgres_table(self.cursor, table_name, database_name) if numbers > 0: db = self.database_or_default(database_name) self.instance.query( f"INSERT INTO {db}.{table_name} SELECT number, number from numbers({numbers})" ) queries = [ "INSERT INTO postgresql_replica_{} select i, i from generate_series(0, 10000) as t(i);", "DELETE FROM postgresql_replica_{} WHERE (value*value) % 3 = 0;", "UPDATE postgresql_replica_{} SET value = value - 125 WHERE key % 2 = 0;", "UPDATE postgresql_replica_{} SET key=key+20000 WHERE key%2=0", "INSERT INTO postgresql_replica_{} select i, i from generate_series(40000, 50000) as t(i);", "DELETE FROM postgresql_replica_{} WHERE key % 10 = 0;", "UPDATE postgresql_replica_{} SET value = value + 101 WHERE key % 2 = 1;", "UPDATE postgresql_replica_{} SET key=key+80000 WHERE key%2=1", "DELETE FROM postgresql_replica_{} WHERE value % 2 = 0;", "UPDATE postgresql_replica_{} SET value = value + 2000 WHERE key % 5 = 0;", "INSERT INTO postgresql_replica_{} select i, i from generate_series(200000, 250000) as t(i);", "DELETE FROM postgresql_replica_{} WHERE value % 3 = 0;", "UPDATE postgresql_replica_{} SET value = value * 2 WHERE key % 3 = 0;", "UPDATE postgresql_replica_{} SET key=key+500000 WHERE key%2=1", "INSERT INTO postgresql_replica_{} select i, i from generate_series(1000000, 1050000) as t(i);", "DELETE FROM postgresql_replica_{} WHERE value % 9 = 2;", "UPDATE postgresql_replica_{} SET key=key+10000000", "UPDATE postgresql_replica_{} SET value = value + 2 WHERE key % 3 = 1;", "DELETE FROM postgresql_replica_{} WHERE value%5 = 0;", ] def assert_nested_table_is_created( instance, table_name, materialized_database="test_database", schema_name="" ): if len(schema_name) == 0: table = table_name else: table = schema_name + "." + table_name print(f"Checking table {table} exists in {materialized_database}") database_tables = instance.query(f"SHOW TABLES FROM {materialized_database}") while table not in database_tables: time.sleep(0.2) database_tables = instance.query(f"SHOW TABLES FROM {materialized_database}") assert table in database_tables def assert_number_of_columns( instance, expected, table_name, database_name="test_database" ): result = instance.query( f"select count() from system.columns where table = '{table_name}' and database = '{database_name}' and not startsWith(name, '_')" ) while int(result) != expected: time.sleep(1) result = instance.query( f"select count() from system.columns where table = '{table_name}' and database = '{database_name}' and not startsWith(name, '_')" ) print("Number of columns ok") def check_tables_are_synchronized( instance, table_name, order_by="key", postgres_database="postgres_database", materialized_database="test_database", schema_name="", ): assert_nested_table_is_created( instance, table_name, materialized_database, schema_name ) table_path = "" if len(schema_name) == 0: table_path = f"{materialized_database}.{table_name}" else: table_path = f"{materialized_database}.`{schema_name}.{table_name}`" print(f"Checking table is synchronized: {table_path}") result_query = f"select * from {table_path} order by {order_by};" expected = instance.query( f"select * from {postgres_database}.{table_name} order by {order_by};" ) result = instance.query(result_query) for _ in range(50): if result == expected: break else: time.sleep(1) result = instance.query(result_query) if result != expected: count = int(instance.query(f"select count() from {table_path}")) expected_count = int( instance.query(f"select count() from {postgres_database}.{table_name}") ) print(f"Having {count}, expected {expected_count}") assert result == expected def check_several_tables_are_synchronized( instance, tables_num, order_by="key", postgres_database="postgres_database", materialized_database="test_database", schema_name="", ): for i in range(tables_num): check_tables_are_synchronized(instance, f"postgresql_replica_{i}")