Skip to content

Commit

Permalink
feat: Refactor for a more efficient prepare_table() (#228)
Browse files Browse the repository at this point in the history
Refactor code in prepare_table() to drastically reduce database queries
required from O(n) db queries to O(1) db queries, where n is the number
of columns in the table.

This change improves test_reserved_keywords() running time and closes
#142

Running `time poetry run pytest -k 'test_reserved_keywords'` on main
branch gave me an average running time of 33.441 seconds across three
runs:

```
real    0m33.539s
user    0m21.294s
sys     0m0.531s

real    0m33.300s
user    0m21.272s
sys     0m0.472s

real    0m33.484s
user    0m21.337s
sys     0m0.503s
```

Running `time poetry run pytest -k 'test_reserved_keywords'` after my
changes gave me an average running time of 1.271 seconds across three
runs:

```
real    0m1.232s
user    0m1.005s
sys     0m0.056s

real    0m1.299s
user    0m1.018s
sys     0m0.105s

real    0m1.281s
user    0m1.024s
sys     0m0.082s
```
  • Loading branch information
sebastianswms authored Nov 28, 2023
1 parent 840181e commit d59fb44
Showing 1 changed file with 29 additions and 7 deletions.
36 changes: 29 additions & 7 deletions target_postgres/connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,13 +118,24 @@ def prepare_table( # type: ignore[override]
table = meta.tables[
full_table_name
] # So we don't mess up the casing of the Table reference

columns = self.get_table_columns(
schema_name=cast(str, schema_name),
table_name=table_name,
connection=connection,
)

for property_name, property_def in schema["properties"].items():
column_object = None
if property_name in columns:
column_object = columns[property_name]
self.prepare_column(
schema_name=cast(str, schema_name),
table=table,
column_name=property_name,
sql_type=self.to_sql_type(property_def),
connection=connection,
column_object=column_object,
)

return meta.tables[full_table_name]
Expand Down Expand Up @@ -366,6 +377,7 @@ def prepare_column( # type: ignore[override]
column_name: str,
sql_type: sqlalchemy.types.TypeEngine,
connection: sqlalchemy.engine.Connection,
column_object: sqlalchemy.Column | None = None,
) -> None:
"""Adapt target table to provided schema if possible.
Expand All @@ -376,7 +388,11 @@ def prepare_column( # type: ignore[override]
sql_type: the SQLAlchemy type.
connection: the database connection.
"""
if not self.column_exists(table.fullname, column_name, connection=connection):
column_exists = column_object is not None or self.column_exists(
table.fullname, column_name, connection=connection
)

if not column_exists:
self._create_empty_column(
# We should migrate every function to use sqlalchemy.Table
# instead of having to know what the function wants
Expand All @@ -394,6 +410,7 @@ def prepare_column( # type: ignore[override]
column_name=column_name,
sql_type=sql_type,
connection=connection,
column_object=column_object,
)

def _create_empty_column( # type: ignore[override]
Expand Down Expand Up @@ -466,6 +483,7 @@ def _adapt_column_type( # type: ignore[override]
column_name: str,
sql_type: sqlalchemy.types.TypeEngine,
connection: sqlalchemy.engine.Connection,
column_object: sqlalchemy.Column | None,
) -> None:
"""Adapt table column type to support the new JSON schema type.
Expand All @@ -477,12 +495,16 @@ def _adapt_column_type( # type: ignore[override]
Raises:
NotImplementedError: if altering columns is not supported.
"""
current_type: sqlalchemy.types.TypeEngine = self._get_column_type(
schema_name=schema_name,
table_name=table_name,
column_name=column_name,
connection=connection,
)
current_type: sqlalchemy.types.TypeEngine
if column_object is not None:
current_type = t.cast(sqlalchemy.types.TypeEngine, column_object.type)
else:
current_type = self._get_column_type(
schema_name=schema_name,
table_name=table_name,
column_name=column_name,
connection=connection,
)

# remove collation if present and save it
current_type_collation = self.remove_collation(current_type)
Expand Down

0 comments on commit d59fb44

Please sign in to comment.