Skip to content

Commit 93b846c

Browse files
vieiralucasvieiralucas
committed
ping and fetch sql server schema with python
1 parent 8a94e95 commit 93b846c

File tree

3 files changed

+40
-186
lines changed

3 files changed

+40
-186
lines changed

‎apps/api/src/datasources/sqlserver.ts‎

Lines changed: 6 additions & 181 deletions
Original file line numberDiff line numberDiff line change
@@ -1,197 +1,22 @@
11
import { config } from '../config/index.js'
2-
import prisma, {
3-
SQLServerDataSource,
4-
getSQLServerCert,
5-
getSQLServerPassword,
6-
} from '@briefer/database'
7-
import sql from 'mssql'
8-
import { logger } from '../logger.js'
9-
import { DataSourceColumn, DataSourceConnectionError } from '@briefer/types'
2+
import prisma, { SQLServerDataSource } from '@briefer/database'
103
import { DataSourceStatus } from './index.js'
11-
import { OnTable } from './structure.js'
12-
13-
type ConnectionConfig = {
14-
user: string
15-
password: string
16-
server: string
17-
database: string
18-
ssl?: { ca: Buffer }
19-
}
20-
21-
async function getSQLServerConfig(
22-
datasource: SQLServerDataSource
23-
): Promise<ConnectionConfig> {
24-
const password = await getSQLServerPassword(
25-
datasource,
26-
config().DATASOURCES_ENCRYPTION_KEY
27-
)
28-
29-
const cert = await getSQLServerCert(
30-
datasource,
31-
config().DATASOURCES_ENCRYPTION_KEY
32-
)
33-
34-
return {
35-
user: datasource.username,
36-
password,
37-
database: datasource.database,
38-
server: datasource.host,
39-
ssl: cert ? { ca: cert } : undefined,
40-
}
41-
}
4+
import { pingSQLServer } from '../python/query/sqlserver.js'
425

436
export async function ping(
44-
datasource: SQLServerDataSource
7+
ds: SQLServerDataSource
458
): Promise<SQLServerDataSource> {
469
const lastConnection = new Date()
47-
const SQLServerConfig = await getSQLServerConfig(datasource)
48-
49-
const err = await pingSQLServerFromConfig(SQLServerConfig)
5010

11+
const err = await pingSQLServer(ds, config().DATASOURCES_ENCRYPTION_KEY)
5112
if (!err) {
52-
return updateConnStatus(datasource, {
13+
return updateConnStatus(ds, {
5314
connStatus: 'online',
5415
lastConnection,
5516
})
5617
}
5718

58-
return updateConnStatus(datasource, { connStatus: 'offline', connError: err })
59-
}
60-
61-
async function createConnection(
62-
config: ConnectionConfig
63-
): Promise<sql.ConnectionPool> {
64-
const mustEncrypt = config.ssl ? true : false
65-
return sql.connect({
66-
user: config.user,
67-
password: config.password,
68-
server: config.server,
69-
database: config.database,
70-
options: {
71-
trustServerCertificate: mustEncrypt ? false : true,
72-
encrypt: mustEncrypt,
73-
cryptoCredentialsDetails: {
74-
ca: config.ssl?.ca,
75-
},
76-
},
77-
})
78-
}
79-
80-
async function pingSQLServerFromConfig(
81-
config: ConnectionConfig
82-
): Promise<DataSourceConnectionError | null> {
83-
try {
84-
const connection = await createConnection(config)
85-
86-
return await Promise.race([
87-
new Promise<DataSourceConnectionError>((resolve) =>
88-
setTimeout(
89-
() =>
90-
resolve({
91-
name: 'TimeoutError',
92-
message: 'Did not receive response from SQLServer within 10s',
93-
}),
94-
10000 // 10s timeout
95-
)
96-
),
97-
new Promise<DataSourceConnectionError | null>(async (resolve) => {
98-
try {
99-
await connection.query('SELECT 1')
100-
resolve(null)
101-
} catch (err) {
102-
logger().info({ err }, 'Error pinging SQLServer')
103-
const parsedError = DataSourceConnectionError.safeParse(err)
104-
if (!parsedError.success) {
105-
logger().error(
106-
{
107-
error: err,
108-
},
109-
'Failed to parse error from SQLServer ping'
110-
)
111-
resolve({ name: 'UnknownError', message: 'Unknown error' })
112-
return
113-
}
114-
115-
resolve(parsedError.data)
116-
} finally {
117-
connection.close()
118-
}
119-
}),
120-
])
121-
} catch (err) {
122-
logger().info({ err }, 'Error pinging SQLServer')
123-
const parsedError = DataSourceConnectionError.safeParse(err)
124-
if (!parsedError.success) {
125-
logger().error(
126-
{
127-
error: err,
128-
},
129-
'Failed to parse error from SQLServer ping'
130-
)
131-
return { name: 'UnknownError', message: 'Unknown error' }
132-
}
133-
134-
return parsedError.data
135-
}
136-
}
137-
138-
export async function getSQLServerSchemaFromConfig(
139-
datasourceId: string,
140-
sqlServerConfig: ConnectionConfig,
141-
onTable: OnTable
142-
): Promise<void> {
143-
const connection = await createConnection(sqlServerConfig)
144-
145-
// select all tables with their column names and types from all schemas
146-
const result = await connection.query(`
147-
SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, DATA_TYPE
148-
FROM INFORMATION_SCHEMA.COLUMNS
149-
WHERE TABLE_SCHEMA NOT IN ('information_schema', 'sys', 'master', 'tempdb', 'model', 'msdb');
150-
`)
151-
await connection.close()
152-
153-
const schemas: Record<
154-
string,
155-
Record<string, { columns: DataSourceColumn[] }>
156-
> = {}
157-
158-
for (const row of result.recordset) {
159-
const schemaName = row.TABLE_SCHEMA
160-
const tableName = row.TABLE_NAME
161-
const columnName = row.COLUMN_NAME
162-
const dataType = row.DATA_TYPE
163-
164-
let schema = schemas[schemaName]
165-
if (!schema) {
166-
schema = {}
167-
schemas[schemaName] = schema
168-
}
169-
170-
let table = schema[tableName]
171-
if (!table) {
172-
table = {
173-
columns: [],
174-
}
175-
schema[tableName] = table
176-
}
177-
178-
table.columns.push({ name: columnName, type: dataType })
179-
}
180-
181-
for (const [schemaName, schema] of Object.entries(schemas)) {
182-
for (const [tableName, table] of Object.entries(schema)) {
183-
onTable(schemaName, tableName, table, datasourceId)
184-
}
185-
}
186-
}
187-
188-
export async function getSqlServerSchema(
189-
datasource: SQLServerDataSource,
190-
onTable: OnTable
191-
): Promise<void> {
192-
const SQLServerConfig = await getSQLServerConfig(datasource)
193-
194-
return getSQLServerSchemaFromConfig(datasource.id, SQLServerConfig, onTable)
19+
return updateConnStatus(ds, { connStatus: 'offline', connError: err })
19520
}
19621

19722
export async function updateConnStatus(

‎apps/api/src/datasources/structure.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,10 @@ import { getAthenaSchema } from './athena.js'
3030
import { getMySQLSchema } from './mysql.js'
3131
import { getDatabricksSQLSchema } from '../python/query/databrickssql.js'
3232
import { PythonExecutionError } from '../python/index.js'
33-
import { getSqlServerSchema } from './sqlserver.js'
3433
import { z } from 'zod'
3534
import { splitEvery } from 'ramda'
3635
import { createEmbedding } from '../embedding.js'
36+
import { getSQLServerSchema } from '../python/query/sqlserver.js'
3737

3838
async function assignDataSourceSchemaId(
3939
dataSourceId: string,
@@ -500,7 +500,11 @@ async function _refreshDataSourceStructure(
500500
await getAthenaSchema(dataSource.config.data, onTable)
501501
break
502502
case 'sqlserver':
503-
await getSqlServerSchema(dataSource.config.data, onTable)
503+
await getSQLServerSchema(
504+
dataSource.config.data,
505+
config().DATASOURCES_ENCRYPTION_KEY,
506+
onTable
507+
)
504508
break
505509
case 'trino':
506510
await getTrinoSchema(

‎apps/api/src/python/query/sqlserver.ts‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,19 @@
11
import { v4 as uuidv4 } from 'uuid'
2-
import { MySQLDataSource, getDatabaseURL } from '@briefer/database'
2+
import { SQLServerDataSource, getDatabaseURL } from '@briefer/database'
33
import { RunQueryResult, SuccessRunQueryResult } from '@briefer/types'
4-
import { makeSQLAlchemyQuery } from './sqlalchemy.js'
4+
import {
5+
getSQLAlchemySchema,
6+
makeSQLAlchemyQuery,
7+
pingSQLAlchemy,
8+
} from './sqlalchemy.js'
9+
import { OnTable } from '../../datasources/structure.js'
510

611
export async function makeSQLServerQuery(
712
workspaceId: string,
813
sessionId: string,
914
queryId: string,
1015
dataframeName: string,
11-
datasource: MySQLDataSource,
16+
datasource: SQLServerDataSource,
1217
encryptionKey: string,
1318
sql: string,
1419
resultOptions: { pageSize: number; dashboardPageSize: number },
@@ -35,3 +40,23 @@ export async function makeSQLServerQuery(
3540
onProgress
3641
)
3742
}
43+
44+
export function pingSQLServer(
45+
ds: SQLServerDataSource,
46+
encryptionKey: string
47+
): Promise<null | Error> {
48+
return pingSQLAlchemy({ type: 'sqlserver', data: ds }, encryptionKey, null)
49+
}
50+
51+
export function getSQLServerSchema(
52+
ds: SQLServerDataSource,
53+
encryptionKey: string,
54+
onTable: OnTable
55+
): Promise<void> {
56+
return getSQLAlchemySchema(
57+
{ type: 'sqlserver', data: ds },
58+
encryptionKey,
59+
null,
60+
onTable
61+
)
62+
}

0 commit comments

Comments
 (0)