Skip to content

Commit 04825d4

Browse files
committed
Update old code
1 parent ed4e617 commit 04825d4

File tree

11 files changed

+37
-30
lines changed

11 files changed

+37
-30
lines changed

text_2_sql/.env.example

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,20 @@ AIService__AzureSearchOptions__Text2SqlQueryCache__Index=<Query cache index name
1717
AIService__AzureSearchOptions__Text2SqlQueryCache__SemanticConfig=<Query cache semantic config. Default is created as "text-2-sql-query-cache-semantic-config">
1818
AIService__AzureSearchOptions__Text2SqlColumnValueStore__Index=<Column value store index name. Default is created as "text-2-sql-column-value-store-index">
1919

20-
# All SQL Engine specific connection details
21-
Text2Sql__DatabaseName=<databaseName>
20+
# TSQL
21+
Text2Sql__Tsql__ConnectionString=<Tsql databaseConnectionString if using Tsql Data Source>
22+
Text2Sql__Database=<Tsql database if using Tsql Data Source>
2223

23-
# TSQL or PostgreSQL Specific Connection Details
24-
Text2Sql__DatabaseConnectionString=<databaseConnectionString>
24+
# PostgreSQL Specific Connection Details
25+
Text2Sql__Postgresql__ConnectionString=<Postgresql databaseConnectionString if using Postgresql Data Source>
26+
Text2Sql__Postgresql__Database=<Postgresql database if using Postgresql Data Source>
2527

2628
# Snowflake Specific Connection Details
2729
Text2Sql__Snowflake__User=<snowflakeUser if using Snowflake Data Source>
2830
Text2Sql__Snowflake__Password=<snowflakePassword if using Snowflake Data Source>
2931
Text2Sql__Snowflake__Account=<snowflakeAccount if using Snowflake Data Source>
3032
Text2Sql__Snowflake__Warehouse=<snowflakeWarehouse if using Snowflake Data Source>
33+
Text2Sql__Snowflake__Database=<snowflakeDatabase if using Snowflake Data Source>
3134

3235
# Databricks Specific Connection Details
3336
Text2Sql__Databricks__Catalog=<databricksCatalog if using Databricks Data Source with Unity Catalog>

text_2_sql/autogen/evaluate_autogen_text2sql.ipynb

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -215,8 +215,8 @@
215215
" \n",
216216
" # Update database connection string for current database\n",
217217
" db_path = DATABASE_DIR / db_id / f\"{db_id}.sqlite\"\n",
218-
" os.environ[\"Text2Sql__DatabaseConnectionString\"] = str(db_path)\n",
219-
" os.environ[\"Text2Sql__DatabaseName\"] = db_id\n",
218+
" os.environ[\"Text2Sql__Tsql__ConnectionString\"] = str(db_path)\n",
219+
" os.environ[\"Text2Sql__Database\"] = db_id\n",
220220
" \n",
221221
" sql = await generate_sql(question)\n",
222222
" predictions.append(f\"{sql}\\t{db_id}\")\n",

text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/parallel_query_solving_agent.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -219,12 +219,12 @@ async def consume_inner_messages_from_agentic_flow(
219219

220220
# Add database connection info to injected parameters
221221
query_params = injected_parameters.copy() if injected_parameters else {}
222-
if "Text2Sql__DatabaseConnectionString" in os.environ:
222+
if "Text2Sql__Tsql__ConnectionString" in os.environ:
223223
query_params["database_connection_string"] = os.environ[
224-
"Text2Sql__DatabaseConnectionString"
224+
"Text2Sql__Tsql__ConnectionString"
225225
]
226-
if "Text2Sql__DatabaseName" in os.environ:
227-
query_params["database_name"] = os.environ["Text2Sql__DatabaseName"]
226+
if "Text2Sql__Database" in os.environ:
227+
query_params["database_name"] = os.environ["Text2Sql__Database"]
228228

229229
# Launch tasks for each sub-query
230230
inner_solving_generators.append(

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -45,27 +45,25 @@ def __init__(self, **kwargs: dict):
4545
self.set_mode()
4646

4747
# Store original environment variables
48-
self.original_db_conn = os.environ.get("Text2Sql__DatabaseConnectionString")
49-
self.original_db_name = os.environ.get("Text2Sql__DatabaseName")
48+
self.original_db_conn = os.environ.get("Text2Sql__Tsql__ConnectionString")
49+
self.original_db_name = os.environ.get("Text2Sql__Database")
5050

5151
def _update_environment(self, injected_parameters: dict = None):
5252
"""Update environment variables with injected parameters."""
5353
if injected_parameters:
5454
if "database_connection_string" in injected_parameters:
55-
os.environ["Text2Sql__DatabaseConnectionString"] = injected_parameters[
55+
os.environ["Text2Sql__Tsql__ConnectionString"] = injected_parameters[
5656
"database_connection_string"
5757
]
5858
if "database_name" in injected_parameters:
59-
os.environ["Text2Sql__DatabaseName"] = injected_parameters[
60-
"database_name"
61-
]
59+
os.environ["Text2Sql__Database"] = injected_parameters["database_name"]
6260

6361
def _restore_environment(self):
6462
"""Restore original environment variables."""
6563
if self.original_db_conn:
66-
os.environ["Text2Sql__DatabaseConnectionString"] = self.original_db_conn
64+
os.environ["Text2Sql__Tsql__ConnectionString"] = self.original_db_conn
6765
if self.original_db_name:
68-
os.environ["Text2Sql__DatabaseName"] = self.original_db_name
66+
os.environ["Text2Sql__Database"] = self.original_db_name
6967

7068
def set_mode(self):
7169
"""Set the mode of the plugin based on the environment variables."""

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/postgresql_sql.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ async def query_execution(
6666
"""
6767
logging.info(f"Running query: {sql_query}")
6868
results = []
69-
connection_string = os.environ["Text2Sql__DatabaseConnectionString"]
69+
connection_string = os.environ["Text2Sql__Postgresql__ConnectionString"]
7070

7171
# Establish an asynchronous connection to the PostgreSQL database
7272
async with await psycopg.AsyncConnection.connect(connection_string) as conn:

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/snowflake_sql.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ async def query_execution(
100100
password=os.environ["Text2Sql__Snowflake__Password"],
101101
account=os.environ["Text2Sql__Snowflake__Account"],
102102
warehouse=os.environ["Text2Sql__Snowflake__Warehouse"],
103-
database=os.environ["Text2Sql__DatabaseName"],
103+
database=os.environ["Text2Sql__Database"],
104104
)
105105

106106
try:

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sqlite_sql.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ async def query_execution(
6363
Returns:
6464
List of dictionaries containing query results.
6565
"""
66-
db_file = os.environ["Text2Sql__DatabaseConnectionString"]
66+
db_file = os.environ["Text2Sql__Tsql__ConnectionString"]
6767

6868
if not os.path.exists(db_file):
6969
raise FileNotFoundError(f"Database file not found: {db_file}")
@@ -127,7 +127,9 @@ def find_matching_tables(self, text: str, table_names: list[str]) -> list[int]:
127127
List of matching table indices
128128
"""
129129
matches = []
130-
logging.info(f"Looking for tables matching '{text}' in tables: {table_names}")
130+
logging.info(
131+
"Looking for tables matching '%s' in tables: %s", text, table_names
132+
)
131133

132134
# First try exact matches
133135
for idx, name in enumerate(table_names):
@@ -144,7 +146,9 @@ def find_matching_tables(self, text: str, table_names: list[str]) -> list[int]:
144146
for idx, name in enumerate(table_names):
145147
table_terms = set(re.split(r"[_\s]+", name.lower()))
146148
if search_terms & table_terms: # If there's any overlap in terms
147-
logging.info(f"Found partial match: '{name}' with terms {table_terms}")
149+
logging.info(
150+
"Found partial match: '%s' with terms %s", name, table_terms
151+
)
148152
matches.append(idx)
149153

150154
return matches
@@ -181,7 +185,7 @@ async def get_entity_schemas(
181185
spider_schemas = json.load(f)
182186

183187
# Get current database name from path
184-
db_path = os.environ["Text2Sql__DatabaseConnectionString"]
188+
db_path = os.environ["Text2Sql__Tsql__ConnectionString"]
185189
db_name = os.path.splitext(os.path.basename(db_path))[0]
186190

187191
logging.info(f"Looking for schemas in database: {db_name}")
@@ -196,7 +200,7 @@ async def get_entity_schemas(
196200
if not db_schema:
197201
raise ValueError(f"Schema not found for database: {db_name}")
198202

199-
logging.info(f"Looking for tables matching '{text}' in database '{db_name}'")
203+
logging.info("Looking for tables matching '%s' in database '%s'", text, db_name)
200204
logging.info(f"Available tables: {db_schema['table_names']}")
201205

202206
# Find all matching tables using flexible matching
@@ -228,7 +232,9 @@ async def get_entity_schemas(
228232
}
229233
schemas.append(schema)
230234
logging.info(
231-
f"Added schema for table '{db_schema['table_names'][table_idx]}': {schema}"
235+
"Added schema for table '%s': %s",
236+
db_schema["table_names"][table_idx],
237+
schema,
232238
)
233239

234240
if as_json:

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/tsql_sql.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ async def query_execution(
8686
"""
8787
logging.info(f"Running query: {sql_query}")
8888
results = []
89-
connection_string = os.environ["Text2Sql__DatabaseConnectionString"]
89+
connection_string = os.environ["Text2Sql__Tsql__ConnectionString"]
9090
async with await aioodbc.connect(dsn=connection_string) as sql_db_client:
9191
async with sql_db_client.cursor() as cursor:
9292
await cursor.execute(sql_query)

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/postgresql_data_dictionary_creator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ def __init__(self, **kwargs):
1616
excluded_schemas = ["information_schema", "pg_catalog"]
1717
super().__init__(excluded_schemas=excluded_schemas, **kwargs)
1818

19-
self.database = os.environ["Text2Sql__DatabaseName"]
19+
self.database = os.environ["Text2Sql__Postgresql__Database"]
2020
self.database_engine = DatabaseEngine.POSTGRESQL
2121

2222
self.sql_connector = PostgresqlSqlConnector()

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/snowflake_data_dictionary_creator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ def __init__(self, **kwargs):
1818
excluded_schemas = ["INFORMATION_SCHEMA"]
1919
super().__init__(excluded_schemas=excluded_schemas, **kwargs)
2020

21-
self.database = os.environ["Text2Sql__DatabaseName"]
21+
self.database = os.environ["Text2Sql__Snowflake__Database"]
2222
self.warehouse = os.environ["Text2Sql__Snowflake__Warehouse"]
2323
self.database_engine = DatabaseEngine.SNOWFLAKE
2424

0 commit comments

Comments
 (0)